Compare commits
2 Commits
8f255a4fe6
...
ab17ff8b79
| Author | SHA1 | Date | |
|---|---|---|---|
| ab17ff8b79 | |||
| 8990f5afb7 |
@@ -150,23 +150,7 @@ func shouldCompressResponse(contentType string, config *CompressionConfig) bool
|
|||||||
}
|
}
|
||||||
|
|
||||||
func DecompressionMiddleware() func(http.Handler) http.Handler {
|
func DecompressionMiddleware() func(http.Handler) http.Handler {
|
||||||
return func(next http.Handler) http.Handler {
|
return DecompressionMiddlewareWithConfig(nil)
|
||||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
||||||
if r.Header.Get("Content-Encoding") == "gzip" {
|
|
||||||
gz, err := gzip.NewReader(r.Body)
|
|
||||||
if err != nil {
|
|
||||||
http.Error(w, "Invalid gzip encoding", http.StatusBadRequest)
|
|
||||||
return
|
|
||||||
}
|
|
||||||
defer gz.Close()
|
|
||||||
|
|
||||||
r.Body = io.NopCloser(gz)
|
|
||||||
r.Header.Del("Content-Encoding")
|
|
||||||
}
|
|
||||||
|
|
||||||
next.ServeHTTP(w, r)
|
|
||||||
})
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
type CompressionConfig struct {
|
type CompressionConfig struct {
|
||||||
@@ -189,3 +173,37 @@ func DefaultCompressionConfig() *CompressionConfig {
|
|||||||
},
|
},
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
type DecompressionConfig struct {
|
||||||
|
MaxDecompressedSize int64
|
||||||
|
}
|
||||||
|
|
||||||
|
func DefaultDecompressionConfig() *DecompressionConfig {
|
||||||
|
return &DecompressionConfig{
|
||||||
|
MaxDecompressedSize: 1024 * 1024, // 1MB
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func DecompressionMiddlewareWithConfig(config *DecompressionConfig) func(http.Handler) http.Handler {
|
||||||
|
if config == nil {
|
||||||
|
config = DefaultDecompressionConfig()
|
||||||
|
}
|
||||||
|
|
||||||
|
return func(next http.Handler) http.Handler {
|
||||||
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
if r.Header.Get("Content-Encoding") == "gzip" {
|
||||||
|
gz, err := gzip.NewReader(r.Body)
|
||||||
|
if err != nil {
|
||||||
|
http.Error(w, "Invalid gzip encoding", http.StatusBadRequest)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
defer gz.Close()
|
||||||
|
|
||||||
|
r.Body = io.NopCloser(io.LimitReader(gz, config.MaxDecompressedSize))
|
||||||
|
r.Header.Del("Content-Encoding")
|
||||||
|
}
|
||||||
|
|
||||||
|
next.ServeHTTP(w, r)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -562,6 +562,41 @@ func TestDecompressionMiddleware(t *testing.T) {
|
|||||||
t.Errorf("Expected empty body, got '%s'", recorder.Body.String())
|
t.Errorf("Expected empty body, got '%s'", recorder.Body.String())
|
||||||
}
|
}
|
||||||
})
|
})
|
||||||
|
|
||||||
|
t.Run("Limits decompressed size", func(t *testing.T) {
|
||||||
|
config := &DecompressionConfig{
|
||||||
|
MaxDecompressedSize: 10,
|
||||||
|
}
|
||||||
|
middleware := DecompressionMiddlewareWithConfig(config)
|
||||||
|
|
||||||
|
var buf bytes.Buffer
|
||||||
|
gz := gzip.NewWriter(&buf)
|
||||||
|
gz.Write([]byte("this is more than ten bytes of data"))
|
||||||
|
gz.Close()
|
||||||
|
|
||||||
|
request := httptest.NewRequest("POST", "/test", &buf)
|
||||||
|
request.Header.Set("Content-Encoding", "gzip")
|
||||||
|
recorder := httptest.NewRecorder()
|
||||||
|
|
||||||
|
handler := middleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
body, err := io.ReadAll(r.Body)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Failed to read request body: %v", err)
|
||||||
|
}
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
w.Write(body)
|
||||||
|
}))
|
||||||
|
|
||||||
|
handler.ServeHTTP(recorder, request)
|
||||||
|
|
||||||
|
if recorder.Code != http.StatusOK {
|
||||||
|
t.Errorf("Expected status 200, got %d", recorder.Code)
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(recorder.Body.String()) > 10 {
|
||||||
|
t.Errorf("Expected body to be truncated to <= 10 bytes, got %d bytes", len(recorder.Body.String()))
|
||||||
|
}
|
||||||
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestShouldCompressWithConfig(t *testing.T) {
|
func TestShouldCompressWithConfig(t *testing.T) {
|
||||||
|
|||||||
Reference in New Issue
Block a user