| 1 | package serve |
| 2 | |
| 3 | import ( |
| 4 | "bytes" |
| 5 | "compress/gzip" |
| 6 | "io" |
| 7 | "net/http" |
| 8 | "net/http/httptest" |
| 9 | "strings" |
| 10 | "testing" |
| 11 | ) |
| 12 | |
| 13 | func TestGzipMiddleware(t *testing.T) { |
| 14 | big := []byte(`{"data":"` + strings.Repeat("x", 8192) + `"}`) |
| 15 | handler := gzipMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 16 | if r.URL.Path == "/events" { |
| 17 | w.Header().Set("Content-Type", "text/event-stream") |
| 18 | _, _ = w.Write([]byte("data: {}\n\n")) |
| 19 | return |
| 20 | } |
| 21 | w.Header().Set("ETag", `"abc"`) |
| 22 | if r.Header.Get("If-None-Match") == `"abc"` { |
| 23 | w.WriteHeader(http.StatusNotModified) |
| 24 | return |
| 25 | } |
| 26 | w.Header().Set("Content-Type", "application/json") |
| 27 | _, _ = w.Write(big) |
| 28 | })) |
| 29 | srv := httptest.NewServer(handler) |
| 30 | defer srv.Close() |
| 31 | |
| 32 | req, _ := http.NewRequest(http.MethodGet, srv.URL+"/history", nil) |
| 33 | req.Header.Set("Accept-Encoding", "gzip") |
| 34 | resp, err := http.DefaultClient.Do(req) |
| 35 | if err != nil { |
| 36 | t.Fatal(err) |
| 37 | } |
| 38 | defer resp.Body.Close() |
| 39 | if got := resp.Header.Get("Content-Encoding"); got != "gzip" { |
| 40 | t.Fatalf("Content-Encoding = %q, want gzip", got) |
| 41 | } |
| 42 | if got := resp.Header.Get("ETag"); got != `"abc"` { |
| 43 | t.Fatalf("ETag = %q, want preserved", got) |
| 44 | } |
| 45 | zr, err := gzip.NewReader(resp.Body) |
| 46 | if err != nil { |
| 47 | t.Fatal(err) |
| 48 | } |
| 49 | plain, err := io.ReadAll(zr) |
| 50 | if err != nil { |
| 51 | t.Fatal(err) |
| 52 | } |
| 53 | if !bytes.Equal(plain, big) { |
| 54 | t.Fatal("decompressed body mismatch") |
| 55 | } |
| 56 | |
| 57 | req, _ = http.NewRequest(http.MethodGet, srv.URL+"/history", nil) |
| 58 | req.Header.Set("Accept-Encoding", "gzip") |
| 59 | req.Header.Set("If-None-Match", `"abc"`) |
| 60 | resp304, err := http.DefaultClient.Do(req) |
| 61 | if err != nil { |
| 62 | t.Fatal(err) |
| 63 | } |
| 64 | defer resp304.Body.Close() |
| 65 | if resp304.StatusCode != http.StatusNotModified || resp304.Header.Get("Content-Encoding") != "" { |
| 66 | t.Fatalf("304 response = status %d encoding %q", resp304.StatusCode, resp304.Header.Get("Content-Encoding")) |
| 67 | } |
| 68 | |
| 69 | req, _ = http.NewRequest(http.MethodGet, srv.URL+"/events", nil) |
| 70 | req.Header.Set("Accept-Encoding", "gzip") |
| 71 | events, err := http.DefaultClient.Do(req) |
| 72 | if err != nil { |
| 73 | t.Fatal(err) |
| 74 | } |
| 75 | defer events.Body.Close() |
| 76 | if events.Header.Get("Content-Encoding") != "" { |
| 77 | t.Fatal("SSE must bypass gzip") |
| 78 | } |
| 79 | } |
| 80 | |
| 81 | func TestGzipMiddlewareKeepsSmallResponsesPlain(t *testing.T) { |
| 82 | h := gzipMiddleware(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { |
| 83 | _, _ = w.Write([]byte("small")) |
| 84 | })) |
| 85 | req := httptest.NewRequest(http.MethodGet, "/status", nil) |
| 86 | req.Header.Set("Accept-Encoding", "gzip") |
| 87 | rr := httptest.NewRecorder() |
| 88 | h.ServeHTTP(rr, req) |
| 89 | if got := rr.Header().Get("Content-Encoding"); got != "" { |
| 90 | t.Fatalf("small response encoding = %q", got) |
| 91 | } |
| 92 | if rr.Body.String() != "small" { |
| 93 | t.Fatalf("small body = %q", rr.Body.String()) |
| 94 | } |
| 95 | if got := rr.Header().Get("Content-Type"); got != "text/plain; charset=utf-8" { |
| 96 | t.Fatalf("small response Content-Type = %q", got) |
| 97 | } |
| 98 | if got := rr.Header().Get("X-Content-Type-Options"); got != "nosniff" { |
| 99 | t.Fatalf("small response X-Content-Type-Options = %q", got) |
| 100 | } |
| 101 | } |
| 102 | |
| 103 | func TestGzipMiddlewareDoesNotRenderUntypedReflectedHTML(t *testing.T) { |
| 104 | h := gzipMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 105 | _, _ = w.Write([]byte(r.URL.Query().Get("value"))) |
| 106 | })) |
| 107 | req := httptest.NewRequest(http.MethodGet, "/status?value=%3Cscript%3Ealert(1)%3C/script%3E", nil) |
| 108 | req.Header.Set("Accept-Encoding", "gzip") |
| 109 | rr := httptest.NewRecorder() |
| 110 | h.ServeHTTP(rr, req) |
| 111 | if got := rr.Header().Get("Content-Type"); got != "text/plain; charset=utf-8" { |
| 112 | t.Fatalf("reflected response Content-Type = %q", got) |
| 113 | } |
| 114 | if got := rr.Header().Get("X-Content-Type-Options"); got != "nosniff" { |
| 115 | t.Fatalf("reflected response X-Content-Type-Options = %q", got) |
| 116 | } |
| 117 | } |
| 118 | |
| 119 | func TestGzipMiddlewareHonorsDisabledEncoding(t *testing.T) { |
| 120 | h := gzipMiddleware(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { |
| 121 | _, _ = w.Write([]byte(strings.Repeat("x", gzipThreshold+1))) |
| 122 | })) |
| 123 | for _, encoding := range []string{"gzip;q=0", "br, *;q=1, gzip;q=0"} { |
| 124 | req := httptest.NewRequest(http.MethodGet, "/history", nil) |
| 125 | req.Header.Set("Accept-Encoding", encoding) |
| 126 | rr := httptest.NewRecorder() |
| 127 | h.ServeHTTP(rr, req) |
| 128 | if got := rr.Header().Get("Content-Encoding"); got != "" { |
| 129 | t.Fatalf("Accept-Encoding %q produced %q", encoding, got) |
| 130 | } |
| 131 | } |
| 132 | } |
| 133 | |
| 134 | func TestGzipPlainWritesPreserveDeclaredMediaType(t *testing.T) { |
| 135 | for _, mediaType := range []string{"", "application/json", "text/html; charset=utf-8", "application/octet-stream"} { |
| 136 | t.Run(mediaType, func(t *testing.T) { |
| 137 | rr := httptest.NewRecorder() |
| 138 | if mediaType != "" { |
| 139 | rr.Header().Set("Content-Type", mediaType) |
| 140 | } |
| 141 | writer := &gzipBufferedWriter{ResponseWriter: rr} |
| 142 | writer.startPlain() |
| 143 | payload := []byte("<script>alert(1)</script>") |
| 144 | if _, err := writer.Write(payload); err != nil { |
| 145 | t.Fatal(err) |
| 146 | } |
| 147 | response := rr.Result() |
| 148 | defer response.Body.Close() |
| 149 | wantType := mediaType |
| 150 | if wantType == "" { |
| 151 | wantType = "text/plain; charset=utf-8" |
| 152 | } |
| 153 | if got := response.Header.Get("Content-Type"); got != wantType { |
| 154 | t.Fatalf("committed content type = %q, want %q", got, wantType) |
| 155 | } |
| 156 | if response.Header.Get("X-Content-Type-Options") != "nosniff" || !bytes.Equal(rr.Body.Bytes(), payload) { |
| 157 | t.Fatal("plain write changed payload or omitted nosniff") |
| 158 | } |
| 159 | }) |
| 160 | } |
| 161 | } |
| 162 |