added caching of images to the drive

cache-images.go

@@ -0,0 +1,223 @@
+package main
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"image"
+	"image/gif"
+	"image/jpeg"
+	"image/png"
+	"io"
+	"net/http"
+	"os"
+	"path/filepath"
+	"strings"
+	"sync"
+
+	"github.com/chai2010/webp"
+	"golang.org/x/image/bmp"
+	"golang.org/x/image/tiff"
+)
+
+var (
+	cachingImages    = make(map[string]*sync.Mutex)
+	cachingImagesMu  sync.Mutex
+	cachingSemaphore = make(chan struct{}, 10) // Limit to 10 concurrent downloads
+)
+
+func cacheImage(imageURL, filename string) (string, error) {
+	cacheDir := "image_cache"
+	cachedImagePath := filepath.Join(cacheDir, filename)
+
+	// Check if the image is already cached
+	if _, err := os.Stat(cachedImagePath); err == nil {
+		return cachedImagePath, nil
+	}
+
+	// Ensure only one goroutine caches the same image
+	cachingImagesMu.Lock()
+	if _, exists := cachingImages[imageURL]; !exists {
+		cachingImages[imageURL] = &sync.Mutex{}
+	}
+	mu := cachingImages[imageURL]
+	cachingImagesMu.Unlock()
+
+	mu.Lock()
+	defer mu.Unlock()
+
+	// Double-check if the image was cached while waiting
+	if _, err := os.Stat(cachedImagePath); err == nil {
+		return cachedImagePath, nil
+	}
+
+	cachingSemaphore <- struct{}{}        // Acquire a token
+	defer func() { <-cachingSemaphore }() // Release the token
+
+	// Download the image
+	resp, err := http.Get(imageURL)
+	if err != nil {
+		return "", err
+	}
+	defer resp.Body.Close()
+
+	// Read the image data into a byte slice
+	data, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return "", err
+	}
+
+	// Detect the content type
+	contentType := http.DetectContentType(data)
+
+	// If content type is HTML, skip caching
+	if strings.HasPrefix(contentType, "text/html") {
+		return "", fmt.Errorf("URL returned HTML content instead of an image: %s", imageURL)
+	}
+
+	// Handle SVG files directly
+	if contentType == "image/svg+xml" {
+		// Ensure the cache directory exists
+		if _, err := os.Stat(cacheDir); os.IsNotExist(err) {
+			os.Mkdir(cacheDir, os.ModePerm)
+		}
+
+		// Save the SVG file as-is
+		err = os.WriteFile(cachedImagePath, data, 0644)
+		if err != nil {
+			return "", err
+		}
+
+		// Clean up mutex
+		cachingImagesMu.Lock()
+		delete(cachingImages, imageURL)
+		cachingImagesMu.Unlock()
+
+		return cachedImagePath, nil
+	}
+
+	// Decode the image based on the content type
+	var img image.Image
+	switch contentType {
+	case "image/jpeg":
+		img, err = jpeg.Decode(bytes.NewReader(data))
+	case "image/png":
+		img, err = png.Decode(bytes.NewReader(data))
+	case "image/gif":
+		img, err = gif.Decode(bytes.NewReader(data))
+	case "image/webp":
+		img, err = webp.Decode(bytes.NewReader(data))
+	case "image/bmp":
+		img, err = bmp.Decode(bytes.NewReader(data))
+	case "image/tiff":
+		img, err = tiff.Decode(bytes.NewReader(data))
+	default:
+		return "", fmt.Errorf("unsupported image type: %s", contentType)
+	}
+
+	if err != nil {
+		return "", err
+	}
+
+	// Ensure the cache directory exists
+	if _, err := os.Stat(cacheDir); os.IsNotExist(err) {
+		os.Mkdir(cacheDir, os.ModePerm)
+	}
+
+	// Open the cached file for writing
+	outFile, err := os.Create(cachedImagePath)
+	if err != nil {
+		return "", err
+	}
+	defer outFile.Close()
+
+	// Encode the image to WebP and save
+	options := &webp.Options{Lossless: false, Quality: 80}
+	err = webp.Encode(outFile, img, options)
+	if err != nil {
+		return "", err
+	}
+
+	// Clean up mutex
+	cachingImagesMu.Lock()
+	delete(cachingImages, imageURL)
+	cachingImagesMu.Unlock()
+
+	return cachedImagePath, nil
+}
+
+func handleCachedImages(w http.ResponseWriter, r *http.Request) {
+	imageName := filepath.Base(r.URL.Path)
+	cacheDir := "image_cache"
+	cachedImagePath := filepath.Join(cacheDir, imageName)
+
+	if _, err := os.Stat(cachedImagePath); os.IsNotExist(err) {
+		// Serve placeholder image with no-store headers
+		placeholderPath := "static/images/placeholder.webp"
+		placeholderContentType := "image/webp"
+
+		// You can also check for SVG placeholder if needed
+		if strings.HasSuffix(imageName, ".svg") {
+			placeholderPath = "static/images/placeholder.svg"
+			placeholderContentType = "image/svg+xml"
+		}
+
+		w.Header().Set("Content-Type", placeholderContentType)
+		w.Header().Set("Cache-Control", "no-store, must-revalidate")
+		w.Header().Set("Pragma", "no-cache")
+		w.Header().Set("Expires", "0")
+		http.ServeFile(w, r, placeholderPath)
+		return
+	}
+
+	// Determine the content type based on the file extension
+	extension := strings.ToLower(filepath.Ext(cachedImagePath))
+	var contentType string
+	switch extension {
+	case ".svg":
+		contentType = "image/svg+xml"
+	case ".jpg", ".jpeg":
+		contentType = "image/jpeg"
+	case ".png":
+		contentType = "image/png"
+	case ".gif":
+		contentType = "image/gif"
+	case ".webp":
+		contentType = "image/webp"
+	default:
+		// Default to binary stream if unknown
+		contentType = "application/octet-stream"
+	}
+
+	w.Header().Set("Content-Type", contentType)
+	w.Header().Set("Cache-Control", "public, max-age=31536000") // Cache the image for 1 year
+	http.ServeFile(w, r, cachedImagePath)
+}
+
+func handleImageStatus(w http.ResponseWriter, r *http.Request) {
+	imageIDs := r.URL.Query().Get("image_ids")
+	ids := strings.Split(imageIDs, ",")
+
+	statusMap := make(map[string]string)
+
+	cacheDir := "image_cache"
+
+	printDebug("Received image status request for IDs: %v", ids)
+	printDebug("Status map: %v", statusMap)
+
+	for _, id := range ids {
+		filename := id + ".webp"
+		cachedImagePath := filepath.Join(cacheDir, filename)
+
+		if _, err := os.Stat(cachedImagePath); err == nil {
+			// Image is cached and ready
+			statusMap[id] = "/image_cache/" + filename
+		} else {
+			// Image is not ready
+			statusMap[id] = ""
+		}
+	}
+
+	w.Header().Set("Content-Type", "application/json")
+	json.NewEncoder(w).Encode(statusMap)
+}