Merge pull request #28 from infinite-iroha/feat/memory-improvements-and-request-size-limit

Feat/memory improvements and request size limit

context.go

@@ -4,11 +4,13 @@ import (
 	"bytes"
 	"context"
 	"encoding/gob"
+	"encoding/xml" // Added for XML binding
 	"errors"
 	"fmt"
 	"html/template"
 	"io"
 	"math"
+	"mime" // Added for Content-Type parsing in ShouldBind
 	"net"
 	"net/http"
 	"net/netip"
@@ -19,6 +21,7 @@ import (
 
 	"github.com/fenthope/reco"
 	"github.com/go-json-experiment/json"
+	"github.com/gorilla/schema" // Added for form binding
 
 	"github.com/WJQSERVER-STUDIO/go-utils/copyb"
 	"github.com/WJQSERVER-STUDIO/httpc"
@@ -298,21 +301,24 @@ func (c *Context) HTML(code int, name string, obj interface{}) {
 	c.Writer.Header().Set("Content-Type", "text/html; charset=utf-8")
 	c.Writer.WriteHeader(code)
 
-	if c.engine != nil && c.engine.HTMLRender != nil {
+	if c.engine == nil || c.engine.HTMLRender == nil {
-		// 假设 HTMLRender 是一个 *template.Template 实例
+		errMsg := "HTML renderer not configured"
-		if tpl, ok := c.engine.HTMLRender.(*template.Template); ok {
+		if c.engine != nil && c.engine.LogReco != nil {
-			err := tpl.ExecuteTemplate(c.Writer, name, obj)
+			c.engine.LogReco.Error("[Context.HTML] HTMLRender not configured on engine")
-			if err != nil {
+		} else {
-				c.AddError(fmt.Errorf("failed to render HTML template '%s': %w", name, err))
+			// Fallback logging if LogReco is also nil, though unlikely if engine is not nil
-				//c.String(http.StatusInternalServerError, "Internal Server Error: Failed to render HTML template")
+			// log.Println("[Context.HTML] HTMLRender not configured on engine")
-				c.ErrorUseHandle(http.StatusInternalServerError, fmt.Errorf("failed to render HTML template '%s': %w", name, err))
 		}
+		c.ErrorUseHandle(http.StatusInternalServerError, errors.New(errMsg))
 		return
 	}
-		// 可以扩展支持其他渲染器接口
+
+	err := c.engine.HTMLRender.Render(c.Writer, name, obj, c)
+	if err != nil {
+		renderErr := fmt.Errorf("failed to render HTML template '%s': %w", name, err)
+		c.AddError(renderErr)
+		c.ErrorUseHandle(http.StatusInternalServerError, renderErr)
 	}
-	// 默认简单输出，用于未配置 HTMLRender 的情况
-	c.Writer.Write([]byte(fmt.Sprintf("<!-- HTML rendered for %s -->\n<pre>%v</pre>", name, obj)))
 }
 
 // Redirect 执行 HTTP 重定向
@@ -330,34 +336,207 @@ func (c *Context) ShouldBindJSON(obj interface{}) error {
 	if c.Request.Body == nil {
 		return errors.New("request body is empty")
 	}
-	/*
+	// defer c.Request.Body.Close() // 通常由调用方或中间件确保关闭，但如果这里是唯一消耗点，可以考虑
-		decoder := json.NewDecoder(c.Request.Body)
+
-		if err := decoder.Decode(obj); err != nil {
+	var reader io.Reader = c.Request.Body
-			return fmt.Errorf("json binding error: %w", err)
+	if c.engine != nil && c.engine.MaxRequestBodySize > 0 {
+		if c.Request.ContentLength != -1 && c.Request.ContentLength > c.engine.MaxRequestBodySize {
+			return fmt.Errorf("request body size (%d bytes) exceeds configured limit (%d bytes)", c.Request.ContentLength, c.engine.MaxRequestBodySize)
 		}
-	*/
+		// 注意：http.MaxBytesReader(nil, ...) 中的 nil ResponseWriter 参数意味着当超出限制时，
-	err := json.UnmarshalRead(c.Request.Body, obj)
+		// MaxBytesReader 会直接返回错误，而不会尝试写入 HTTP 错误响应。这对于 API 来说是合适的。
+		reader = http.MaxBytesReader(nil, c.Request.Body, c.engine.MaxRequestBodySize)
+	}
+
+	err := json.UnmarshalRead(reader, obj)
 	if err != nil {
+		// 检查错误类型是否为 http.MaxBytesError
+		var maxBytesErr *http.MaxBytesError
+		if errors.As(err, &maxBytesErr) {
+			return fmt.Errorf("request body size exceeds configured limit (%d bytes): %w", c.engine.MaxRequestBodySize, err)
+		}
+		// 检查是否是 json 相关的错误，可能需要更细致的错误处理
+		// 例如，json.SyntaxError, json.UnmarshalTypeError 等
 		return fmt.Errorf("json binding error: %w", err)
 	}
 	return nil
 }
 
+// ShouldBindXML 尝试将请求体中的 XML 数据绑定到 obj。
+func (c *Context) ShouldBindXML(obj interface{}) error {
+	if c.Request == nil || c.Request.Body == nil {
+		return errors.New("request body is empty for XML binding")
+	}
+	// defer c.Request.Body.Close() // Caller is responsible for closing the body
+
+	var reader io.Reader = c.Request.Body
+	if c.engine != nil && c.engine.MaxRequestBodySize > 0 {
+		if c.Request.ContentLength != -1 && c.Request.ContentLength > c.engine.MaxRequestBodySize {
+			return fmt.Errorf("request body size (%d bytes) exceeds XML binding limit (%d bytes)", c.Request.ContentLength, c.engine.MaxRequestBodySize)
+		}
+		reader = http.MaxBytesReader(nil, c.Request.Body, c.engine.MaxRequestBodySize)
+	}
+
+	decoder := xml.NewDecoder(reader)
+	if err := decoder.Decode(obj); err != nil {
+		// Check for MaxBytesError specifically
+		var maxBytesErr *http.MaxBytesError
+		if errors.As(err, &maxBytesErr) {
+			return fmt.Errorf("request body size exceeds XML binding limit (%d bytes): %w", c.engine.MaxRequestBodySize, err)
+		}
+		return fmt.Errorf("xml binding error: %w", err)
+	}
+	return nil
+}
+
+// ShouldBindQuery 尝试将 URL 查询参数绑定到 obj。
+// 它使用 gorilla/schema 来执行绑定。
+func (c *Context) ShouldBindQuery(obj interface{}) error {
+	if c.Request == nil {
+		return errors.New("request is nil")
+	}
+
+	if c.queryCache == nil {
+		c.queryCache = c.Request.URL.Query()
+	}
+	values := c.queryCache
+
+	if len(values) == 0 {
+		// No query parameters to bind
+		return nil
+	}
+
+	decoder := schema.NewDecoder()
+	// decoder.IgnoreUnknownKeys(true) // Optional
+
+	if err := decoder.Decode(obj, values); err != nil {
+		return fmt.Errorf("query parameter binding error using schema: %w", err)
+	}
+
+	return nil
+}
+
 // ShouldBind 尝试将请求体绑定到各种类型（JSON, Form, XML 等）
-// 这是一个复杂的通用绑定接口，通常根据 Content-Type 或其他头部来判断绑定方式
+// 根据请求的 Content-Type 自动选择合适的绑定器。
-// 预留接口，可根据项目需求进行扩展
 func (c *Context) ShouldBind(obj interface{}) error {
-	// TODO: 完整的通用绑定逻辑
+	if c.Request == nil {
-	// 可以根据 c.Request.Header.Get("Content-Type") 来判断是 JSON, Form, XML 等
+		return errors.New("request is nil for binding")
-	// 例如：
+	}
-	// contentType := c.Request.Header.Get("Content-Type")
+
-	// if strings.HasPrefix(contentType, "application/json") {
+	// If there's no body, no binding from body can occur.
-	//     return c.ShouldBindJSON(obj)
+	if c.Request.Body == nil || c.Request.Body == http.NoBody {
-	// }
+		// Consider if query binding should be attempted for GET requests by default.
-	// if strings.HasPrefix(contentType, "application/x-www-form-urlencoded") || strings.HasPrefix(contentType, "multipart/form-data") {
+		// For now, if no body, assume successful (empty) binding from body perspective.
-	//     return c.ShouldBindForm(obj) // 需要实现 ShouldBindForm
+		return nil
-	// }
+	}
-	return errors.New("generic binding not fully implemented yet, implement based on Content-Type")
+
+	contentType := c.ContentType() // This uses c.GetReqHeader("Content-Type")
+	if contentType == "" {
+		// If there is a body (ContentLength > 0 or chunked) but no Content-Type, this is an issue.
+		if c.Request.ContentLength > 0 || len(c.Request.TransferEncoding) > 0 {
+			return errors.New("missing Content-Type header for request body binding")
+		}
+		// If no Content-Type and no actual body content indicated, effectively no body to bind.
+		return nil
+	}
+
+	mimeType, _, err := mime.ParseMediaType(contentType)
+	if err != nil {
+		return fmt.Errorf("error parsing Content-Type header '%s': %w", contentType, err)
+	}
+
+	switch mimeType {
+	case "application/json":
+		return c.ShouldBindJSON(obj)
+	case "application/xml", "text/xml":
+		return c.ShouldBindXML(obj)
+	case "application/x-www-form-urlencoded":
+		return c.ShouldBindForm(obj)
+	case "multipart/form-data":
+		return c.ShouldBindForm(obj) // ShouldBindForm handles multipart fields
+	default:
+		return fmt.Errorf("unsupported Content-Type for binding: %s", mimeType)
+	}
+}
+
+// ShouldBindForm 尝试将请求体中的表单数据绑定到 obj。
+// 它使用 gorilla/schema 来执行绑定。
+// 注意：此方法期望请求体是 application/x-www-form-urlencoded 或 multipart/form-data 类型。
+// 调用此方法前，应确保请求的 Content-Type 是合适的。
+func (c *Context) ShouldBindForm(obj interface{}) error {
+	if c.Request == nil {
+		return errors.New("request is nil")
+	}
+
+	// ParseMultipartForm populates c.Request.PostForm and c.Request.MultipartForm.
+	// defaultMemory is used to limit the size of memory used for storing file parts.
+	// If the form data exceeds this, it will be stored in temporary files.
+	// MaxBytesReader applied earlier (if any) would have limited the total body size.
+	if err := c.Request.ParseMultipartForm(defaultMemory); err != nil {
+		// Ignore "http: multipart handled by ParseMultipartForm" error, which means it was already parsed.
+		// This can happen if a middleware or previous ShouldBind call already parsed the form.
+		// Other errors (e.g., I/O errors during parsing) should be returned.
+		// Note: Gorilla schema might not need this if it directly uses r.Form or r.PostForm
+		// which are populated by ParseForm/ParseMultipartForm.
+		// For now, we ensure it's parsed. A more specific check might be needed if this causes issues.
+		// A common error to ignore here is `http.ErrNotMultipart` if the content type isn't multipart,
+		// as ParseMultipartForm expects that. ParseForm might be more general if we only expect
+		// x-www-form-urlencoded, but ParseMultipartForm handles both.
+		// Let's proceed and let schema decoder handle empty PostForm if parsing wasn't applicable.
+		// However, a direct "request body too large" from MaxBytesReader should have priority if it happened before.
+		// This specific error from ParseMultipartForm might relate to parts of a valid multipart form being too large for memory,
+		// not the overall request size.
+		// For simplicity in this step, we'll return the error unless it's a known "already parsed" scenario (which is not standard).
+		// A better approach for "already parsed" would be to check if c.Request.PostForm is already populated.
+		if c.formCache == nil && c.Request.PostForm == nil { // Attempt to parse only if not already cached or populated
+			if perr := c.Request.ParseMultipartForm(defaultMemory); perr != nil {
+				// http.ErrNotMultipart is returned if Content-Type is not multipart/form-data
+				// For x-www-form-urlencoded, ParseForm() is implicitly called by accessing PostForm
+				// Let's try to populate PostForm if it's not already
+				if c.Request.PostForm == nil {
+					if perr2 := c.Request.ParseForm(); perr2 != nil {
+						return fmt.Errorf("form parse error (ParseForm): %w", perr2)
+					}
+				}
+				// If it was not multipart and ParseForm also failed or PostForm is still nil,
+				// then we might have an issue. However, gorilla/schema works on `url.Values`
+				// which `c.Request.PostForm` provides.
+				// If `Content-Type` was not form-like, `PostForm` would be empty.
+			}
+		}
+		// If `c.formCache` is not nil, `PostForm()` would have already tried parsing.
+		// We will use `c.Request.PostForm` which gets populated by `ParseMultipartForm` or `ParseForm`.
+	}
+
+	// Initialize schema decoder
+	decoder := schema.NewDecoder()
+	// decoder.IgnoreUnknownKeys(true) // Optional: if you want to ignore fields in form not in struct
+
+	// Get form values. c.Request.PostForm includes values from both
+	// application/x-www-form-urlencoded and multipart/form-data bodies.
+	// It needs to be called after ParseMultipartForm or ParseForm.
+	// Accessing c.Request.PostForm itself can trigger ParseForm if not already parsed and content type is x-www-form-urlencoded.
+	if err := c.Request.ParseForm(); err != nil && c.Request.PostForm == nil {
+		 // If ParseForm itself errors and PostForm is still nil, then return error.
+		 // This ensures that for x-www-form-urlencoded, parsing is attempted.
+		 // ParseMultipartForm handles multipart, and its error is handled above.
+		return fmt.Errorf("form parse error (PostForm init): %w", err)
+	}
+
+	values := c.Request.PostForm
+	if len(values) == 0 {
+		// If PostForm is empty, there's nothing to bind from the POST body.
+		// This is not necessarily an error for schema.Decode, it will just bind zero values.
+		// Depending on requirements, one might want to return an error here if binding is mandatory.
+		// For now, we let schema.Decode handle it (it will likely do nothing or bind zero values).
+	}
+
+	// Decode the form values into the object
+	if err := decoder.Decode(obj, values); err != nil {
+		return fmt.Errorf("form binding error using schema: %w", err)
+	}
+
+	return nil
 }
 
 // AddError 添加一个错误到 Context
@@ -441,13 +620,31 @@ func (c *Context) GetReqBody() io.ReadCloser {
 // 注意：请求体只能读取一次
 func (c *Context) GetReqBodyFull() ([]byte, error) {
 	if c.Request.Body == nil {
-		return nil, nil
+		return nil, nil // 或者返回一个错误: errors.New("request body is nil")
 	}
 	defer c.Request.Body.Close() // 确保请求体被关闭
-	data, err := io.ReadAll(c.Request.Body)
+
+	var reader io.Reader = c.Request.Body
+	if c.engine != nil && c.engine.MaxRequestBodySize > 0 {
+		if c.Request.ContentLength != -1 && c.Request.ContentLength > c.engine.MaxRequestBodySize {
+			err := fmt.Errorf("request body size (%d bytes) exceeds configured limit (%d bytes)", c.Request.ContentLength, c.engine.MaxRequestBodySize)
+			c.AddError(err)
+			return nil, err
+		}
+		reader = http.MaxBytesReader(nil, c.Request.Body, c.engine.MaxRequestBodySize)
+	}
+
+	data, err := io.ReadAll(reader)
 	if err != nil {
-		c.AddError(fmt.Errorf("failed to read request body: %w", err))
+		// 检查错误类型是否为 http.MaxBytesError，如果是，则表示超出了限制
-		return nil, fmt.Errorf("failed to read request body: %w", err)
+		var maxBytesErr *http.MaxBytesError
+		if errors.As(err, &maxBytesErr) {
+			err = fmt.Errorf("request body size exceeds configured limit (%d bytes): %w", c.engine.MaxRequestBodySize, err)
+		} else {
+			err = fmt.Errorf("failed to read request body: %w", err)
+		}
+		c.AddError(err)
+		return nil, err
 	}
 	return data, nil
 }
@@ -455,13 +652,30 @@ func (c *Context) GetReqBodyFull() ([]byte, error) {
 // 类似 GetReqBodyFull, 返回 *bytes.Buffer
 func (c *Context) GetReqBodyBuffer() (*bytes.Buffer, error) {
 	if c.Request.Body == nil {
-		return nil, nil
+		return nil, nil // 或者返回一个错误: errors.New("request body is nil")
 	}
 	defer c.Request.Body.Close() // 确保请求体被关闭
-	data, err := io.ReadAll(c.Request.Body)
+
+	var reader io.Reader = c.Request.Body
+	if c.engine != nil && c.engine.MaxRequestBodySize > 0 {
+		if c.Request.ContentLength != -1 && c.Request.ContentLength > c.engine.MaxRequestBodySize {
+			err := fmt.Errorf("request body size (%d bytes) exceeds configured limit (%d bytes)", c.Request.ContentLength, c.engine.MaxRequestBodySize)
+			c.AddError(err)
+			return nil, err
+		}
+		reader = http.MaxBytesReader(nil, c.Request.Body, c.engine.MaxRequestBodySize)
+	}
+
+	data, err := io.ReadAll(reader)
 	if err != nil {
-		c.AddError(fmt.Errorf("failed to read request body: %w", err))
+		var maxBytesErr *http.MaxBytesError
-		return nil, fmt.Errorf("failed to read request body: %w", err)
+		if errors.As(err, &maxBytesErr) {
+			err = fmt.Errorf("request body size exceeds configured limit (%d bytes): %w", c.engine.MaxRequestBodySize, err)
+		} else {
+			err = fmt.Errorf("failed to read request body: %w", err)
+		}
+		c.AddError(err)
+		return nil, err
 	}
 	return bytes.NewBuffer(data), nil
 }

engine.go

@@ -9,15 +9,31 @@ import (
 	"runtime"
 	"strings"
 
+	"html/template"
+	"io"
 	"net/http"
 	"path"
-
 	"sync"
 
 	"github.com/WJQSERVER-STUDIO/httpc"
 	"github.com/fenthope/reco"
 )
 
+// HTMLRender defines the interface for HTML rendering.
+type HTMLRender interface {
+	Render(writer io.Writer, name string, data interface{}, c *Context) error
+}
+
+// DefaultHTMLRenderer is a basic implementation of HTMLRender using html/template.
+type DefaultHTMLRenderer struct {
+	Templates *template.Template
+}
+
+// Render executes the template and writes to the writer.
+func (r *DefaultHTMLRenderer) Render(writer io.Writer, name string, data interface{}, c *Context) error {
+	return r.Templates.ExecuteTemplate(writer, name, data)
+}
+
 // Last 返回链中的最后一个处理函数
 // 如果链为空,则返回 nil
 func (c HandlersChain) Last() HandlerFunc {
@@ -50,7 +66,7 @@ type Engine struct {
 
 	LogReco *reco.Logger
 
-	HTMLRender interface{} // 用于 HTML 模板渲染,可以设置为 *template.Template 或自定义渲染器接口
+	HTMLRender HTMLRender // 用于 HTML 模板渲染
 
 	routesInfo []RouteInfo // 存储所有注册的路由信息
 
@@ -74,6 +90,8 @@ type Engine struct {
 	// 如果设置了此回调,它将优先于 ServerConfigurator 被用于 HTTPS 服务器
 	// 如果未设置,HTTPS 服务器将回退使用 ServerConfigurator (如果已设置)
 	TLSServerConfigurator func(*http.Server)
+
+	MaxRequestBodySize int64 // 限制读取Body的最大字节数
 }
 
 type ErrorHandle struct {
@@ -87,12 +105,39 @@ type ErrorHandler func(c *Context, code int, err error)
 func defaultErrorHandle(c *Context, code int, err error) { // 检查客户端是否已断开连接
 	select {
 	case <-c.Request.Context().Done():
-
+		// 客户端断开连接,无需进一步处理
 		return
 	default:
+		// 检查响应是否已经写入
 		if c.Writer.Written() {
 			return
 		}
+
+		// 收集错误信息用于日志记录
+		primaryErrStr := "none"
+		if err != nil {
+			primaryErrStr = err.Error()
+		}
+
+		var collectedErrors []string
+		for _, e := range c.GetErrors() {
+			collectedErrors = append(collectedErrors, e.Error())
+		}
+		collectedErrorsStr := strings.Join(collectedErrors, "; ")
+		if collectedErrorsStr == "" {
+			collectedErrorsStr = "none"
+		}
+
+		// 记录错误日志
+		logMessage := fmt.Sprintf("[Touka ErrorHandler] Request: [%s] %s | Primary Error: %s | Collected Errors: %s",
+			c.Request.Method, c.Request.URL.Path, primaryErrStr, collectedErrorsStr)
+
+		if c.engine != nil && c.engine.LogReco != nil {
+			c.engine.LogReco.Error(logMessage)
+		} else {
+			log.Println(logMessage) // Fallback to standard logger
+		}
+
 		// 输出json 状态码与状态码对应描述
 		var errMsg string
 		if err != nil {
@@ -160,6 +205,7 @@ func New() *Engine {
 		noRoutes:              make(HandlersChain, 0),
 		ServerConfigurator:    nil,
 		TLSServerConfigurator: nil,
+		MaxRequestBodySize:    10 * 1024 * 1024, // 默认 10MB
 	}
 	//engine.SetProtocols(GetDefaultProtocolsConfig())
 	engine.SetDefaultProtocols()
@@ -189,6 +235,23 @@ func Default() *Engine {
 
 // === 外部操作方法 ===
 
+// LoadHTMLGlob loads HTML templates from a glob pattern and sets them as the HTML renderer.
+func (engine *Engine) LoadHTMLGlob(pattern string) {
+	tpl := template.Must(template.ParseGlob(pattern))
+	engine.HTMLRender = &DefaultHTMLRenderer{Templates: tpl}
+}
+
+// SetHTMLTemplate sets a custom *template.Template as the HTML renderer.
+// This will wrap the *template.Template with the DefaultHTMLRenderer.
+func (engine *Engine) SetHTMLTemplate(tpl *template.Template) {
+	engine.HTMLRender = &DefaultHTMLRenderer{Templates: tpl}
+}
+
+// SetMaxRequestBodySize 设置读取Body的最大字节数
+func (engine *Engine) SetMaxRequestBodySize(size int64) {
+	engine.MaxRequestBodySize = size
+}
+
 // SetServerConfigurator 设置一个函数,该函数将在任何 HTTP 或 HTTPS 服务器
 // (通过 RunShutdown, RunTLS, RunTLSRedir) 启动前被调用,
 // 允许用户对底层的 *http.Server 实例进行自定义配置

go.mod

@@ -9,4 +9,7 @@ require (
 	github.com/go-json-experiment/json v0.0.0-20250517221953-25912455fbc8
 )
 
-require github.com/valyala/bytebufferpool v1.0.0 // indirect
+require (
+	github.com/gorilla/schema v1.4.1 // indirect
+	github.com/valyala/bytebufferpool v1.0.0 // indirect
+)

go.sum

@@ -6,5 +6,7 @@ github.com/fenthope/reco v0.0.3 h1:RmnQ0D9a8PWtwOODawitTe4BztTnS9wYwrDbipISNq4=
 github.com/fenthope/reco v0.0.3/go.mod h1:mDkGLHte5udWTIcjQTxrABRcf56SSdxBOCLgrRDwI/Y=
 github.com/go-json-experiment/json v0.0.0-20250517221953-25912455fbc8 h1:o8UqXPI6SVwQt04RGsqKp3qqmbOfTNMqDrWsc4O47kk=
 github.com/go-json-experiment/json v0.0.0-20250517221953-25912455fbc8/go.mod h1:TiCD2a1pcmjd7YnhGH0f/zKNcCD06B029pHhzV23c2M=
+github.com/gorilla/schema v1.4.1 h1:jUg5hUjCSDZpNGLuXQOgIWGdlgrIdYvgQ0wZtdK1M3E=
+github.com/gorilla/schema v1.4.1/go.mod h1:Dg5SSm5PV60mhF2NFaTV1xuYYj8tV8NOPRo4FggUMnM=
 github.com/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6KllzawFIhcdPw=
 github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc=