agentmon/cmd/web-ui/main.go

package main

import (
	"embed"
	"io"
	"io/fs"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"strings"
	"sync"

	"github.com/gorilla/websocket"
)

//go:embed static
var staticFiles embed.FS

var wsUpgrader = websocket.Upgrader{
	// Only accept WebSocket upgrades whose Origin matches the Host header.
	// This blocks CSWSH: a third-party page can't open the live event
	// firehose just because the user has this UI in another tab.
	CheckOrigin: func(r *http.Request) bool {
		origin := r.Header.Get("Origin")
		if origin == "" {
			return false
		}
		u, err := url.Parse(origin)
		if err != nil || u.Host == "" {
			return false
		}
		return u.Host == r.Host
	},
}

// securityHeaders sets sensible defaults on every response. The CSP allows
// the external CDNs currently referenced from index.html (fonts, uPlot).
// Remove 'unsafe-inline' from script-src once the theme-init snippet in
// index.html is moved to a file.
func securityHeaders(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		h := w.Header()
		h.Set("X-Content-Type-Options", "nosniff")
		h.Set("X-Frame-Options", "DENY")
		h.Set("Referrer-Policy", "no-referrer")
		h.Set("Content-Security-Policy", strings.Join([]string{
			"default-src 'self'",
			"script-src 'self' https://cdn.jsdelivr.net 'unsafe-inline'",
			"style-src 'self' https://fonts.googleapis.com https://cdn.jsdelivr.net 'unsafe-inline'",
			"font-src 'self' https://fonts.gstatic.com",
			"img-src 'self' data:",
			"connect-src 'self' ws: wss:",
			"frame-ancestors 'none'",
			"base-uri 'self'",
			"form-action 'self'",
		}, "; "))
		next.ServeHTTP(w, r)
	})
}

func main() {
	addr := envDefault("AGENTMON_UI_ADDR", ":8082")
	queryAPIBase := envDefault("AGENTMON_QUERY_BASE", "http://query-api")

	queryURL, err := url.Parse(queryAPIBase)
	if err != nil {
		log.Fatalf("invalid AGENTMON_QUERY_BASE %q: %v", queryAPIBase, err)
	}
	if queryURL.Host == "" {
		log.Fatalf("AGENTMON_QUERY_BASE %q has no host", queryAPIBase)
	}

	mux := http.NewServeMux()

	// Health check
	mux.HandleFunc("/healthz", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte("ok"))
	})

	// API proxy to query-api
	proxy := httputil.NewSingleHostReverseProxy(queryURL)
	mux.HandleFunc("/api/", func(w http.ResponseWriter, r *http.Request) {
		r.URL.Path = strings.TrimPrefix(r.URL.Path, "/api")
		r.Host = queryURL.Host
		proxy.ServeHTTP(w, r)
	})

	// WebSocket proxy to query-api
	mux.HandleFunc("/api/v1/ws", func(w http.ResponseWriter, r *http.Request) {
		// Upgrade first so CheckOrigin runs before we touch the upstream.
		// On rejection, Upgrade has already written the response; just bail.
		uiConn, err := wsUpgrader.Upgrade(w, r, nil)
		if err != nil {
			return
		}
		defer uiConn.Close()

		queryWSURL := "ws://" + queryURL.Host + "/v1/ws"
		conn, _, err := websocket.DefaultDialer.Dial(queryWSURL, nil)
		if err != nil {
			_ = uiConn.WriteMessage(websocket.CloseMessage,
				websocket.FormatCloseMessage(websocket.CloseTryAgainLater, "upstream unavailable"))
			return
		}
		defer conn.Close()

		var uiMu, upstreamMu sync.Mutex

		// upstream → UI
		done := make(chan struct{})
		go func() {
			defer close(done)
			for {
				_, msg, err := conn.ReadMessage()
				if err != nil {
					break
				}
				uiMu.Lock()
				err = uiConn.WriteMessage(websocket.TextMessage, msg)
				uiMu.Unlock()
				if err != nil {
					break
				}
			}
		}()

		// UI → upstream
		for {
			_, msg, err := uiConn.ReadMessage()
			if err != nil {
				break
			}
			upstreamMu.Lock()
			err = conn.WriteMessage(websocket.TextMessage, msg)
			upstreamMu.Unlock()
			if err != nil {
				break
			}
		}
		<-done
	})

	// Static files
	staticFS, _ := fs.Sub(staticFiles, "static")
	fileServer := http.FileServer(http.FS(staticFS))

	mux.HandleFunc("/favicon.svg", func(w http.ResponseWriter, r *http.Request) {
		r.URL.Path = "/favicon.svg"
		fileServer.ServeHTTP(w, r)
	})

	mux.HandleFunc("/favicon.ico", func(w http.ResponseWriter, r *http.Request) {
		http.Redirect(w, r, "/favicon.svg", http.StatusMovedPermanently)
	})

	mux.HandleFunc("/static/", func(w http.ResponseWriter, r *http.Request) {
		r.URL.Path = strings.TrimPrefix(r.URL.Path, "/static")
		fileServer.ServeHTTP(w, r)
	})

	// SPA catch-all: serve index.html for routes without file extensions
	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		// If the path contains a dot, it's likely a missing static asset
		if strings.Contains(r.URL.Path, ".") {
			http.NotFound(w, r)
			return
		}

		f, err := staticFiles.Open("static/index.html")
		if err != nil {
			http.Error(w, "index.html not found", http.StatusInternalServerError)
			return
		}
		defer f.Close()

		w.Header().Set("Content-Type", "text/html; charset=utf-8")
		_, _ = io.Copy(w, f)
	})

	log.Printf("web-ui listening on %s", addr)
	log.Fatal(http.ListenAndServe(addr, securityHeaders(mux)))
}

func envDefault(key, def string) string {
	if v := os.Getenv(key); v != "" {
		return v
	}
	return def
}