Assistant/cmd/creator-web/server.go

package main

import (
	"embed"
	"encoding/json"
	"fmt"
	"io"
	"io/fs"
	"net/http"
	"net/url"
	"strings"
)

const (
	maxRequestBytes  = 8 << 10
	maxResponseBytes = 1 << 20
)

//go:embed static
var staticFiles embed.FS

type gateway struct {
	upstream *url.URL
	client   *http.Client
	static   http.Handler
}

func newHandler(upstream string, client *http.Client) (http.Handler, error) {
	parsed, err := url.Parse(upstream)
	if err != nil || parsed.Scheme == "" || parsed.Host == "" {
		return nil, fmt.Errorf("invalid core runtime URL")
	}
	assets, err := fs.Sub(staticFiles, "static")
	if err != nil {
		return nil, fmt.Errorf("load embedded console: %w", err)
	}
	if client == nil {
		client = http.DefaultClient
	}
	return securityHeaders(&gateway{upstream: parsed, client: client, static: http.FileServer(http.FS(assets))}), nil
}

func (g *gateway) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path == "/healthz" {
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(map[string]string{"status": "ok", "runtime": "creator-web"})
		return
	}
	if strings.HasPrefix(r.URL.Path, "/api/") {
		g.forwardAPI(w, r)
		return
	}
	if r.Method != http.MethodGet && r.Method != http.MethodHead {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}
	if r.URL.Path == "/" {
		w.Header().Set("Cache-Control", "no-store")
	}
	g.static.ServeHTTP(w, r)
}

func (g *gateway) forwardAPI(w http.ResponseWriter, r *http.Request) {
	upstreamPath, allowed := creatorRoute(r.Method, r.URL.Path)
	if !allowed {
		http.Error(w, "creator route denied", http.StatusNotFound)
		return
	}
	body, err := io.ReadAll(io.LimitReader(r.Body, maxRequestBytes+1))
	if err != nil {
		http.Error(w, "invalid request", http.StatusBadRequest)
		return
	}
	if len(body) > maxRequestBytes {
		http.Error(w, "request too large", http.StatusRequestEntityTooLarge)
		return
	}
	target := *g.upstream
	target.Path = upstreamPath
	target.RawQuery = r.URL.RawQuery
	request, err := http.NewRequestWithContext(r.Context(), r.Method, target.String(), strings.NewReader(string(body)))
	if err != nil {
		http.Error(w, "creator gateway unavailable", http.StatusBadGateway)
		return
	}
	copyRequestHeader(request.Header, r.Header, "Authorization")
	copyRequestHeader(request.Header, r.Header, "Accept")
	copyRequestHeader(request.Header, r.Header, "Content-Type")
	response, err := g.client.Do(request)
	if err != nil {
		http.Error(w, "creator gateway unavailable", http.StatusBadGateway)
		return
	}
	defer response.Body.Close()
	copyResponseHeader(w.Header(), response.Header, "Content-Type")
	w.Header().Set("Cache-Control", "no-store")
	w.WriteHeader(response.StatusCode)
	_, _ = io.Copy(w, io.LimitReader(response.Body, maxResponseBytes))
}

func creatorRoute(method, path string) (string, bool) {
	switch {
	case method == http.MethodPost && path == "/api/session":
		return "/v1/creator/sessions", true
	case method == http.MethodGet && path == "/api/session/current":
		return "/v1/creator/sessions/current", true
	case method == http.MethodGet && path == "/api/focus":
		return "/v1/creator/focus", true
	case method == http.MethodGet && strings.HasPrefix(path, "/api/tasks/"):
		id := strings.TrimPrefix(path, "/api/tasks/")
		if validID(id) {
			return "/v1/creator/tasks/" + id, true
		}
	case method == http.MethodGet && strings.HasPrefix(path, "/api/capabilities/") && strings.HasSuffix(path, "/health"):
		id := strings.TrimSuffix(strings.TrimPrefix(path, "/api/capabilities/"), "/health")
		if validID(id) {
			return "/v1/creator/capabilities/" + id + "/health", true
		}
	}
	return "", false
}

func validID(value string) bool {
	if value == "" || len(value) > 128 {
		return false
	}
	for _, char := range value {
		if (char >= 'a' && char <= 'z') || (char >= 'A' && char <= 'Z') || (char >= '0' && char <= '9') || strings.ContainsRune("._:-", char) {
			continue
		}
		return false
	}
	return true
}

func copyRequestHeader(destination, source http.Header, name string) {
	if value := source.Get(name); value != "" {
		destination.Set(name, value)
	}
}

func copyResponseHeader(destination, source http.Header, name string) {
	if value := source.Get(name); value != "" {
		destination.Set(name, value)
	}
}

func securityHeaders(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Security-Policy", "default-src 'self'; script-src 'self'; style-src 'self'; img-src 'self' data:; connect-src 'self'; frame-ancestors 'none'; base-uri 'none'; form-action 'self'")
		w.Header().Set("Referrer-Policy", "no-referrer")
		w.Header().Set("X-Content-Type-Options", "nosniff")
		w.Header().Set("X-Frame-Options", "DENY")
		next.ServeHTTP(w, r)
	})
}