Compare commits

...

6 Commits

6 changed files with 117 additions and 29 deletions
+35 -17
View File
@@ -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)
})
}
}
+35
View File
@@ -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) {
+6 -1
View File
@@ -71,7 +71,7 @@ func CSRFMiddleware() func(http.Handler) http.Handler {
return return
} }
if strings.HasPrefix(r.URL.Path, "/api/") { if strings.HasPrefix(r.URL.Path, "/api/") && hasBearerToken(r) {
next.ServeHTTP(w, r) next.ServeHTTP(w, r)
return return
} }
@@ -86,6 +86,11 @@ func CSRFMiddleware() func(http.Handler) http.Handler {
} }
} }
func hasBearerToken(r *http.Request) bool {
auth := strings.TrimSpace(r.Header.Get("Authorization"))
return strings.HasPrefix(auth, "Bearer ")
}
func IsHTTPS(r *http.Request) bool { func IsHTTPS(r *http.Request) bool {
if r.TLS != nil { if r.TLS != nil {
return true return true
+18 -2
View File
@@ -186,8 +186,9 @@ func TestCSRFMiddlewareAllowsValidToken(t *testing.T) {
} }
} }
func TestCSRFMiddlewareSkipsAPI(t *testing.T) { func TestCSRFMiddlewareSkipsAPIWithBearerToken(t *testing.T) {
request := httptest.NewRequest("POST", "/api/test", nil) request := httptest.NewRequest("POST", "/api/test", nil)
request.Header.Set("Authorization", "Bearer valid-token")
recorder := httptest.NewRecorder() recorder := httptest.NewRecorder()
handler := CSRFMiddleware()(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { handler := CSRFMiddleware()(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -197,7 +198,22 @@ func TestCSRFMiddlewareSkipsAPI(t *testing.T) {
handler.ServeHTTP(recorder, request) handler.ServeHTTP(recorder, request)
if recorder.Code != http.StatusOK { if recorder.Code != http.StatusOK {
t.Errorf("API requests should skip CSRF validation, got status %d", recorder.Code) t.Errorf("API requests with Bearer token should skip CSRF validation, got status %d", recorder.Code)
}
}
func TestCSRFMiddlewareBlocksAPIWithoutBearerToken(t *testing.T) {
request := httptest.NewRequest("POST", "/api/test", nil)
recorder := httptest.NewRecorder()
handler := CSRFMiddleware()(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusOK)
}))
handler.ServeHTTP(recorder, request)
if recorder.Code != http.StatusForbidden {
t.Errorf("API requests without Bearer token should require CSRF validation, got status %d", recorder.Code)
} }
} }
+19 -7
View File
@@ -5,6 +5,7 @@ import (
"net/http" "net/http"
"os" "os"
"strings" "strings"
"sync"
"time" "time"
) )
@@ -220,18 +221,29 @@ func isSuspiciousUserAgent(userAgent string) bool {
return false return false
} }
var requestCounts = make(map[string]int) type rapidRequestTracker struct {
var lastReset = time.Now() mu sync.Mutex
counts map[string]int
lastReset time.Time
}
var rapidRequests = rapidRequestTracker{
counts: make(map[string]int),
lastReset: time.Now(),
}
func isRapidRequest(ip string) bool { func isRapidRequest(ip string) bool {
rapidRequests.mu.Lock()
defer rapidRequests.mu.Unlock()
now := time.Now() now := time.Now()
if now.Sub(lastReset) > time.Minute { if now.Sub(rapidRequests.lastReset) > time.Minute {
requestCounts = make(map[string]int) rapidRequests.counts = make(map[string]int)
lastReset = now rapidRequests.lastReset = now
} }
requestCounts[ip]++ rapidRequests.counts[ip]++
return requestCounts[ip] > 100 return rapidRequests.counts[ip] > 100
} }
+4 -2
View File
@@ -453,8 +453,10 @@ func TestIsSuspiciousUserAgent(t *testing.T) {
func TestIsRapidRequest(t *testing.T) { func TestIsRapidRequest(t *testing.T) {
requestCounts = make(map[string]int) rapidRequests.mu.Lock()
lastReset = time.Now() rapidRequests.counts = make(map[string]int)
rapidRequests.lastReset = time.Now()
rapidRequests.mu.Unlock()
ip := "192.168.1.1" ip := "192.168.1.1"