Enhance Response with body handlers

response.go

@@ -77,24 +77,10 @@ func (resp *Response) Headers() http.Header {
 	return resp.headers
 }
 
-// Write writes the given data to the response body. It appends the data to any existing content in
-// the body. It returns the number of bytes written and any error encountered during the write
-// operation.
-func (resp *Response) Write(data []byte) (int, error) {
-	if len(data) == 0 {
-		return 0, nil
-	}
-
-	maxBytes := resp.maxBodyBytes.Load()
-
-	if maxBytes == MaxResponseBodyBytesUnlimited {
-		return resp.body.Write(data)
-	}
-
-	if int64(resp.body.Len())+int64(len(data)) > maxBytes {
-		return 0, ErrResponseTooLarge
-	}
-	return resp.body.Write(data)
+// Size returns the current size of the response body in bytes. This allows you to check how much data
+// has been written to the response body before sending it back to the client.
+func (resp *Response) Size() int {
+	return resp.body.Len()
 }
 
 // Status sets the HTTP status code for the response. This method allows you to specify the status
@@ -111,6 +97,43 @@ func (resp *Response) StatusCode() int {
 	return resp.statusCode
 }
 
+// Write writes the given data to the response body. It appends the data to any existing content in
+// the body. It returns the number of bytes written and any error encountered during the write
+// operation.
+func (resp *Response) Write(data []byte) (int, error) {
+	if len(data) == 0 {
+		return 0, nil
+	}
+
+	if err := resp.checkBodySize(len(data)); err != nil {
+		return 0, err
+	}
+	return resp.body.Write(data)
+}
+
+// WriteString writes the given string data to the response body. It appends the string to any
+// existing content in the body. It returns the number of bytes written and any error encountered
+// during the write operation.
+func (resp *Response) WriteString(data string) (int, error) {
+	if len(data) == 0 {
+		return 0, nil
+	}
+
+	if err := resp.checkBodySize(len(data)); err != nil {
+		return 0, err
+	}
+	return resp.body.WriteString(data)
+}
+
+// Written returns the current content of the response body as a byte slice. It creates a copy of
+// the underlying buffer to ensure that the returned data is not modified externally.
+func (resp *Response) Written() []byte {
+	data := resp.body.Bytes()
+	copied := make([]byte, len(data))
+	copy(copied, data)
+	return copied
+}
+
 // Response returns the underlying http.ResponseWriter associated with this Response. This can be
 // used to access any additional information or functionality provided by the http.ResponseWriter
 // that is not exposed through the Response interface.
@@ -153,3 +176,18 @@ func (resp *Response) send() error {
 
 	return nil
 }
+
+// checkBodySize checks if the size of the response body exceeds the maximum allowed size. If the
+// size exceeds the limit, it returns an error indicating that the response body is too large.
+func (resp *Response) checkBodySize(size int) error {
+	maxBytes := resp.maxBodyBytes.Load()
+
+	if maxBytes == MaxResponseBodyBytesUnlimited {
+		return nil
+	}
+
+	if int64(resp.body.Len())+int64(size) > maxBytes {
+		return ErrResponseTooLarge
+	}
+	return nil
+}

response_test.go

@@ -219,3 +219,126 @@ func TestResponse_SetMaxBodyBytes(t *testing.T) {
 		t.Fatalf("expected ErrResponseTooLarge when maxBodyBytes=0, got n=%d err=%v", n, err)
 	}
 }
+
+func TestResponse_WriteString(t *testing.T) {
+	rr := httptest.NewRecorder()
+	resp := newResponse(nil, rr)
+
+	n, err := resp.WriteString("")
+	if err != nil || n != 0 {
+		t.Fatalf("WriteString empty: n=%d err=%v", n, err)
+	}
+
+	n, err = resp.WriteString("hello")
+	if err != nil || n != 5 {
+		t.Fatalf("WriteString: n=%d err=%v", n, err)
+	}
+
+	if resp.body.String() != "hello" {
+		t.Fatalf("expected body 'hello', got %q", resp.body.String())
+	}
+
+	n, err = resp.WriteString(" world")
+	if err != nil || n != 6 {
+		t.Fatalf("WriteString: n=%d err=%v", n, err)
+	}
+
+	if resp.body.String() != "hello world" {
+		t.Fatalf("expected body 'hello world', got %q", resp.body.String())
+	}
+
+	resp.send()
+
+	if rr.Body.String() != "hello world" {
+		t.Fatalf("expected sent body 'hello world', got %q", rr.Body.String())
+	}
+}
+
+func TestResponse_WriteString_LargeBody(t *testing.T) {
+	app := New(WithMaxResponseBodyBytes(11))
+	rr := httptest.NewRecorder()
+	resp := newResponse(app, rr)
+
+	n, err := resp.WriteString("hello")
+	if err != nil || n != 5 {
+		t.Fatalf("WriteString: n=%d err=%v", n, err)
+	}
+
+	n, err = resp.WriteString(" world!")
+	if err == nil || !errors.Is(err, ErrResponseTooLarge) || n != 0 {
+		t.Fatalf("expected ErrResponseTooLarge, got n=%d err=%v", n, err)
+	}
+
+	// body should remain unchanged after failed write
+	if resp.body.String() != "hello" {
+		t.Fatalf("expected body 'hello' after failed write, got %q", resp.body.String())
+	}
+
+	resp.send()
+
+	if rr.Body.String() != "hello" {
+		t.Fatalf("expected sent body 'hello', got %q", rr.Body.String())
+	}
+}
+
+func TestResponse_WriteString_Unlimited(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping test in short mode")
+	}
+
+	app := New(WithMaxResponseBodyBytes(MaxResponseBodyBytesUnlimited))
+	rr := httptest.NewRecorder()
+	resp := newResponse(app, rr)
+
+	largeStr := string(make([]byte, MaxResponseBodyBytesDefault+1))
+
+	n, err := resp.WriteString(largeStr)
+	if err != nil || n != len(largeStr) {
+		t.Fatalf("WriteString unlimited: n=%d err=%v", n, err)
+	}
+}
+
+func TestResponse_Size(t *testing.T) {
+	rr := httptest.NewRecorder()
+	resp := newResponse(nil, rr)
+
+	if resp.Size() != 0 {
+		t.Fatalf("expected Size 0, got %d", resp.Size())
+	}
+
+	resp.Write([]byte("hello"))
+	if resp.Size() != 5 {
+		t.Fatalf("expected Size 5, got %d", resp.Size())
+	}
+
+	resp.WriteString(" world")
+	if resp.Size() != 11 {
+		t.Fatalf("expected Size 11, got %d", resp.Size())
+	}
+}
+
+func TestResponse_Written(t *testing.T) {
+	rr := httptest.NewRecorder()
+	resp := newResponse(nil, rr)
+
+	if len(resp.Written()) != 0 {
+		t.Fatalf("expected empty Written, got %v", resp.Written())
+	}
+
+	resp.Write([]byte("hello"))
+	if string(resp.Written()) != "hello" {
+		t.Fatalf("expected Written 'hello', got %q", string(resp.Written()))
+	}
+
+	// verify that mutating the returned slice does not affect the body
+	data := resp.Written()
+	data[0] = 'x'
+	if string(resp.Written()) != "hello" {
+		t.Fatalf("Written should return a copy; mutation must not affect body, got %q", string(resp.Written()))
+	}
+
+	resp.WriteString(" world")
+	if string(resp.Written()) != "hello world" {
+		t.Fatalf("expected Written 'hello world', got %q", string(resp.Written()))
+	}
+}