From 230432db448df88e8113b116af150e7504544351 Mon Sep 17 00:00:00 2001 From: Chen Su Date: Fri, 24 Jul 2026 09:52:36 +0800 Subject: [PATCH 1/2] feat: add String, JSON, and Redirect Signed-off-by: Chen Su --- context.go | 56 +++++++++ context_test.go | 326 ++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 382 insertions(+) diff --git a/context.go b/context.go index e85f6f1..acc6137 100644 --- a/context.go +++ b/context.go @@ -2,7 +2,10 @@ package engine import ( "context" + "encoding/json" "errors" + "fmt" + "html" "io" "mime" "net/http" @@ -306,6 +309,59 @@ func (ctx *Context) Write(data []byte) (int, error) { return ctx.res.Write(data) } +// String writes a string to the response body. +func (ctx *Context) String(s string) (int, error) { + if ctx.GetHeader("Content-Type") == "" { + ctx.SetHeader("Content-Type", "text/plain; charset=utf-8") + } + + return ctx.res.Write([]byte(s)) +} + +// JSON serializes the given value to JSON and writes it to the response body. +func (ctx *Context) JSON(v any) (int, error) { + jsonData, err := json.Marshal(v) + if err != nil { + return 0, err + } + + if ctx.GetHeader("Content-Type") == "" { + ctx.SetHeader("Content-Type", "application/json") + } + + return ctx.res.Write(jsonData) +} + +// Redirect sets the Location header and writes a redirect response with the specified status code, +// defaulting to 302 Found if no code is provided. +func (ctx *Context) Redirect(link string, code ...int) error { + statusCode := http.StatusFound + if len(code) > 0 { + statusCode = code[0] + } + switch statusCode { + case http.StatusMultipleChoices, http.StatusMovedPermanently, http.StatusFound, http.StatusSeeOther, + http.StatusTemporaryRedirect, http.StatusPermanentRedirect: + ctx.Status(statusCode) + default: + return fmt.Errorf("invalid redirect status code: %d", statusCode) + } + + ctx.SetHeader("Location", link) + if ctx.GetHeader("Content-Type") == "" { + ctx.SetHeader("Content-Type", "text/html; charset=utf-8") + } + + str := fmt.Sprintf("%s", html.EscapeString(link), http.StatusText(statusCode)) + n, err := ctx.res.Write([]byte(str)) + if err != nil { + return err + } else if n != len(str) { + return errors.New("failed to write complete redirect response") + } + return nil +} + // Request returns the wrapped Request object associated with the context. func (ctx *Context) Request() core.Request { return ctx.req diff --git a/context_test.go b/context_test.go index 4fa46b9..5684d31 100644 --- a/context_test.go +++ b/context_test.go @@ -3,10 +3,12 @@ package engine_test import ( "bytes" "errors" + "fmt" "io" "net/http" "net/http/httptest" "reflect" + "strings" "testing" "github.com/go-amwk/core" @@ -622,6 +624,330 @@ func TestContext_Write(t *testing.T) { } } +func TestContext_String(t *testing.T) { + rr := httptest.NewRecorder() + resp := NewResponse(rr) + ctx := getDefaultContext(nil, resp) + n, err := ctx.String("hello") + if err != nil { + t.Fatalf("String returned error: %v", err) + } + if n != 5 { + t.Errorf("Expected String to write 5 bytes, got %d", n) + } + + resp.send() + + if rr.Body.String() != "hello" { + t.Errorf("Expected sent body 'hello', got %v", rr.Body.String()) + } + if rr.Header().Get("Content-Type") != "text/plain; charset=utf-8" { + t.Errorf("Expected Content-Type 'text/plain; charset=utf-8', got %v", rr.Header().Get("Content-Type")) + } +} + +func TestContext_String_PreserveContentType(t *testing.T) { + rr := httptest.NewRecorder() + resp := NewResponse(rr) + ctx := getDefaultContext(nil, resp) + ctx.SetHeader("Content-Type", "text/html; charset=utf-8") + n, err := ctx.String("

hello

") + if err != nil { + t.Fatalf("String returned error: %v", err) + } + if n != 12 { + t.Errorf("Expected String to write 12 bytes, got %d", n) + } + + resp.send() + + if rr.Header().Get("Content-Type") != "text/html; charset=utf-8" { + t.Errorf("Expected Content-Type to be preserved as 'text/html; charset=utf-8', got %v", rr.Header().Get("Content-Type")) + } +} + +func TestContext_JSON(t *testing.T) { + rr := httptest.NewRecorder() + resp := NewResponse(rr) + ctx := getDefaultContext(nil, resp) + data := map[string]string{"key": "value"} + n, err := ctx.JSON(data) + if err != nil { + t.Fatalf("JSON returned error: %v", err) + } + expectedJSON := `{"key":"value"}` + if n != len(expectedJSON) { + t.Errorf("Expected JSON to write %d bytes, got %d", len(expectedJSON), n) + } + + resp.send() + + if rr.Body.String() != expectedJSON { + t.Errorf("Expected sent body '%s', got '%s'", expectedJSON, rr.Body.String()) + } + if rr.Header().Get("Content-Type") != "application/json" { + t.Errorf("Expected Content-Type 'application/json', got %v", rr.Header().Get("Content-Type")) + } +} + +func TestContext_JSON_PreserveContentType(t *testing.T) { + rr := httptest.NewRecorder() + resp := NewResponse(rr) + ctx := getDefaultContext(nil, resp) + ctx.SetHeader("Content-Type", "application/json; charset=utf-8") + data := map[string]string{"key": "value"} + _, err := ctx.JSON(data) + if err != nil { + t.Fatalf("JSON returned error: %v", err) + } + + resp.send() + + if rr.Header().Get("Content-Type") != "application/json; charset=utf-8" { + t.Errorf("Expected Content-Type to be preserved as 'application/json; charset=utf-8', got %v", rr.Header().Get("Content-Type")) + } +} + +func TestContext_JSON_MarshalError(t *testing.T) { + resp := NewResponse(httptest.NewRecorder()) + ctx := getDefaultContext(nil, resp) + // channel and func types cannot be marshaled to JSON + _, err := ctx.JSON(make(chan int)) + if err == nil { + t.Fatalf("Expected JSON to return error for unmarshalable type") + } +} + +func TestContext_String_Empty(t *testing.T) { + rr := httptest.NewRecorder() + resp := NewResponse(rr) + ctx := getDefaultContext(nil, resp) + n, err := ctx.String("") + if err != nil { + t.Fatalf("String returned error: %v", err) + } + if n != 0 { + t.Errorf("Expected String to write 0 bytes, got %d", n) + } + + resp.send() + + if rr.Body.String() != "" { + t.Errorf("Expected empty body, got '%s'", rr.Body.String()) + } + if rr.Header().Get("Content-Type") != "text/plain; charset=utf-8" { + t.Errorf("Expected Content-Type 'text/plain; charset=utf-8', got %v", rr.Header().Get("Content-Type")) + } +} + +func TestContext_JSON_Nil(t *testing.T) { + rr := httptest.NewRecorder() + resp := NewResponse(rr) + ctx := getDefaultContext(nil, resp) + n, err := ctx.JSON(nil) + if err != nil { + t.Fatalf("JSON returned error: %v", err) + } + expectedJSON := "null" + if n != len(expectedJSON) { + t.Errorf("Expected JSON to write %d bytes, got %d", len(expectedJSON), n) + } + + resp.send() + + if rr.Body.String() != expectedJSON { + t.Errorf("Expected sent body 'null', got '%s'", rr.Body.String()) + } + if rr.Header().Get("Content-Type") != "application/json" { + t.Errorf("Expected Content-Type 'application/json', got %v", rr.Header().Get("Content-Type")) + } +} + +func TestContext_JSON_MarshalError_NoContentTypeSideEffect(t *testing.T) { + resp := NewResponse(httptest.NewRecorder()) + ctx := getDefaultContext(nil, resp) + // channel types cannot be marshaled — Content-Type must NOT be set on failure + _, err := ctx.JSON(make(chan int)) + if err == nil { + t.Fatalf("Expected JSON to return error for unmarshalable type") + } + if ctx.GetHeader("Content-Type") != "" { + t.Errorf("Expected no Content-Type header to be set on marshal error, got '%s'", ctx.GetHeader("Content-Type")) + } +} + +func TestContext_Redirect(t *testing.T) { + rr := httptest.NewRecorder() + resp := NewResponse(rr) + ctx := getDefaultContext(nil, resp) + + err := ctx.Redirect("https://example.com", http.StatusMovedPermanently) + if err != nil { + t.Fatalf("Redirect returned error: %v", err) + } + + resp.send() + + if rr.Header().Get("Location") != "https://example.com" { + t.Errorf("Expected Location header to be set to 'https://example.com', got %v", rr.Header().Get("Location")) + } + if rr.Header().Get("Content-Type") != "text/html; charset=utf-8" { + t.Errorf("Expected Content-Type 'text/html; charset=utf-8', got %v", rr.Header().Get("Content-Type")) + } + + if rr.Code != http.StatusMovedPermanently { + t.Errorf("Expected status code %d, got %d", http.StatusMovedPermanently, rr.Code) + } + expectedBody := `Moved Permanently` + if rr.Body.String() != expectedBody { + t.Errorf("Expected body '%s', got '%s'", expectedBody, rr.Body.String()) + } +} + +func TestContext_Redirect_DefaultStatus(t *testing.T) { + rr := httptest.NewRecorder() + resp := NewResponse(rr) + ctx := getDefaultContext(nil, resp) + + err := ctx.Redirect("https://example.com") + if err != nil { + t.Fatalf("Redirect returned error: %v", err) + } + + resp.send() + + if rr.Header().Get("Location") != "https://example.com" { + t.Errorf("Expected Location header to be set to 'https://example.com', got %v", rr.Header().Get("Location")) + } + if rr.Header().Get("Content-Type") != "text/html; charset=utf-8" { + t.Errorf("Expected Content-Type 'text/html; charset=utf-8', got %v", rr.Header().Get("Content-Type")) + } + if rr.Code != http.StatusFound { + t.Errorf("Expected status code %d, got %d", http.StatusFound, rr.Code) + } + expectedBody := `Found` + if rr.Body.String() != expectedBody { + t.Errorf("Expected body '%s', got '%s'", expectedBody, rr.Body.String()) + } +} + +func TestContext_Redirect_InvalidCode(t *testing.T) { + rr := httptest.NewRecorder() + resp := NewResponse(rr) + ctx := getDefaultContext(nil, resp) + + err := ctx.Redirect("https://example.com", http.StatusBadRequest) + if err == nil { + t.Fatalf("Expected Redirect to return error for non-redirect status code 400") + } + + err = ctx.Redirect("https://example.com", 200) + if err == nil { + t.Fatalf("Expected Redirect to return error for non-redirect status code 200") + } + + err = ctx.Redirect("https://example.com", 500) + if err == nil { + t.Fatalf("Expected Redirect to return error for non-redirect status code 500") + } + + // Verify no headers were set on error + if rr.Header().Get("Location") != "" { + t.Errorf("Expected no Location header to be set on error") + } + if rr.Header().Get("Content-Type") != "" { + t.Errorf("Expected no Content-Type header to be set on error") + } +} + +func TestContext_Redirect_XSSEscaping(t *testing.T) { + rr := httptest.NewRecorder() + resp := NewResponse(rr) + ctx := getDefaultContext(nil, resp) + + dangerousLink := `https://example.com/` + err := ctx.Redirect(dangerousLink) + if err != nil { + t.Fatalf("Redirect returned error: %v", err) + } + + resp.send() + + // Verify dangerous characters are escaped in body + body := rr.Body.String() + if strings.Contains(body, "` - err := ctx.Redirect(dangerousLink) - if err != nil { - t.Fatalf("Redirect returned error: %v", err) - } - - resp.send() - - // Verify dangerous characters are escaped in body - body := rr.Body.String() - if strings.Contains(body, "