vibedit/app.go

package main

import (
	"crypto/rand"
	"crypto/subtle"
	"embed"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"io/fs"
	"log"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"sync"
	"time"
)

//go:embed web
var webFiles embed.FS

// idleExitDelay is how long the program waits after the last editor window
// disconnects before quitting (long enough to survive a page reload).
const idleExitDelay = 10 * time.Second

type App struct {
	tools     *Tools
	media     *MediaStore
	jobs      *JobManager
	hub       *Hub
	mux       *http.ServeMux
	token     string
	cookie    string
	hostOK    map[string]bool // allowed Host headers; nil allows any
	configDir string
	keepAlive bool

	quit     chan struct{}
	quitOnce sync.Once

	mu            sync.Mutex
	clients       int
	everConnected bool

	dialogMu sync.Mutex
}

func newApp(ffmpegPath, ffprobePath string, keepAlive bool) *App {
	a := &App{
		tools:     findTools(ffmpegPath, ffprobePath),
		hub:       newHub(),
		token:     randomHex(16),
		keepAlive: keepAlive,
		quit:      make(chan struct{}),
	}
	a.jobs = newJobManager(a.hub, a.scheduleIdleCheck)
	a.media = newMediaStore(a.tools, a.jobs)
	if dir, err := os.UserConfigDir(); err == nil {
		a.configDir = filepath.Join(dir, "vibedit")
	}
	if a.tools.OK() {
		go a.tools.detectEncoders()
	}
	a.routes()
	return a
}

func (a *App) configureHost(addr *net.TCPAddr) {
	port := fmt.Sprint(addr.Port)
	a.cookie = "vibedit_" + port
	if addr.IP.IsLoopback() {
		a.hostOK = map[string]bool{
			"127.0.0.1:" + port: true,
			"localhost:" + port: true,
			"[::1]:" + port:     true,
		}
	}
}

func (a *App) routes() {
	m := http.NewServeMux()
	static, _ := fs.Sub(webFiles, "web")
	files := http.FileServerFS(static)
	m.HandleFunc("GET /", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Cache-Control", "no-cache")
		files.ServeHTTP(w, r)
	})
	m.HandleFunc("GET /api/info", a.handleInfo)
	m.HandleFunc("GET /api/events", a.handleEvents)
	m.HandleFunc("GET /api/session", a.handleSessionGet)
	m.HandleFunc("POST /api/session", a.handleSessionPut)
	m.HandleFunc("GET /api/fs/list", a.handleFSList)
	m.HandleFunc("POST /api/dialog/open", a.handleDialogOpen)
	m.HandleFunc("POST /api/dialog/save", a.handleDialogSave)
	m.HandleFunc("POST /api/reveal", a.handleReveal)
	m.HandleFunc("POST /api/probe", a.handleProbe)
	m.HandleFunc("POST /api/captions", a.handleCaptions)
	m.HandleFunc("GET /api/media/{id}", a.handleMedia)
	m.HandleFunc("GET /api/thumb", a.handleThumb)
	m.HandleFunc("GET /api/keyframes", a.handleKeyframes)
	m.HandleFunc("POST /api/proxy", a.handleProxy)
	m.HandleFunc("POST /api/export", a.handleExport)
	m.HandleFunc("POST /api/job/cancel", a.handleCancel)
	a.mux = m
}

// ServeHTTP guards every request: only loopback Host headers (DNS rebinding)
// and only browsers holding the session cookie, which is handed out in
// exchange for the random token in the URL printed at startup.
func (a *App) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if a.hostOK != nil && !a.hostOK[r.Host] {
		http.Error(w, "forbidden host", http.StatusForbidden)
		return
	}
	if t := r.URL.Query().Get("t"); r.URL.Path == "/" && t != "" {
		if subtle.ConstantTimeCompare([]byte(t), []byte(a.token)) == 1 {
			http.SetCookie(w, &http.Cookie{Name: a.cookie, Value: a.token, Path: "/", HttpOnly: true, SameSite: http.SameSiteLaxMode})
			http.Redirect(w, r, "/", http.StatusFound)
			return
		}
	}
	c, err := r.Cookie(a.cookie)
	if err != nil || subtle.ConstantTimeCompare([]byte(c.Value), []byte(a.token)) != 1 {
		w.Header().Set("Content-Type", "text/html; charset=utf-8")
		w.WriteHeader(http.StatusForbidden)
		io.WriteString(w, `<!doctype html><title>VibEdit</title><body style="font:16px system-ui;padding:40px;background:#16181d;color:#ddd">`+
			`<h2>This VibEdit link has expired</h2><p>Use the address shown in the VibEdit console window, or restart VibEdit.</p></body>`)
		return
	}
	if r.Method == http.MethodPost {
		if o := r.Header.Get("Origin"); o != "" && !strings.HasSuffix(o, "//"+r.Host) {
			http.Error(w, "forbidden origin", http.StatusForbidden)
			return
		}
	}
	a.mux.ServeHTTP(w, r)
}

func (a *App) requestQuit(reason string) {
	a.quitOnce.Do(func() {
		log.Printf("shutting down (%s)", reason)
		close(a.quit)
	})
}

func (a *App) shutdown() {
	a.jobs.CancelAll()
	a.jobs.Wait(5 * time.Second)
}

func (a *App) clientConnected() {
	a.mu.Lock()
	a.clients++
	a.everConnected = true
	a.mu.Unlock()
}

func (a *App) clientDisconnected() {
	a.mu.Lock()
	a.clients--
	a.mu.Unlock()
	a.scheduleIdleCheck()
}

// scheduleIdleCheck quits the program once no editor window has been
// connected for a while and no export is running.
func (a *App) scheduleIdleCheck() {
	if a.keepAlive {
		return
	}
	time.AfterFunc(idleExitDelay, func() {
		a.mu.Lock()
		idle := a.everConnected && a.clients == 0
		a.mu.Unlock()
		if idle && !a.jobs.AnyRunning("export") {
			a.requestQuit("editor window closed")
		}
	})
}

// --- handlers ---

func (a *App) handleInfo(w http.ResponseWriter, r *http.Request) {
	home, _ := os.UserHomeDir()
	writeJSON(w, map[string]any{
		"version":       version,
		"os":            runtime.GOOS,
		"sep":           string(filepath.Separator),
		"home":          home,
		"ffmpeg":        a.tools.FFmpeg,
		"ffprobe":       a.tools.FFprobe,
		"ffmpegOK":      a.tools.OK(),
		"ffmpegVersion": a.tools.Version,
		"nativeDialogs": nativeDialogs,
		"codecs":        a.tools.CodecList(),
		"codecsReady":   a.tools.EncodersReady(),
		"captionsOK":    a.tools.HasFilter("subtitles"),
		"containers":    containers,
		"copyExts":      copyExts,
		"mediaExts":     mediaExts,
	})
}

func (a *App) handleEvents(w http.ResponseWriter, r *http.Request) {
	fl, ok := w.(http.Flusher)
	if !ok {
		http.Error(w, "streaming unsupported", http.StatusInternalServerError)
		return
	}
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	ch := a.hub.Subscribe()
	defer a.hub.Unsubscribe(ch)
	a.clientConnected()
	defer a.clientDisconnected()

	for _, j := range a.jobs.Snapshot() {
		w.Write(sseMessage("job", j))
	}
	w.Write([]byte(": hello\n\n"))
	fl.Flush()

	ping := time.NewTicker(20 * time.Second)
	defer ping.Stop()
	for {
		var err error
		select {
		case <-r.Context().Done():
			return
		case <-a.quit:
			return
		case msg := <-ch:
			_, err = w.Write(msg)
		case <-ping.C:
			_, err = w.Write([]byte(": ping\n\n"))
		}
		if err != nil {
			return
		}
		fl.Flush()
	}
}

func (a *App) sessionPath() string {
	if a.configDir == "" {
		return ""
	}
	return filepath.Join(a.configDir, "session.json")
}

func (a *App) handleSessionGet(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	if p := a.sessionPath(); p != "" {
		if b, err := os.ReadFile(p); err == nil && json.Valid(b) {
			w.Write(b)
			return
		}
	}
	io.WriteString(w, "{}")
}

func (a *App) handleSessionPut(w http.ResponseWriter, r *http.Request) {
	p := a.sessionPath()
	if p == "" {
		writeError(w, http.StatusInternalServerError, "no config directory available")
		return
	}
	b, err := io.ReadAll(io.LimitReader(r.Body, 20<<20))
	if err != nil || !json.Valid(b) {
		writeError(w, http.StatusBadRequest, "invalid session data")
		return
	}
	if err := writeFileAtomic(p, b); err != nil {
		writeError(w, http.StatusInternalServerError, err.Error())
		return
	}
	writeJSON(w, map[string]any{"ok": true})
}

func (a *App) handleDialogOpen(w http.ResponseWriter, r *http.Request) {
	var req struct {
		Dir  string `json:"dir"`
		Kind string `json:"kind"` // a key of fileKinds
	}
	readJSON(r, &req)
	if !nativeDialogs {
		writeJSON(w, map[string]any{"supported": false})
		return
	}
	if !a.dialogMu.TryLock() {
		writeError(w, http.StatusConflict, "a file dialog is already open")
		return
	}
	defer a.dialogMu.Unlock()
	paths, err := nativeOpenDialog(fileKindOf(req.Kind), req.Dir)
	if err != nil {
		log.Printf("native open dialog failed: %v", err)
		writeJSON(w, map[string]any{"supported": false, "error": err.Error()})
		return
	}
	writeJSON(w, map[string]any{"supported": true, "paths": paths})
}

func (a *App) handleDialogSave(w http.ResponseWriter, r *http.Request) {
	var req struct {
		Path string `json:"path"`
	}
	readJSON(r, &req)
	if !nativeDialogs {
		writeJSON(w, map[string]any{"supported": false})
		return
	}
	if !a.dialogMu.TryLock() {
		writeError(w, http.StatusConflict, "a file dialog is already open")
		return
	}
	defer a.dialogMu.Unlock()
	path, err := nativeSaveDialog("Export as", req.Path)
	if err != nil {
		log.Printf("native save dialog failed: %v", err)
		writeJSON(w, map[string]any{"supported": false, "error": err.Error()})
		return
	}
	writeJSON(w, map[string]any{"supported": true, "path": path})
}

func (a *App) handleReveal(w http.ResponseWriter, r *http.Request) {
	var req struct {
		Path string `json:"path"`
	}
	if err := readJSON(r, &req); err != nil || req.Path == "" {
		writeError(w, http.StatusBadRequest, "missing path")
		return
	}
	if _, err := os.Stat(req.Path); err != nil {
		writeError(w, http.StatusNotFound, "file not found")
		return
	}
	if err := revealInFileManager(req.Path); err != nil {
		writeError(w, http.StatusInternalServerError, err.Error())
		return
	}
	writeJSON(w, map[string]any{"ok": true})
}

func (a *App) handleCancel(w http.ResponseWriter, r *http.Request) {
	var req struct {
		ID string `json:"id"`
	}
	readJSON(r, &req)
	writeJSON(w, map[string]any{"ok": a.jobs.Cancel(req.ID)})
}

// --- helpers ---

func randomHex(n int) string {
	b := make([]byte, n)
	if _, err := rand.Read(b); err != nil {
		panic(err)
	}
	return hex.EncodeToString(b)
}

func writeJSON(w http.ResponseWriter, v any) {
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(v)
}

func writeError(w http.ResponseWriter, code int, msg string) {
	writeErrorData(w, code, msg, nil)
}

func writeErrorData(w http.ResponseWriter, code int, msg string, extra map[string]any) {
	body := map[string]any{"error": msg}
	for k, v := range extra {
		body[k] = v
	}
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(code)
	json.NewEncoder(w).Encode(body)
}

func readJSON(r *http.Request, v any) error {
	return json.NewDecoder(io.LimitReader(r.Body, 20<<20)).Decode(v)
}

func writeFileAtomic(path string, data []byte) error {
	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
		return err
	}
	tmp := path + ".tmp"
	if err := os.WriteFile(tmp, data, 0o644); err != nil {
		return err
	}
	return os.Rename(tmp, path)
}