From 2db34e600c0199ffa0a25d98815eff137c4eadf2 Mon Sep 17 00:00:00 2001 From: efiten Date: Thu, 30 Apr 2026 16:21:25 +0200 Subject: [PATCH] feat(server): add opt-in gzip and WebSocket compression config Adds a `compression` config block (disabled by default) to let operators enable HTTP gzip and WebSocket permessage-deflate when sitting behind a reverse proxy that does not already compress upstream responses. "compression": { "gzip": true, "websocket": true } Co-Authored-By: Claude Sonnet 4.6 --- cmd/server/compress.go | 46 +++++++++++++++ cmd/server/compress_test.go | 109 ++++++++++++++++++++++++++++++++++++ cmd/server/config.go | 18 ++++++ cmd/server/main.go | 11 +++- cmd/server/websocket.go | 18 +++--- 5 files changed, 192 insertions(+), 10 deletions(-) create mode 100644 cmd/server/compress.go create mode 100644 cmd/server/compress_test.go diff --git a/cmd/server/compress.go b/cmd/server/compress.go new file mode 100644 index 00000000..f0ae50fc --- /dev/null +++ b/cmd/server/compress.go @@ -0,0 +1,46 @@ +package main + +import ( + "compress/gzip" + "net/http" + "strings" +) + +type gzipResponseWriter struct { + http.ResponseWriter + gz *gzip.Writer + wrote bool +} + +func (g *gzipResponseWriter) WriteHeader(code int) { + g.ResponseWriter.Header().Del("Content-Length") + g.ResponseWriter.WriteHeader(code) +} + +func (g *gzipResponseWriter) Write(b []byte) (int, error) { + if !g.wrote { + g.wrote = true + g.ResponseWriter.Header().Del("Content-Length") + } + return g.gz.Write(b) +} + +// gzipMiddleware compresses HTTP responses when the client supports gzip. +// WebSocket upgrade requests are passed through unmodified. +func gzipMiddleware(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if !strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") { + next.ServeHTTP(w, r) + return + } + if strings.EqualFold(r.Header.Get("Upgrade"), "websocket") { + next.ServeHTTP(w, r) + return + } + w.Header().Set("Content-Encoding", "gzip") + w.Header().Set("Vary", "Accept-Encoding") + gz, _ := gzip.NewWriterLevel(w, gzip.DefaultCompression) + defer gz.Close() + next.ServeHTTP(&gzipResponseWriter{ResponseWriter: w, gz: gz}, r) + }) +} diff --git a/cmd/server/compress_test.go b/cmd/server/compress_test.go new file mode 100644 index 00000000..1301ce60 --- /dev/null +++ b/cmd/server/compress_test.go @@ -0,0 +1,109 @@ +package main + +import ( + "compress/gzip" + "io" + "net/http" + "net/http/httptest" + "testing" +) + +func TestCompressionConfigDefaults(t *testing.T) { + cfg := &Config{} + if cfg.GZipEnabled() { + t.Error("GZipEnabled should be false when compression is nil") + } + if cfg.WSCompressionEnabled() { + t.Error("WSCompressionEnabled should be false when compression is nil") + } +} + +func TestCompressionConfigExplicitFalse(t *testing.T) { + cfg := &Config{Compression: &CompressionConfig{GZip: false, Websocket: false}} + if cfg.GZipEnabled() { + t.Error("GZipEnabled should be false") + } + if cfg.WSCompressionEnabled() { + t.Error("WSCompressionEnabled should be false") + } +} + +func TestCompressionConfigEnabled(t *testing.T) { + cfg := &Config{Compression: &CompressionConfig{GZip: true, Websocket: true}} + if !cfg.GZipEnabled() { + t.Error("GZipEnabled should be true") + } + if !cfg.WSCompressionEnabled() { + t.Error("WSCompressionEnabled should be true") + } +} + +func TestGZipMiddlewareCompresses(t *testing.T) { + body := `{"nodes":[{"id":"abc"}]}` + handler := gzipMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(body)) + })) + + req := httptest.NewRequest("GET", "/api/nodes", nil) + req.Header.Set("Accept-Encoding", "gzip") + rr := httptest.NewRecorder() + handler.ServeHTTP(rr, req) + + if rr.Header().Get("Content-Encoding") != "gzip" { + t.Errorf("expected Content-Encoding: gzip, got %q", rr.Header().Get("Content-Encoding")) + } + if rr.Header().Get("Vary") != "Accept-Encoding" { + t.Errorf("expected Vary: Accept-Encoding, got %q", rr.Header().Get("Vary")) + } + gz, err := gzip.NewReader(rr.Body) + if err != nil { + t.Fatalf("response is not valid gzip: %v", err) + } + defer gz.Close() + decoded, err := io.ReadAll(gz) + if err != nil { + t.Fatalf("reading gzip: %v", err) + } + if string(decoded) != body { + t.Errorf("decompressed body = %q, want %q", string(decoded), body) + } +} + +func TestGZipMiddlewareSkipsNoAcceptEncoding(t *testing.T) { + handler := gzipMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Write([]byte("hello")) + })) + + req := httptest.NewRequest("GET", "/api/nodes", nil) + rr := httptest.NewRecorder() + handler.ServeHTTP(rr, req) + + if rr.Header().Get("Content-Encoding") != "" { + t.Errorf("expected no Content-Encoding, got %q", rr.Header().Get("Content-Encoding")) + } + if rr.Body.String() != "hello" { + t.Errorf("expected plain body, got %q", rr.Body.String()) + } +} + +func TestGZipMiddlewareSkipsWebSocket(t *testing.T) { + called := false + handler := gzipMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + called = true + w.Write([]byte("ws")) + })) + + req := httptest.NewRequest("GET", "/ws", nil) + req.Header.Set("Accept-Encoding", "gzip") + req.Header.Set("Upgrade", "websocket") + rr := httptest.NewRecorder() + handler.ServeHTTP(rr, req) + + if !called { + t.Error("expected next handler to be called") + } + if rr.Header().Get("Content-Encoding") != "" { + t.Errorf("WebSocket should not be gzip-encoded, got %q", rr.Header().Get("Content-Encoding")) + } +} diff --git a/cmd/server/config.go b/cmd/server/config.go index 6039d41e..6931e150 100644 --- a/cmd/server/config.go +++ b/cmd/server/config.go @@ -70,6 +70,7 @@ type Config struct { DebugAffinity bool `json:"debugAffinity,omitempty"` + Compression *CompressionConfig `json:"compression,omitempty"` ResolvedPath *ResolvedPathConfig `json:"resolvedPath,omitempty"` NeighborGraph *NeighborGraphConfig `json:"neighborGraph,omitempty"` } @@ -102,6 +103,23 @@ func IsWeakAPIKey(key string) bool { return false } +// CompressionConfig controls HTTP gzip and WebSocket permessage-deflate compression. +// Both are disabled by default — enable only when the upstream proxy does not already compress. +type CompressionConfig struct { + GZip bool `json:"gzip"` + Websocket bool `json:"websocket"` +} + +// GZipEnabled returns true when HTTP gzip compression is explicitly enabled. +func (c *Config) GZipEnabled() bool { + return c.Compression != nil && c.Compression.GZip +} + +// WSCompressionEnabled returns true when WebSocket permessage-deflate is explicitly enabled. +func (c *Config) WSCompressionEnabled() bool { + return c.Compression != nil && c.Compression.Websocket +} + // ResolvedPathConfig controls async backfill behavior. type ResolvedPathConfig struct { BackfillHours int `json:"backfillHours"` // how far back (hours) to scan for NULL resolved_path (default 24) diff --git a/cmd/server/main.go b/cmd/server/main.go index 22dc600e..dbcf1a84 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -232,6 +232,7 @@ func main() { // WebSocket hub hub := NewHub() + hub.upgrader.EnableCompression = cfg.WSCompressionEnabled() // HTTP server srv := NewServer(database, cfg, hub) @@ -404,9 +405,17 @@ func main() { } // Graceful shutdown + var handler http.Handler = router + if cfg.GZipEnabled() { + handler = gzipMiddleware(router) + log.Printf("[server] HTTP gzip compression enabled") + } + if cfg.WSCompressionEnabled() { + log.Printf("[server] WebSocket permessage-deflate compression enabled") + } httpServer := &http.Server{ Addr: fmt.Sprintf(":%d", cfg.Port), - Handler: router, + Handler: handler, ReadTimeout: 30 * time.Second, WriteTimeout: 60 * time.Second, IdleTimeout: 120 * time.Second, diff --git a/cmd/server/websocket.go b/cmd/server/websocket.go index 713ebe73..c9ba1aae 100644 --- a/cmd/server/websocket.go +++ b/cmd/server/websocket.go @@ -11,16 +11,11 @@ import ( "github.com/gorilla/websocket" ) -var upgrader = websocket.Upgrader{ - ReadBufferSize: 1024, - WriteBufferSize: 4096, - CheckOrigin: func(r *http.Request) bool { return true }, -} - // Hub manages WebSocket clients and broadcasts. type Hub struct { - mu sync.RWMutex - clients map[*Client]bool + mu sync.RWMutex + clients map[*Client]bool + upgrader websocket.Upgrader } // Client is a single WebSocket connection. @@ -33,6 +28,11 @@ type Client struct { func NewHub() *Hub { return &Hub{ clients: make(map[*Client]bool), + upgrader: websocket.Upgrader{ + ReadBufferSize: 1024, + WriteBufferSize: 4096, + CheckOrigin: func(r *http.Request) bool { return true }, + }, } } @@ -95,7 +95,7 @@ func (h *Hub) Broadcast(msg interface{}) { // ServeWS handles the WebSocket upgrade and runs the client. func (h *Hub) ServeWS(w http.ResponseWriter, r *http.Request) { - conn, err := upgrader.Upgrade(w, r, nil) + conn, err := h.upgrader.Upgrade(w, r, nil) if err != nil { log.Printf("[ws] upgrade error: %v", err) return