byre/internal/server/images.go

package server

import (
	"errors"
	"fmt"
	"io/fs"
	"net/http"
	"os"
	"path/filepath"
	"strings"

	"g.fc.run/ai/byre/internal/api"
)

var worktreeImageTypes = map[string]string{
	".png":  "image/png",
	".jpg":  "image/jpeg",
	".jpeg": "image/jpeg",
	".gif":  "image/gif",
	".webp": "image/webp",
	".svg":  "image/svg+xml",
}

func (s *Server) GetTaskImage(w http.ResponseWriter, r *http.Request, taskID string, params api.GetTaskImageParams) {
	t, ok := s.activeTask(w, r, taskID)
	if !ok {
		return
	}
	serveImage(w, r, t.WorktreePath, params.Path)
}

func (s *Server) GetProjectImage(w http.ResponseWriter, r *http.Request, projectID string, params api.GetProjectImageParams) {
	p, err := s.store.GetProject(r.Context(), projectID)
	if err != nil {
		writeStoreError(w, err)
		return
	}
	serveImage(w, r, p.RepoPath, params.Path)
}

func serveImage(w http.ResponseWriter, r *http.Request, root, rel string) {
	path := rel
	if !filepath.IsAbs(path) {
		path = filepath.Join(root, path)
	}
	real, err := filepath.EvalSymlinks(path)
	if errors.Is(err, fs.ErrNotExist) {
		writeError(w, http.StatusNotFound, fmt.Errorf("%s does not exist", rel))
		return
	}
	if err != nil {
		writeError(w, http.StatusBadRequest, err)
		return
	}

	contentType := worktreeImageTypes[strings.ToLower(filepath.Ext(real))]
	if contentType == "" {
		writeError(w, http.StatusBadRequest, fmt.Errorf("%s is not a PNG, JPEG, GIF, WebP or SVG file", rel))
		return
	}

	f, err := os.Open(real)
	if err != nil {
		writeError(w, http.StatusInternalServerError, err)
		return
	}
	defer f.Close()

	info, err := f.Stat()
	if err != nil {
		writeError(w, http.StatusInternalServerError, err)
		return
	}
	if !info.Mode().IsRegular() {
		writeError(w, http.StatusBadRequest, fmt.Errorf("%s is not a file", rel))
		return
	}

	w.Header().Set("Content-Type", contentType)
	w.Header().Set("X-Content-Type-Options", "nosniff")
	w.Header().Set("Content-Security-Policy", "default-src 'none'; style-src 'unsafe-inline'; sandbox")
	w.Header().Set("Cache-Control", "private, no-cache")
	http.ServeContent(w, r, "", info.ModTime(), f)
}