datascape/cmd/companion/server.go

package main

import (
	"embed"
	"encoding/json"
	"errors"
	"fmt"
	"html/template"
	"net/http"
	"path/filepath"
	"strings"
	"sync"
)

//go:embed config.html
var templates embed.FS

var configTmpl = template.Must(template.ParseFS(templates, "config.html"))

type server struct {
	mu      sync.Mutex
	cfg     *config
	cfgPath string
	port    int
}

func newServer(cfg *config, cfgPath string) *server {
	return &server{cfg: cfg, cfgPath: cfgPath, port: cfg.Port}
}

func (s *server) snapshot() config {
	s.mu.Lock()
	defer s.mu.Unlock()
	c := *s.cfg
	c.AllowedOrigins = append([]string(nil), s.cfg.AllowedOrigins...)
	return c
}

func (s *server) run() error {
	mux := http.NewServeMux()
	mux.HandleFunc("/status", s.handleStatus)
	mux.HandleFunc("/open-file", s.handleOpenFile)
	mux.HandleFunc("/open-folder", s.handleOpenFolder)
	mux.HandleFunc("/config", s.handleConfig)

	addr := fmt.Sprintf("127.0.0.1:%d", s.port)
	return http.ListenAndServe(addr, mux)
}

// methodAllowed enforces a single allowed method, sending 405 otherwise.
func methodAllowed(w http.ResponseWriter, r *http.Request, method string) bool {
	if r.Method == method {
		return true
	}
	w.Header().Set("Allow", method)
	http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
	return false
}

// requireAllowedOrigin checks the Origin header against the allowlist.
// Sets the matching CORS header on success. Returns false (and writes a
// 403) when no match is found.
func (s *server) requireAllowedOrigin(w http.ResponseWriter, r *http.Request) bool {
	origin := r.Header.Get("Origin")
	if origin == "" {
		http.Error(w, "origin required", http.StatusForbidden)
		return false
	}
	cfg := s.snapshot()
	for _, allowed := range cfg.AllowedOrigins {
		if origin == allowed {
			w.Header().Set("Access-Control-Allow-Origin", origin)
			w.Header().Set("Vary", "Origin")
			return true
		}
	}
	http.Error(w, "origin not allowed", http.StatusForbidden)
	return false
}

func (s *server) handleStatus(w http.ResponseWriter, r *http.Request) {
	if !methodAllowed(w, r, http.MethodGet) {
		return
	}
	if !s.requireAllowedOrigin(w, r) {
		return
	}
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(map[string]string{
		"name":    "datascape-companion",
		"version": version,
	})
}

func (s *server) handleOpenFile(w http.ResponseWriter, r *http.Request) {
	s.handleOpen(w, r, false)
}

func (s *server) handleOpenFolder(w http.ResponseWriter, r *http.Request) {
	s.handleOpen(w, r, true)
}

func (s *server) handleOpen(w http.ResponseWriter, r *http.Request, isFolder bool) {
	if !methodAllowed(w, r, http.MethodGet) {
		return
	}
	if !s.requireAllowedOrigin(w, r) {
		return
	}
	wikiPath := r.URL.Query().Get("path")
	cfg := s.snapshot()
	if cfg.WikiRoot == "" {
		writeJSONError(w, http.StatusBadRequest, "wikiRoot is not configured")
		return
	}
	resolved, err := resolveWikiPath(cfg.WikiRoot, wikiPath)
	if err != nil {
		writeJSONError(w, http.StatusBadRequest, err.Error())
		return
	}
	info, err := statPath(resolved)
	if err != nil {
		writeJSONError(w, http.StatusNotFound, "path not found")
		return
	}
	if isFolder && !info.IsDir() {
		// Reveal: if the user asked to open a file's folder, walk up.
		resolved = filepath.Dir(resolved)
	}
	if !isFolder && info.IsDir() {
		writeJSONError(w, http.StatusBadRequest, "path is a directory")
		return
	}
	if err := openOSPath(resolved, isFolder); err != nil {
		writeJSONError(w, http.StatusInternalServerError, err.Error())
		return
	}
	w.WriteHeader(http.StatusNoContent)
}

func (s *server) handleConfig(w http.ResponseWriter, r *http.Request) {
	switch r.Method {
	case http.MethodGet:
		s.renderConfig(w, "")
	case http.MethodPost:
		s.saveConfig(w, r)
	default:
		w.Header().Set("Allow", "GET, POST")
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
	}
}

type configPageData struct {
	WikiRoot       string
	AllowedOrigins string
	Port           int
	ConfigPath     string
	Version        string
	Notice         string
}

func (s *server) renderConfig(w http.ResponseWriter, notice string) {
	cfg := s.snapshot()
	data := configPageData{
		WikiRoot:       cfg.WikiRoot,
		AllowedOrigins: strings.Join(cfg.AllowedOrigins, "\n"),
		Port:           cfg.Port,
		ConfigPath:     s.cfgPath,
		Version:        version,
		Notice:         notice,
	}
	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	if err := configTmpl.Execute(w, data); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

func (s *server) saveConfig(w http.ResponseWriter, r *http.Request) {
	expected := fmt.Sprintf("http://127.0.0.1:%d", s.port)
	if r.Header.Get("Origin") != expected {
		http.Error(w, "origin mismatch", http.StatusForbidden)
		return
	}
	if err := r.ParseForm(); err != nil {
		http.Error(w, "bad form", http.StatusBadRequest)
		return
	}
	wikiRoot := strings.TrimSpace(r.FormValue("wikiRoot"))
	originsRaw := r.FormValue("allowedOrigins")
	var origins []string
	for _, line := range strings.Split(originsRaw, "\n") {
		line = strings.TrimSpace(line)
		line = strings.TrimRight(line, "/")
		if line == "" {
			continue
		}
		origins = append(origins, line)
	}
	if origins == nil {
		origins = []string{}
	}

	s.mu.Lock()
	newCfg := *s.cfg
	newCfg.WikiRoot = wikiRoot
	newCfg.AllowedOrigins = origins
	if err := writeConfigFile(s.cfgPath, &newCfg); err != nil {
		s.mu.Unlock()
		http.Error(w, "save failed: "+err.Error(), http.StatusInternalServerError)
		return
	}
	s.cfg = &newCfg
	s.mu.Unlock()

	s.renderConfig(w, "Saved.")
}

func writeJSONError(w http.ResponseWriter, code int, msg string) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(code)
	_ = json.NewEncoder(w).Encode(map[string]string{"error": msg})
}

// resolveWikiPath joins wikiRoot with a wiki-relative path after rejecting
// absolute paths, traversal segments, and null bytes. The cleaned result
// must remain inside wikiRoot.
func resolveWikiPath(wikiRoot, wikiPath string) (string, error) {
	if strings.ContainsRune(wikiPath, 0) {
		return "", errors.New("invalid path")
	}
	// Reject absolute paths from either family before any cleaning so we
	// don't depend on filepath.IsAbs's per-OS behavior.
	if strings.HasPrefix(wikiPath, "/") || strings.HasPrefix(wikiPath, `\`) ||
		(len(wikiPath) >= 2 && wikiPath[1] == ':') {
		return "", errors.New("absolute path not allowed")
	}
	clean := filepath.Clean(filepath.FromSlash(wikiPath))
	if clean == ".." || strings.HasPrefix(clean, ".."+string(filepath.Separator)) {
		return "", errors.New("path escapes wikiRoot")
	}
	rootAbs, err := filepath.Abs(wikiRoot)
	if err != nil {
		return "", err
	}
	full := filepath.Join(rootAbs, clean)
	rel, err := filepath.Rel(rootAbs, full)
	if err != nil {
		return "", err
	}
	if rel == ".." || strings.HasPrefix(rel, ".."+string(filepath.Separator)) {
		return "", errors.New("path escapes wikiRoot")
	}
	return full, nil
}