@samitouri / QOSAMI-X402Tunnel / commits / 1833a14f

feat: add Brotli compression support and related tests

Kim committed Apr 6, 2026 at 18:55 UTC 1833a14f5a4de8021a6fc8c18232154fce0e1d1b
5 files changed +453 -101
go.mod
+1
@@ -4,6 +4,7 @@ go 1.26.1
4
5 require (
6 cloud.google.com/go/compute/metadata v0.9.0
7 + github.com/andybalholm/brotli v1.2.1
8 github.com/aws/aws-sdk-go-v2 v1.41.1
9 github.com/aws/aws-sdk-go-v2/config v1.32.8
10 github.com/aws/aws-sdk-go-v2/credentials v1.19.8
go.sum
+4
@@ -6,6 +6,8 @@ cloud.google.com/go/compute/metadata v0.9.0 h1:pDUj4QMoPejqq20dK0Pg2N4yG9zIkYGdB
6 cloud.google.com/go/compute/metadata v0.9.0/go.mod h1:E0bWwX5wTnLPedCKqk3pJmVgCBSM6qQI1yTBdEb3C10=
7 github.com/ProjectZKM/Ziren/crates/go-runtime/zkvm_runtime v0.0.0-20251001021608-1fe7b43fc4d6 h1:1zYrtlhrZ6/b6SAjLSfKzWtdgqK0U+HtH/VcBWh1BaU=
8 github.com/ProjectZKM/Ziren/crates/go-runtime/zkvm_runtime v0.0.0-20251001021608-1fe7b43fc4d6/go.mod h1:ioLG6R+5bUSO1oeGSDxOV3FADARuMoytZCSX6MEMQkI=
9 +github.com/andybalholm/brotli v1.2.1 h1:R+f5xP285VArJDRgowrfb9DqL18yVK0gKAW/F+eTWro=
10 +github.com/andybalholm/brotli v1.2.1/go.mod h1:rzTDkvFWvIrjDXZHkuS16NPggd91W3kUSvPlQ1pLaKY=
11 github.com/aws/aws-sdk-go-v2 v1.41.1 h1:ABlyEARCDLN034NhxlRUSZr4l71mh+T5KAeGh6cerhU=
12 github.com/aws/aws-sdk-go-v2 v1.41.1/go.mod h1:MayyLB8y+buD9hZqkCW3kX1AKq07Y5pXxtgB+rRFhz0=
13 github.com/aws/aws-sdk-go-v2/config v1.32.8 h1:iu+64gwDKEoKnyTQskSku72dAwggKI5sV6rNvgSMpMs=
@@ -101,6 +103,8 @@ github.com/spruceid/siwe-go v0.2.1 h1:BroySys6CyUzeyNppTseEOT/w56xTdOfcmECTI7rnu
103 github.com/spruceid/siwe-go v0.2.1/go.mod h1:MHpHbptGsM3lHth2L8quhZ9ipiwST8zsJH1CjWpeO1k=
104 github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
105 github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
106 +github.com/xyproto/randomstring v1.0.5 h1:YtlWPoRdgMu3NZtP45drfy1GKoojuR7hmRcnhZqKjWU=
107 +github.com/xyproto/randomstring v1.0.5/go.mod h1:rgmS5DeNXLivK7YprL0pY+lTuhNQW3iGxZ18UQApw/E=
108 go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64=
109 go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y=
110 go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.61.0 h1:q4XOmH/0opmeuJtPsbFNivyl7bCt7yRBbeEm2sC/XtQ=
sdk/expose.go
-101
@@ -472,104 +472,3 @@ func (e *Exposure) RunHTTP(ctx context.Context, handler http.Handler, localAddr
472 }
473 return RunHTTP(ctx, relayListener, handler, localAddr)
474 }
475 -
476 -func RunHTTP(ctx context.Context, relayListener net.Listener, handler http.Handler, localAddr string) error {
477 - localAddr = strings.TrimSpace(localAddr)
478 -
479 - if relayListener == nil && localAddr == "" {
480 - return errors.New("relay listener or local address is required")
481 - }
482 -
483 - var relaySrv *http.Server
484 - if relayListener != nil {
485 - relaySrv = &http.Server{
486 - Handler: handler,
487 - ReadHeaderTimeout: defaultRequestTimeout,
488 - }
489 - }
490 -
491 - var localSrv *http.Server
492 - if localAddr != "" {
493 - localSrv = &http.Server{
494 - Addr: localAddr,
495 - Handler: handler,
496 - ReadHeaderTimeout: defaultRequestTimeout,
497 - }
498 - }
499 -
500 - serverCount := 0
501 - if relaySrv != nil {
502 - serverCount++
503 - }
504 - if localSrv != nil {
505 - serverCount++
506 - }
507 -
508 - results := make(chan error, serverCount)
509 - normalizeServeErr := func(err error, prefix string) error {
510 - if err == nil || errors.Is(err, http.ErrServerClosed) || errors.Is(err, net.ErrClosed) {
511 - return nil
512 - }
513 - return fmt.Errorf("%s: %w", prefix, err)
514 - }
515 -
516 - var (
517 - shutdownOnce sync.Once
518 - shutdownErr error
519 - )
520 - shutdown := func() error {
521 - shutdownOnce.Do(func() {
522 - shutdownCtx, cancel := context.WithTimeout(context.Background(), defaultHTTPShutdownTimeout)
523 - defer cancel()
524 -
525 - var localErr error
526 - if localSrv != nil {
527 - localErr = localSrv.Shutdown(shutdownCtx)
528 - if errors.Is(localErr, http.ErrServerClosed) {
529 - localErr = nil
530 - }
531 - }
532 -
533 - var relayErr error
534 - if relaySrv != nil {
535 - relayErr = relaySrv.Shutdown(shutdownCtx)
536 - if errors.Is(relayErr, http.ErrServerClosed) {
537 - relayErr = nil
538 - }
539 - }
540 -
541 - shutdownErr = errors.Join(localErr, relayErr)
542 - })
543 - return shutdownErr
544 - }
545 -
546 - if localSrv != nil {
547 - go func() {
548 - results <- normalizeServeErr(localSrv.ListenAndServe(), "serve local http")
549 - }()
550 - }
551 - if relaySrv != nil {
552 - go func() {
553 - results <- normalizeServeErr(relaySrv.Serve(relayListener), "serve relay http")
554 - }()
555 - }
556 -
557 - var serveErr error
558 - remaining := serverCount
559 - ctxDone := ctx.Done()
560 - for remaining > 0 {
561 - select {
562 - case err := <-results:
563 - remaining--
564 - if err != nil {
565 - serveErr = errors.Join(serveErr, err)
566 - _ = shutdown()
567 - }
568 - case <-ctxDone:
569 - _ = shutdown()
570 - ctxDone = nil
571 - }
572 - }
573 -
574 - return errors.Join(serveErr, shutdownErr)
575 -}
sdk/http.go new
+322
@@ -0,0 +1,322 @@
1 +package sdk
2 +
3 +import (
4 + "bufio"
5 + "compress/gzip"
6 + "context"
7 + "errors"
8 + "fmt"
9 + "io"
10 + "net"
11 + "net/http"
12 + "strconv"
13 + "strings"
14 + "sync"
15 +
16 + "github.com/andybalholm/brotli"
17 +)
18 +
19 +func RunHTTP(ctx context.Context, relayListener net.Listener, handler http.Handler, localAddr string) error {
20 + if relayListener == nil && localAddr == "" {
21 + return errors.New("relay listener or local address is required")
22 + }
23 +
24 + serverHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
25 + serveCompressedHTTP(handler, w, r)
26 + })
27 +
28 + var relaySrv *http.Server
29 + if relayListener != nil {
30 + relaySrv = &http.Server{
31 + Handler: serverHandler,
32 + ReadHeaderTimeout: defaultRequestTimeout,
33 + }
34 + }
35 +
36 + var localSrv *http.Server
37 + if localAddr != "" {
38 + localSrv = &http.Server{
39 + Addr: localAddr,
40 + Handler: serverHandler,
41 + ReadHeaderTimeout: defaultRequestTimeout,
42 + }
43 + }
44 +
45 + serverCount := 0
46 + if relaySrv != nil {
47 + serverCount++
48 + }
49 + if localSrv != nil {
50 + serverCount++
51 + }
52 +
53 + results := make(chan error, serverCount)
54 + normalizeServeErr := func(err error, prefix string) error {
55 + if err == nil || errors.Is(err, http.ErrServerClosed) || errors.Is(err, net.ErrClosed) {
56 + return nil
57 + }
58 + return fmt.Errorf("%s: %w", prefix, err)
59 + }
60 +
61 + var (
62 + shutdownOnce sync.Once
63 + shutdownErr error
64 + )
65 + shutdown := func() error {
66 + shutdownOnce.Do(func() {
67 + shutdownCtx, cancel := context.WithTimeout(context.Background(), defaultHTTPShutdownTimeout)
68 + defer cancel()
69 +
70 + var localErr error
71 + if localSrv != nil {
72 + localErr = localSrv.Shutdown(shutdownCtx)
73 + if errors.Is(localErr, http.ErrServerClosed) {
74 + localErr = nil
75 + }
76 + }
77 +
78 + var relayErr error
79 + if relaySrv != nil {
80 + relayErr = relaySrv.Shutdown(shutdownCtx)
81 + if errors.Is(relayErr, http.ErrServerClosed) {
82 + relayErr = nil
83 + }
84 + }
85 +
86 + shutdownErr = errors.Join(localErr, relayErr)
87 + })
88 + return shutdownErr
89 + }
90 +
91 + if localSrv != nil {
92 + go func() {
93 + results <- normalizeServeErr(localSrv.ListenAndServe(), "serve local http")
94 + }()
95 + }
96 + if relaySrv != nil {
97 + go func() {
98 + results <- normalizeServeErr(relaySrv.Serve(relayListener), "serve relay http")
99 + }()
100 + }
101 +
102 + var serveErr error
103 + remaining := serverCount
104 + ctxDone := ctx.Done()
105 + for remaining > 0 {
106 + select {
107 + case err := <-results:
108 + remaining--
109 + if err != nil {
110 + serveErr = errors.Join(serveErr, err)
111 + _ = shutdown()
112 + }
113 + case <-ctxDone:
114 + _ = shutdown()
115 + ctxDone = nil
116 + }
117 + }
118 +
119 + return errors.Join(serveErr, shutdownErr)
120 +}
121 +
122 +func serveCompressedHTTP(handler http.Handler, w http.ResponseWriter, r *http.Request) {
123 + if handler == nil {
124 + http.NotFound(w, r)
125 + return
126 + }
127 +
128 + format := ""
129 + if r != nil {
130 + parseQuality := func(params string) float64 {
131 + for param := range strings.SplitSeq(params, ";") {
132 + key, value, ok := strings.Cut(strings.TrimSpace(param), "=")
133 + if !ok || !strings.EqualFold(strings.TrimSpace(key), "q") {
134 + continue
135 + }
136 +
137 + q, err := strconv.ParseFloat(strings.TrimSpace(value), 64)
138 + if err != nil || q < 0 {
139 + return 0
140 + }
141 + if q > 1 {
142 + return 1
143 + }
144 + return q
145 + }
146 + return 1
147 + }
148 +
149 + bestQ := 0.0
150 + for rawPart := range strings.SplitSeq(r.Header.Get("Accept-Encoding"), ",") {
151 + part := strings.TrimSpace(strings.ToLower(rawPart))
152 + if part == "" {
153 + continue
154 + }
155 +
156 + name, params, _ := strings.Cut(part, ";")
157 + candidate := strings.TrimSpace(name)
158 + if candidate != "br" && candidate != "gzip" {
159 + continue
160 + }
161 +
162 + q := parseQuality(params)
163 + if q <= 0 {
164 + continue
165 + }
166 +
167 + if q > bestQ || (q == bestQ && candidate == "br") {
168 + format = candidate
169 + bestQ = q
170 + }
171 + }
172 + }
173 + if format == "" || strings.TrimSpace(r.Header.Get("Range")) != "" {
174 + handler.ServeHTTP(w, r)
175 + return
176 + }
177 + if headerContainsToken(r.Header.Values("Connection"), "upgrade") && strings.TrimSpace(r.Header.Get("Upgrade")) != "" {
178 + handler.ServeHTTP(w, r)
179 + return
180 + }
181 +
182 + writer := &compressedResponseWriter{
183 + ResponseWriter: w,
184 + format: format,
185 + }
186 + defer func() {
187 + _ = writer.Close()
188 + }()
189 +
190 + handler.ServeHTTP(writer, r)
191 +}
192 +
193 +func headerContainsToken(values []string, target string) bool {
194 + target = strings.ToLower(strings.TrimSpace(target))
195 + for _, value := range values {
196 + for _, part := range strings.Split(value, ",") {
197 + if strings.ToLower(strings.TrimSpace(part)) == target {
198 + return true
199 + }
200 + }
201 + }
202 + return false
203 +}
204 +
205 +type compressedResponseWriter struct {
206 + http.ResponseWriter
207 + format string
208 + writer io.WriteCloser
209 + flushWriter func() error
210 + wroteHeader bool
211 + passthrough bool
212 +}
213 +
214 +func (w *compressedResponseWriter) WriteHeader(statusCode int) {
215 + if w.wroteHeader {
216 + return
217 + }
218 + w.wroteHeader = true
219 +
220 + header := w.Header()
221 + contentType, _, _ := strings.Cut(strings.ToLower(strings.TrimSpace(header.Get("Content-Type"))), ";")
222 + contentType = strings.TrimSpace(contentType)
223 + compressible := strings.HasPrefix(contentType, "text/")
224 + switch contentType {
225 + case "application/json", "application/javascript", "application/xml", "image/svg+xml":
226 + compressible = true
227 + }
228 + smallResponse := false
229 + if contentLength := strings.TrimSpace(header.Get("Content-Length")); contentLength != "" {
230 + if n, err := strconv.ParseInt(contentLength, 10, 64); err == nil && n >= 0 && n <= 1024 {
231 + smallResponse = true
232 + }
233 + }
234 + switch {
235 + case statusCode >= 100 && statusCode < 200:
236 + w.passthrough = true
237 + case statusCode == http.StatusNoContent || statusCode == http.StatusNotModified:
238 + w.passthrough = true
239 + case !compressible:
240 + w.passthrough = true
241 + case smallResponse:
242 + w.passthrough = true
243 + case strings.TrimSpace(header.Get("Content-Encoding")) != "":
244 + w.passthrough = true
245 + case strings.TrimSpace(header.Get("Content-Range")) != "":
246 + w.passthrough = true
247 + case strings.HasPrefix(contentType, "text/event-stream"):
248 + w.passthrough = true
249 + case headerContainsToken(header.Values("Cache-Control"), "no-transform"):
250 + w.passthrough = true
251 + }
252 + if w.passthrough {
253 + w.ResponseWriter.WriteHeader(statusCode)
254 + return
255 + }
256 +
257 + switch w.format {
258 + case "br":
259 + writer := brotli.NewWriter(w.ResponseWriter)
260 + w.writer = writer
261 + w.flushWriter = writer.Flush
262 + case "gzip":
263 + writer := gzip.NewWriter(w.ResponseWriter)
264 + w.writer = writer
265 + w.flushWriter = writer.Flush
266 + default:
267 + w.passthrough = true
268 + w.ResponseWriter.WriteHeader(statusCode)
269 + return
270 + }
271 +
272 + header.Del("Content-Length")
273 + header.Set("Content-Encoding", w.format)
274 + if !headerContainsToken(header.Values("Vary"), "accept-encoding") {
275 + header.Add("Vary", "Accept-Encoding")
276 + }
277 + w.ResponseWriter.WriteHeader(statusCode)
278 +}
279 +
280 +func (w *compressedResponseWriter) Write(p []byte) (int, error) {
281 + if !w.wroteHeader {
282 + header := w.Header()
283 + if strings.TrimSpace(header.Get("Content-Type")) == "" && len(p) > 0 {
284 + header.Set("Content-Type", http.DetectContentType(p))
285 + }
286 + w.WriteHeader(http.StatusOK)
287 + }
288 + if w.passthrough {
289 + return w.ResponseWriter.Write(p)
290 + }
291 + return w.writer.Write(p)
292 +}
293 +
294 +func (w *compressedResponseWriter) Flush() {
295 + if !w.wroteHeader {
296 + w.WriteHeader(http.StatusOK)
297 + }
298 + if !w.passthrough && w.flushWriter != nil {
299 + _ = w.flushWriter()
300 + }
301 + if flusher, ok := w.ResponseWriter.(http.Flusher); ok {
302 + flusher.Flush()
303 + }
304 +}
305 +
306 +func (w *compressedResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
307 + hijacker, ok := w.ResponseWriter.(http.Hijacker)
308 + if !ok {
309 + return nil, nil, http.ErrNotSupported
310 + }
311 + return hijacker.Hijack()
312 +}
313 +
314 +func (w *compressedResponseWriter) Close() error {
315 + if w.writer == nil {
316 + return nil
317 + }
318 + err := w.writer.Close()
319 + w.writer = nil
320 + w.flushWriter = nil
321 + return err
322 +}
sdk/http_test.go new
+126
@@ -0,0 +1,126 @@
1 +package sdk
2 +
3 +import (
4 + "net/http"
5 + "net/http/httptest"
6 + "testing"
7 +)
8 +
9 +func TestServeCompressedHTTPChoosesAcceptedEncoding(t *testing.T) {
10 + t.Parallel()
11 +
12 + tests := []struct {
13 + name string
14 + acceptEncoding string
15 + want string
16 + }{
17 + {name: "missing header", want: ""},
18 + {name: "unsupported encoding only", acceptEncoding: "deflate", want: ""},
19 + {name: "gzip accepted", acceptEncoding: "gzip", want: "gzip"},
20 + {name: "brotli preferred on tie", acceptEncoding: "gzip, br", want: "br"},
21 + {name: "quality chooses gzip", acceptEncoding: "gzip;q=1, br;q=0.5", want: "gzip"},
22 + {name: "zero quality disables format", acceptEncoding: "gzip;q=0, br;q=0", want: ""},
23 + {name: "wildcard ignored", acceptEncoding: "*", want: ""},
24 + }
25 +
26 + for _, tt := range tests {
27 + t.Run(tt.name, func(t *testing.T) {
28 + t.Parallel()
29 +
30 + req := httptest.NewRequest("GET", "/", nil)
31 + if tt.acceptEncoding != "" {
32 + req.Header.Set("Accept-Encoding", tt.acceptEncoding)
33 + }
34 + rec := httptest.NewRecorder()
35 +
36 + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
37 + w.Header().Set("Content-Type", "text/plain; charset=utf-8")
38 + w.Header().Set("Content-Length", "2048")
39 + _, _ = w.Write([]byte("hello world"))
40 + })
41 +
42 + serveCompressedHTTP(handler, rec, req)
43 +
44 + if got := rec.Header().Get("Content-Encoding"); got != tt.want {
45 + t.Fatalf("Content-Encoding = %q, want %q", got, tt.want)
46 + }
47 + })
48 + }
49 +}
50 +
51 +func TestServeCompressedHTTPCompressesTextResponses(t *testing.T) {
52 + t.Parallel()
53 +
54 + req := httptest.NewRequest("GET", "/", nil)
55 + req.Header.Set("Accept-Encoding", "gzip")
56 + rec := httptest.NewRecorder()
57 +
58 + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
59 + w.Header().Set("Content-Type", "text/plain; charset=utf-8")
60 + _, _ = w.Write([]byte("hello world"))
61 + })
62 +
63 + serveCompressedHTTP(handler, rec, req)
64 +
65 + if got := rec.Header().Get("Content-Encoding"); got != "gzip" {
66 + t.Fatalf("Content-Encoding = %q, want gzip", got)
67 + }
68 +}
69 +
70 +func TestServeCompressedHTTPBypassesBinaryResponses(t *testing.T) {
71 + t.Parallel()
72 +
73 + req := httptest.NewRequest("GET", "/", nil)
74 + req.Header.Set("Accept-Encoding", "gzip")
75 + rec := httptest.NewRecorder()
76 +
77 + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
78 + w.Header().Set("Content-Type", "image/png")
79 + _, _ = w.Write([]byte("not-really-a-png"))
80 + })
81 +
82 + serveCompressedHTTP(handler, rec, req)
83 +
84 + if got := rec.Header().Get("Content-Encoding"); got != "" {
85 + t.Fatalf("Content-Encoding = %q, want empty", got)
86 + }
87 +}
88 +
89 +func TestServeCompressedHTTPBypassesSmallResponsesWithContentLength(t *testing.T) {
90 + t.Parallel()
91 +
92 + req := httptest.NewRequest("GET", "/", nil)
93 + req.Header.Set("Accept-Encoding", "gzip")
94 + rec := httptest.NewRecorder()
95 +
96 + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
97 + w.Header().Set("Content-Type", "application/json")
98 + w.Header().Set("Content-Length", "12")
99 + _, _ = w.Write([]byte(`{"ok":true}`))
100 + })
101 +
102 + serveCompressedHTTP(handler, rec, req)
103 +
104 + if got := rec.Header().Get("Content-Encoding"); got != "" {
105 + t.Fatalf("Content-Encoding = %q, want empty", got)
106 + }
107 +}
108 +
109 +func TestServeCompressedHTTPIgnoresSmallThresholdWithoutContentLength(t *testing.T) {
110 + t.Parallel()
111 +
112 + req := httptest.NewRequest("GET", "/", nil)
113 + req.Header.Set("Accept-Encoding", "gzip")
114 + rec := httptest.NewRecorder()
115 +
116 + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
117 + w.Header().Set("Content-Type", "application/json")
118 + _, _ = w.Write([]byte(`{"ok":true}`))
119 + })
120 +
121 + serveCompressedHTTP(handler, rec, req)
122 +
123 + if got := rec.Header().Get("Content-Encoding"); got != "gzip" {
124 + t.Fatalf("Content-Encoding = %q, want gzip", got)
125 + }
126 +}