lijiaoqiao/gateway/internal/handler/handler_security_test.go

package handler

import (
	"bytes"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"lijiaoqiao/gateway/internal/router"
)

func TestMED05_RequestBodySizeLimit(t *testing.T) {
	// MED-05: Request body size should be limited to prevent DoS attacks
	// json.Decoder should use MaxBytes to limit request body size

	r := router.NewRouter(router.StrategyLatency)
	h := NewHandler(r)

	// Create a very large request body (exceeds 1MB limit)
	largeContent := strings.Repeat("a", 2*1024*1024) // 2MB
	largeBody := `{"model": "gpt-4", "messages": [{"role": "user", "content": "` + largeContent + `"}]}`

	req := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewBufferString(largeBody))
	req.Header.Set("Content-Type", "application/json")
	rr := httptest.NewRecorder()

	h.ChatCompletionsHandle(rr, req)

	// After fix: should return 413 Request Entity Too Large
	if rr.Code != http.StatusRequestEntityTooLarge {
		t.Errorf("expected status 413 for large request body, got %d", rr.Code)
	}
}

func TestMED05_NormalRequestShouldPass(t *testing.T) {
	// Normal requests should still work
	r := router.NewRouter(router.StrategyLatency)
	prov := &mockProvider{name: "test", models: []string{"gpt-4"}, healthy: true}
	r.RegisterProvider("test", prov)

	h := NewHandler(r)

	body := `{"model": "gpt-4", "messages": [{"role": "user", "content": "hello"}]}`
	req := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewBufferString(body))
	req.Header.Set("Content-Type", "application/json")
	rr := httptest.NewRecorder()

	h.ChatCompletionsHandle(rr, req)

	// Should succeed (status 200)
	if rr.Code != http.StatusOK {
		t.Errorf("expected status 200 for normal request, got %d", rr.Code)
	}
}

func TestMED05_EmptyBodyShouldFail(t *testing.T) {
	// Empty request body should fail
	r := router.NewRouter(router.StrategyLatency)
	h := NewHandler(r)

	req := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewBufferString(""))
	req.Header.Set("Content-Type", "application/json")
	rr := httptest.NewRecorder()

	h.ChatCompletionsHandle(rr, req)

	// Should fail with 400 Bad Request
	if rr.Code != http.StatusBadRequest {
		t.Errorf("expected status 400 for empty body, got %d", rr.Code)
	}
}

func TestMED05_InvalidJSONShouldFail(t *testing.T) {
	// Invalid JSON should fail
	r := router.NewRouter(router.StrategyLatency)
	h := NewHandler(r)

	body := `{invalid json}`
	req := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewBufferString(body))
	req.Header.Set("Content-Type", "application/json")
	rr := httptest.NewRecorder()

	h.ChatCompletionsHandle(rr, req)

	// Should fail with 400 Bad Request
	if rr.Code != http.StatusBadRequest {
		t.Errorf("expected status 400 for invalid JSON, got %d", rr.Code)
	}
}

// TestMaxBytesReaderWrapper tests the MaxBytes reader wrapper behavior
func TestMaxBytesReaderWrapper(t *testing.T) {
	// Test that limiting reader works correctly
	content := "hello world"
	limitedReader := io.LimitReader(bytes.NewReader([]byte(content)), 5)

	buf := make([]byte, 20)
	n, err := limitedReader.Read(buf)

	// Should only read 5 bytes
	if n != 5 {
		t.Errorf("expected to read 5 bytes, got %d", n)
	}
	if err != nil && err != io.EOF {
		t.Errorf("expected no error or EOF, got %v", err)
	}

	// Reading again should return 0 with EOF
	n2, err2 := limitedReader.Read(buf)
	if n2 != 0 {
		t.Errorf("expected 0 bytes on second read, got %d", n2)
	}
	if err2 != io.EOF {
		t.Errorf("expected EOF on second read, got %v", err2)
	}
}