Files
reproxy/vendor/github.com/go-pkgz/rest/sizelimit.go
2022-01-06 00:16:18 -06:00

41 lines
909 B
Go

package rest
import (
"bytes"
"io"
"net/http"
)
// SizeLimit middleware checks if body size is above the limit and returns StatusRequestEntityTooLarge (413)
func SizeLimit(size int64) func(http.Handler) http.Handler {
return func(h http.Handler) http.Handler {
fn := func(w http.ResponseWriter, r *http.Request) {
// check ContentLength
if r.ContentLength > size {
w.WriteHeader(http.StatusRequestEntityTooLarge)
return
}
// check size of the actual body
content, err := io.ReadAll(io.LimitReader(r.Body, size+1))
if err != nil {
w.WriteHeader(http.StatusServiceUnavailable)
return
}
_ = r.Body.Close() // the original body already consumed
if int64(len(content)) > size {
w.WriteHeader(http.StatusRequestEntityTooLarge)
return
}
r.Body = io.NopCloser(bytes.NewReader(content))
h.ServeHTTP(w, r)
}
return http.HandlerFunc(fn)
}
}