61 lines
1.9 KiB
Go
61 lines
1.9 KiB
Go
package httpx
|
|
|
|
import (
|
|
"errors"
|
|
apierrors "git.yuntuo.hk/k2safe/go-omni/errors"
|
|
"git.yuntuo.hk/k2safe/go-omni/response"
|
|
"net/http"
|
|
)
|
|
|
|
// LimitBody is useful for handwritten routes. Built-in App.HandleWithBodyLimit
|
|
// additionally installs the reader cap before global middleware can read it.
|
|
func LimitBody(maxBytes int64) Middleware {
|
|
if maxBytes <= 0 || maxBytes > 1<<40 {
|
|
panic("httpx: body limit must be 1..1 TiB")
|
|
}
|
|
return func(next Handler) Handler {
|
|
return func(c *Context) error {
|
|
r := c.Request()
|
|
if r.ContentLength > maxBytes {
|
|
return bodyTooLarge()
|
|
}
|
|
if r.Body != nil {
|
|
r.Body = http.MaxBytesReader(c.Raw(), r.Body, maxBytes)
|
|
}
|
|
defer CleanupMultipart(r)
|
|
err := next(c)
|
|
var limitErr *http.MaxBytesError
|
|
if errors.As(err, &limitErr) {
|
|
return bodyTooLarge()
|
|
}
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
func bodyTooLarge() error { return apierrors.New(413, "BODY_TOO_LARGE", "request body exceeds limit") }
|
|
|
|
// CleanupMultipart removes disk spill files created by ParseMultipartForm.
|
|
// Callers still close individual multipart files. Streaming MultipartReader
|
|
// handlers own any files they create themselves.
|
|
func CleanupMultipart(r *http.Request) {
|
|
if r != nil && r.MultipartForm != nil {
|
|
_ = r.MultipartForm.RemoveAll()
|
|
}
|
|
}
|
|
|
|
type BodyLimitRegistrar interface {
|
|
HandleWithBodyLimit(string, string, int64, response.Encoder, Handler, ...Middleware)
|
|
}
|
|
|
|
func HandleWithBodyLimit(router interface {
|
|
Handle(string, string, Handler, ...Middleware)
|
|
}, method, path string, maxBytes int64, encoder response.Encoder, handler Handler, middleware ...Middleware) {
|
|
if r, ok := router.(BodyLimitRegistrar); ok {
|
|
r.HandleWithBodyLimit(method, path, maxBytes, encoder, handler, middleware...)
|
|
return
|
|
}
|
|
// Custom registrars without the early hook cap at route middleware entry.
|
|
all := append([]Middleware{LimitBody(maxBytes)}, middleware...)
|
|
HandleWithResponse(router, method, path, encoder, handler, all...)
|
|
}
|