datascape/thumb.go

package main

import (
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"path"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"
)

// Thumbnailer produces a thumbnail of a source file. Implementations register
// themselves in init() by appending to thumbnailers. The first registered
// handler whose CanHandle returns true is used.
type Thumbnailer interface {
	CanHandle(ext string) bool
	Generate(src io.Reader, dst io.Writer, width int) error
}

var thumbnailers []Thumbnailer

// thumbCacheDir is set from the -cache flag at startup.
var thumbCacheDir string

const thumbURLPrefix = "/_thumb"

// Cache is content-addressed: the cache path is derived from the SHA-256 of
// the source file. Renames and moves reuse the same cache entry; overwriting
// a file with new content produces a new digest and regenerates.
var (
	thumbLocks   = map[string]*sync.Mutex{}
	thumbLocksMu sync.Mutex

	digestCache  = map[string]digestEntry{}
	digestCacheMu sync.Mutex
)

// digestEntry remembers the digest of a source file so repeated requests do
// not re-hash the whole file. The (mtime, size) pair invalidates the cache
// when the file is overwritten in place.
type digestEntry struct {
	mtime time.Time
	size  int64
	hex   string
}

func findThumbnailer(name string) Thumbnailer {
	ext := strings.ToLower(filepath.Ext(name))
	if ext == "" {
		return nil
	}
	for _, t := range thumbnailers {
		if t.CanHandle(ext) {
			return t
		}
	}
	return nil
}

// hasThumbnail reports whether a file name has a registered thumbnailer.
func hasThumbnail(name string) bool {
	return findThumbnailer(name) != nil
}

// thumbURL builds a thumbnail URL for a wiki file. filePath must be URL-style
// (slash-separated, leading slash), as already used on page links.
func thumbURL(filePath string, width int) string {
	return thumbURLPrefix + filePath + "?w=" + strconv.Itoa(width)
}

func (h *handler) handleThumb(w http.ResponseWriter, r *http.Request) {
	raw := strings.TrimPrefix(r.URL.Path, thumbURLPrefix)
	if raw == "" || raw == "/" {
		http.NotFound(w, r)
		return
	}
	decoded, err := url.PathUnescape(raw)
	if err != nil {
		http.Error(w, "bad path", http.StatusBadRequest)
		return
	}
	cleanPath := path.Clean(decoded)

	srcFS := filepath.Join(h.root, filepath.FromSlash(cleanPath))
	rel, err := filepath.Rel(h.root, srcFS)
	if err != nil || strings.HasPrefix(rel, "..") {
		http.Error(w, "Forbidden", http.StatusForbidden)
		return
	}

	srcInfo, err := os.Stat(srcFS)
	if err != nil || srcInfo.IsDir() {
		http.NotFound(w, r)
		return
	}

	t := findThumbnailer(srcFS)
	if t == nil {
		http.NotFound(w, r)
		return
	}

	width := 300
	if s := r.URL.Query().Get("w"); s != "" {
		if n, err := strconv.Atoi(s); err == nil && n > 0 && n <= 2000 {
			width = n
		}
	}

	digest, data, err := sourceDigest(srcFS, srcInfo)
	if err != nil {
		log.Printf("thumb digest %s: %v", rel, err)
		http.Error(w, "thumbnail failed", http.StatusInternalServerError)
		return
	}

	cacheFS := filepath.Join(thumbCacheDir, digest[:2], fmt.Sprintf("%s.%d.jpg", digest, width))
	if _, err := os.Stat(cacheFS); err == nil {
		serveThumb(w, r, cacheFS)
		return
	}

	lock := thumbLock(cacheFS)
	lock.Lock()
	defer lock.Unlock()

	if _, err := os.Stat(cacheFS); err == nil {
		serveThumb(w, r, cacheFS)
		return
	}

	var src io.Reader
	if data != nil {
		src = bytes.NewReader(data)
	} else {
		f, err := os.Open(srcFS)
		if err != nil {
			log.Printf("thumb open %s: %v", rel, err)
			http.Error(w, "thumbnail failed", http.StatusInternalServerError)
			return
		}
		defer f.Close()
		src = f
	}

	if err := generateThumb(t, src, cacheFS, width); err != nil {
		log.Printf("thumb %s: %v", rel, err)
		http.Error(w, "thumbnail failed", http.StatusInternalServerError)
		return
	}
	serveThumb(w, r, cacheFS)
}

// sourceDigest returns the SHA-256 hex digest of a source file's content.
// On a cache hit (path + mtime + size unchanged) the returned data is nil,
// so the caller knows to open the file itself. On a miss the file is read
// once and the contents are returned for the caller to reuse.
func sourceDigest(srcFS string, info os.FileInfo) (string, []byte, error) {
	digestCacheMu.Lock()
	d, ok := digestCache[srcFS]
	digestCacheMu.Unlock()
	if ok && d.mtime.Equal(info.ModTime()) && d.size == info.Size() {
		return d.hex, nil, nil
	}

	data, err := os.ReadFile(srcFS)
	if err != nil {
		return "", nil, err
	}
	sum := sha256.Sum256(data)
	h := hex.EncodeToString(sum[:])

	digestCacheMu.Lock()
	digestCache[srcFS] = digestEntry{mtime: info.ModTime(), size: info.Size(), hex: h}
	digestCacheMu.Unlock()

	return h, data, nil
}

func serveThumb(w http.ResponseWriter, r *http.Request, cacheFS string) {
	w.Header().Set("Cache-Control", "public, max-age=86400")
	http.ServeFile(w, r, cacheFS)
}

func generateThumb(t Thumbnailer, src io.Reader, cacheFS string, width int) error {
	if err := os.MkdirAll(filepath.Dir(cacheFS), 0755); err != nil {
		return err
	}
	tmp, err := os.CreateTemp(filepath.Dir(cacheFS), ".thumb-*")
	if err != nil {
		return err
	}
	tmpName := tmp.Name()
	if err := t.Generate(src, tmp, width); err != nil {
		tmp.Close()
		os.Remove(tmpName)
		return err
	}
	if err := tmp.Close(); err != nil {
		os.Remove(tmpName)
		return err
	}
	return os.Rename(tmpName, cacheFS)
}

func thumbLock(key string) *sync.Mutex {
	thumbLocksMu.Lock()
	defer thumbLocksMu.Unlock()
	m, ok := thumbLocks[key]
	if !ok {
		m = &sync.Mutex{}
		thumbLocks[key] = m
	}
	return m
}