diff options
| author | google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com> | 2026-02-18 03:20:37 +0000 |
|---|---|---|
| committer | google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com> | 2026-02-18 03:20:37 +0000 |
| commit | a6974c3af00a66870424aa971ae4dd38c49847bf (patch) | |
| tree | 120e4dbdef90081e0208f9154197faeb057d74e2 | |
| parent | 9db36ae402dbb74f7223a4efc8b2483086684e38 (diff) | |
| download | neko-a6974c3af00a66870424aa971ae4dd38c49847bf.tar.gz neko-a6974c3af00a66870424aa971ae4dd38c49847bf.tar.bz2 neko-a6974c3af00a66870424aa971ae4dd38c49847bf.zip | |
Fix Unbounded Memory Usage in Image Proxy
- Added `maxImageProxySize` constant (10MB) to limit memory usage.
- Used `io.LimitReader` in `imageProxyHandler` to enforce the limit.
- Added regression test `web/proxy_limit_test.go` to verify the fix.
Co-authored-by: adammathes <868470+adammathes@users.noreply.github.com>
| -rw-r--r-- | web/proxy_limit_test.go | 44 | ||||
| -rw-r--r-- | web/web.go | 4 |
2 files changed, 47 insertions, 1 deletions
diff --git a/web/proxy_limit_test.go b/web/proxy_limit_test.go new file mode 100644 index 0000000..2bd0637 --- /dev/null +++ b/web/proxy_limit_test.go @@ -0,0 +1,44 @@ +package web + +import ( + "encoding/base64" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +// TestImageProxyHandlerUnboundedRead attempts to reproduce the issue by +// serving a large response and checking if the handler reads it all. +// Since we can't easily check memory usage in a test without potentially crashing, +// this test will serve as a baseline. +// +// After the fix, we will modify this test or add a new one to verify that +// the handler stops reading after a certain limit. +func TestImageProxyHandlerLargeResponse(t *testing.T) { + // Create a mock server that returns a large response + largeContent := strings.Repeat("A", 1024*1024*15) // 15MB + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "image/jpeg") + w.Write([]byte(largeContent)) + })) + defer ts.Close() + + // Encode the URL + encodedURL := base64.URLEncoding.EncodeToString([]byte(ts.URL)) + // The handler expects the path to be the encoded URL (after StripPrefix) + req := httptest.NewRequest("GET", "/"+encodedURL, nil) + rr := httptest.NewRecorder() + + // Call the handler directly (skipping auth/routing middleware for simplicity) + http.HandlerFunc(imageProxyHandler).ServeHTTP(rr, req) + + if rr.Code != http.StatusOK { + t.Errorf("Expected status 200, got %d", rr.Code) + } + + // With the fix, the handler should read only up to the limit. + if len(rr.Body.Bytes()) != maxImageProxySize { + t.Errorf("Expected response body length %d, got %d", maxImageProxySize, len(rr.Body.Bytes())) + } +} @@ -31,6 +31,8 @@ var gzPool = sync.Pool{ }, } +const maxImageProxySize = 10 * 1024 * 1024 // 10 MB + var ( //go:embed static/* staticFiles embed.FS @@ -89,7 +91,7 @@ func imageProxyHandler(w http.ResponseWriter, r *http.Request) { return } - bts, err := io.ReadAll(resp.Body) + bts, err := io.ReadAll(io.LimitReader(resp.Body, maxImageProxySize)) if err != nil { http.Error(w, "failed to read proxy image", http.StatusNotFound) return |
