33 lines
1.3 KiB
Go
33 lines
1.3 KiB
Go
package app
|
|
|
|
import (
|
|
"git.yuntuo.hk/k2safe/go-omni/httpx"
|
|
"git.yuntuo.hk/k2safe/go-omni/response"
|
|
"net/http"
|
|
)
|
|
|
|
func (a *App) bodyHandler(maxBytes int64, encoder response.Encoder, handler httpx.Handler, middleware ...httpx.Middleware) http.Handler {
|
|
checked := httpx.LimitBody(maxBytes)(handler)
|
|
if encoder == nil {
|
|
encoder = a.responseEncoder
|
|
}
|
|
next := a.adaptResponse(encoder, checked, middleware...)
|
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.Body != nil {
|
|
r = r.WithContext(r.Context())
|
|
r.Body = http.MaxBytesReader(w, r.Body, maxBytes)
|
|
}
|
|
next.ServeHTTP(w, r)
|
|
})
|
|
}
|
|
|
|
// HandleWithBodyLimit caps the reader before application middleware and preserves
|
|
// the response policy, including auth failures and parameter errors.
|
|
func (a *App) HandleWithBodyLimit(method, path string, maxBytes int64, encoder response.Encoder, handler httpx.Handler, middleware ...httpx.Middleware) {
|
|
a.router.Handle(method, path, a.bodyHandler(maxBytes, encoder, handler, middleware...))
|
|
}
|
|
func (g *RouterGroup) HandleWithBodyLimit(method, path string, maxBytes int64, encoder response.Encoder, handler httpx.Handler, middleware ...httpx.Middleware) {
|
|
combined := append(append([]httpx.Middleware{}, g.middlewares...), middleware...)
|
|
g.group.Handle(method, path, g.app.bodyHandler(maxBytes, encoder, handler, combined...))
|
|
}
|