Created
October 8, 2025 02:56
-
-
Save hnakamur/5c73c0726e73e0ff1935e62ea3ae63ce to your computer and use it in GitHub Desktop.
Buffer the first part of the response body and pass through the rest of it
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| package main | |
| import ( | |
| "fmt" | |
| "io" | |
| "net/http" | |
| "net/http/httptest" | |
| "testing" | |
| ) | |
| func TestHTTPMiddleware(t *testing.T) { | |
| content := "<html><body>Hello World!</body></html>" | |
| testCases := []struct { | |
| testName string | |
| originHandler func(http.ResponseWriter, *http.Request) | |
| }{ | |
| { | |
| testName: "WriteAtOnce", | |
| originHandler: func(w http.ResponseWriter, _ *http.Request) { | |
| io.WriteString(w, content) | |
| }, | |
| }, | |
| { | |
| testName: "WriteInChunks", | |
| originHandler: func(w http.ResponseWriter, _ *http.Request) { | |
| contentLenThird := len(content) / 3 | |
| io.WriteString(w, content[:contentLenThird]) | |
| if f, ok := w.(http.Flusher); ok { | |
| f.Flush() | |
| } | |
| io.WriteString(w, content[contentLenThird:2*contentLenThird]) | |
| if f, ok := w.(http.Flusher); ok { | |
| f.Flush() | |
| } | |
| io.WriteString(w, content[2*contentLenThird:]) | |
| if f, ok := w.(http.Flusher); ok { | |
| f.Flush() | |
| } | |
| }, | |
| }, | |
| } | |
| for _, testCase := range testCases { | |
| t.Run(testCase.testName, func(t *testing.T) { | |
| // NOTE: capacity must be >= 6 for the above content | |
| // in order to set content-type correctly. | |
| capacities := []int{4, 5, 6, 7, 8, len(content)/3 - 1, len(content) / 3, len(content)/3 + 1, len(content) / 2, len(content), len(content) + 1} | |
| for _, capacity := range capacities { | |
| t.Run(fmt.Sprintf("capacity%d", capacity), func(t *testing.T) { | |
| handler := func(w http.ResponseWriter, r *http.Request) { | |
| i := NewRwInterceptorWithBufCap(w, capacity) | |
| testCase.originHandler(i, r) | |
| _ = i.processAndFlushBuffer() | |
| } | |
| req := httptest.NewRequest("GET", "http://example.com/foo", nil) | |
| w := httptest.NewRecorder() | |
| handler(w, req) | |
| resp := w.Result() | |
| body, _ := io.ReadAll(resp.Body) | |
| if got, want := resp.StatusCode, http.StatusOK; got != want { | |
| t.Errorf("statusCode mismatch, got=%v, want=%v", got, want) | |
| } | |
| if got, want := resp.Header.Get("Content-Type"), "text/html; charset=utf-8"; got != want { | |
| t.Errorf("content-type mismatch, got=%v, want=%v", got, want) | |
| } | |
| if got, want := string(body), content; got != want { | |
| t.Errorf("content mismatch, got=%v, want=%v", got, want) | |
| } | |
| }) | |
| } | |
| }) | |
| } | |
| } | |
| type rwInterceptor struct { | |
| w http.ResponseWriter | |
| buf []byte | |
| } | |
| var _ http.ResponseWriter = (*rwInterceptor)(nil) | |
| func NewRwInterceptorWithBufCap(w http.ResponseWriter, capacity int) *rwInterceptor { | |
| i := &rwInterceptor{ | |
| w: w, | |
| buf: make([]byte, 0, capacity), | |
| } | |
| return i | |
| } | |
| // Header implements http.ResponseWriter. | |
| func (i *rwInterceptor) Header() http.Header { | |
| return i.w.Header() | |
| } | |
| // Write implements http.ResponseWriter. | |
| func (i *rwInterceptor) Write(p []byte) (int, error) { | |
| // log.Printf("rwInterceptor Write start, len(p)=%d, cap(buf)=%d, len(buf)=%d", len(p), cap(i.buf), len(i.buf)) | |
| if i.buf == nil { | |
| return i.w.Write(p) | |
| } | |
| toWrite := min(len(p), cap(i.buf)-len(i.buf)) | |
| i.buf = append(i.buf, p[:toWrite]...) | |
| if cap(i.buf)-len(i.buf) > 0 { | |
| return toWrite, nil | |
| } | |
| if err := i.processAndFlushBuffer(); err != nil || toWrite == len(p) { | |
| return toWrite, err | |
| } | |
| n2, err := i.w.Write(p[toWrite:]) | |
| return toWrite + n2, err | |
| } | |
| // WriteHeader implements http.ResponseWriter. | |
| func (i *rwInterceptor) WriteHeader(statusCode int) { | |
| i.w.WriteHeader(statusCode) | |
| } | |
| func (i *rwInterceptor) processAndFlushBuffer() error { | |
| if i.buf == nil { | |
| return nil | |
| } | |
| // process buffer here. | |
| if _, err := i.w.Write(i.buf); err != nil { | |
| return err | |
| } | |
| i.buf = nil | |
| return nil | |
| } |
Author
hnakamur
commented
Oct 8, 2025
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment