canon/internal/imgcache/imgcache.go

// Package imgcache downloads and stores artwork locally so pages never
// hotlink third-party image servers.
package imgcache

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"time"
)

const MaxBytes = 15 << 20

var ErrNotImage = errors.New("that file isn't a supported image (JPEG, PNG, WebP, GIF or AVIF)")

type Cache struct {
	dir    string
	client *http.Client
	ua     string
}

func New(dataDir, userAgent string) (*Cache, error) {
	dir := filepath.Join(dataDir, "images")
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return nil, err
	}
	return &Cache{dir: dir, ua: userAgent, client: &http.Client{Timeout: 30 * time.Second}}, nil
}

func (c *Cache) Dir() string { return c.dir }

var nameRe = regexp.MustCompile(`^[0-9a-f]{40}\.(jpg|png|webp|gif|avif)$`)

// ValidName reports whether name looks like a file this cache created.
func ValidName(name string) bool { return nameRe.MatchString(name) }

// Fetch downloads an image URL and returns the cached file name.
func (c *Cache) Fetch(ctx context.Context, url string) (string, error) {
	if !strings.HasPrefix(url, "https://") && !strings.HasPrefix(url, "http://") {
		return "", fmt.Errorf("image URL must start with http:// or https://")
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return "", err
	}
	req.Header.Set("User-Agent", c.ua)
	req.Header.Set("Accept", "image/avif,image/webp,image/png,image/jpeg,image/*;q=0.8")
	resp, err := c.client.Do(req)
	if err != nil {
		return "", fmt.Errorf("downloading image: %w", err)
	}
	defer resp.Body.Close()
	if resp.StatusCode == http.StatusNotFound {
		return "", fmt.Errorf("no artwork found at that address")
	}
	if resp.StatusCode != http.StatusOK {
		return "", fmt.Errorf("downloading image: server returned %d", resp.StatusCode)
	}
	return c.Save(resp.Body)
}

// Save stores image data from r and returns the cached file name. The name
// is derived from the content, so identical images are stored once.
func (c *Cache) Save(r io.Reader) (string, error) {
	data, err := io.ReadAll(io.LimitReader(r, MaxBytes+1))
	if err != nil {
		return "", err
	}
	if len(data) > MaxBytes {
		return "", fmt.Errorf("image is larger than %d MB", MaxBytes>>20)
	}
	ext := sniff(data)
	if ext == "" {
		return "", ErrNotImage
	}
	sum := sha256.Sum256(data)
	name := hex.EncodeToString(sum[:20]) + "." + ext
	path := filepath.Join(c.dir, name)
	if _, err := os.Stat(path); err == nil {
		return name, nil
	}
	tmp, err := os.CreateTemp(c.dir, ".upload-*")
	if err != nil {
		return "", err
	}
	if _, err := tmp.Write(data); err != nil {
		tmp.Close()
		os.Remove(tmp.Name())
		return "", err
	}
	if err := tmp.Close(); err != nil {
		os.Remove(tmp.Name())
		return "", err
	}
	if err := os.Chmod(tmp.Name(), 0o644); err != nil {
		os.Remove(tmp.Name())
		return "", err
	}
	if err := os.Rename(tmp.Name(), path); err != nil {
		os.Remove(tmp.Name())
		return "", err
	}
	return name, nil
}

// Remove deletes a cached image. Missing files are not an error.
func (c *Cache) Remove(name string) error {
	if !ValidName(name) {
		return nil
	}
	err := os.Remove(filepath.Join(c.dir, name))
	if errors.Is(err, os.ErrNotExist) {
		return nil
	}
	return err
}

func sniff(b []byte) string {
	switch http.DetectContentType(b) {
	case "image/jpeg":
		return "jpg"
	case "image/png":
		return "png"
	case "image/webp":
		return "webp"
	case "image/gif":
		return "gif"
	}
	if len(b) > 12 && bytes.Equal(b[4:8], []byte("ftyp")) && (bytes.Equal(b[8:12], []byte("avif")) || bytes.Equal(b[8:12], []byte("avis"))) {
		return "avif"
	}
	return ""
}