Restrict SSE events endpoint to trusted origins

internal/sse/server.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"fmt"
 	"net/http"
+	"net/url"
+	"strings"
 	"time"
 
 	cblog "github.com/charmbracelet/log"
@@ -40,7 +42,11 @@ func (s *Server) events(w http.ResponseWriter, r *http.Request) {
 	w.Header().Set("Content-Type", "text/event-stream")
 	w.Header().Set("Cache-Control", "no-cache")
 	w.Header().Set("Connection", "keep-alive")
-	w.Header().Set("Access-Control-Allow-Origin", "*")
+
+	if !allowOrigin(w, r) {
+		http.Error(w, "forbidden", http.StatusForbidden)
+		return
+	}
 
 	if err := Flush(fl); err != nil {
 		cblog.Errorf("flush headers: %v", err)
@@ -76,3 +82,37 @@ func (s *Server) events(w http.ResponseWriter, r *http.Request) {
 		}
 	}
 }
+
+func allowOrigin(w http.ResponseWriter, r *http.Request) bool {
+	origin := r.Header.Get("Origin")
+	if origin == "" {
+		return true
+	}
+
+	w.Header().Add("Vary", "Origin")
+
+	if origin == "null" {
+		w.Header().Set("Access-Control-Allow-Origin", "null")
+		return true
+	}
+
+	u, err := url.Parse(origin)
+	if err != nil {
+		return false
+	}
+
+	if u.Scheme != "http" && u.Scheme != "https" {
+		return false
+	}
+
+	if u.Host == "" {
+		return false
+	}
+
+	if strings.EqualFold(u.Host, r.Host) {
+		w.Header().Set("Access-Control-Allow-Origin", origin)
+		return true
+	}
+
+	return false
+}

internal/sse/server_test.go

@@ -95,6 +95,29 @@ func (f *flushRecorder) WriteHeader(statusCode int)  {}
 func (f *flushRecorder) Flush()                      { f.count++ }
 func (f *flushRecorder) FlushErr() error             { return nil }
 
+type corsRecorder struct {
+	header http.Header
+	status int
+	body   bytes.Buffer
+}
+
+func newCorsRecorder() *corsRecorder {
+	return &corsRecorder{header: make(http.Header)}
+}
+
+func (c *corsRecorder) Header() http.Header { return c.header }
+
+func (c *corsRecorder) Write(p []byte) (int, error) {
+	if c.status == 0 {
+		c.status = http.StatusOK
+	}
+	return c.body.Write(p)
+}
+
+func (c *corsRecorder) WriteHeader(statusCode int) { c.status = statusCode }
+
+func (c *corsRecorder) Flush() {}
+
 func TestEventsFlushHeaders(t *testing.T) {
 	hub := NewHub()
 	srv := NewServer(":0", hub)
@@ -197,3 +220,100 @@ func TestServerStartAndShutdown(t *testing.T) {
 		t.Fatal("server did not shut down")
 	}
 }
+
+func TestEventsRejectsCrossOrigin(t *testing.T) {
+	hub := NewHub()
+	srv := NewServer(":0", hub)
+
+	req := httptest.NewRequest("GET", "http://127.0.0.1/events", nil)
+	req.Header.Set("Origin", "https://evil.example")
+
+	rr := newCorsRecorder()
+	srv.events(rr, req)
+
+	if rr.status != http.StatusForbidden {
+		t.Fatalf("expected status %d, got %d", http.StatusForbidden, rr.status)
+	}
+	if acao := rr.Header().Get("Access-Control-Allow-Origin"); acao != "" {
+		t.Fatalf("unexpected allow origin header %q", acao)
+	}
+}
+
+func TestEventsAllowsMatchingOrigin(t *testing.T) {
+	hub := NewHub()
+	srv := NewServer(":0", hub)
+
+	req := httptest.NewRequest("GET", "http://localhost:9036/events", nil)
+	req.Header.Set("Origin", "http://localhost:9036")
+	ctx, cancel := context.WithCancel(req.Context())
+	req = req.WithContext(ctx)
+
+	rr := newCorsRecorder()
+	done := make(chan struct{})
+	go func() {
+		srv.events(rr, req)
+		close(done)
+	}()
+
+	time.Sleep(10 * time.Millisecond)
+
+	if acao := rr.Header().Get("Access-Control-Allow-Origin"); acao != "http://localhost:9036" {
+		t.Fatalf("expected allow origin header, got %q", acao)
+	}
+	if vary := rr.Header().Values("Vary"); len(vary) == 0 {
+		t.Fatal("expected vary header")
+	} else {
+		found := false
+		for _, v := range vary {
+			if strings.Contains(v, "Origin") {
+				found = true
+				break
+			}
+		}
+		if !found {
+			t.Fatalf("vary header missing Origin: %v", vary)
+		}
+	}
+
+	cancel()
+
+	select {
+	case <-done:
+	case <-time.After(100 * time.Millisecond):
+		t.Fatal("handler did not exit")
+	}
+}
+
+func TestEventsAllowsNullOrigin(t *testing.T) {
+	hub := NewHub()
+	srv := NewServer(":0", hub)
+
+	req := httptest.NewRequest("GET", "http://127.0.0.1/events", nil)
+	req.Header.Set("Origin", "null")
+	ctx, cancel := context.WithCancel(req.Context())
+	req = req.WithContext(ctx)
+
+	rr := newCorsRecorder()
+	done := make(chan struct{})
+	go func() {
+		srv.events(rr, req)
+		close(done)
+	}()
+
+	time.Sleep(10 * time.Millisecond)
+
+	if acao := rr.Header().Get("Access-Control-Allow-Origin"); acao != "null" {
+		t.Fatalf("expected allow origin header of null, got %q", acao)
+	}
+	if vary := rr.Header().Get("Vary"); !strings.Contains(vary, "Origin") {
+		t.Fatalf("vary header missing Origin: %q", vary)
+	}
+
+	cancel()
+
+	select {
+	case <-done:
+	case <-time.After(100 * time.Millisecond):
+		t.Fatal("handler did not exit")
+	}
+}