github.com/lingyao2333/mo-zero@v1.4.1/rest/handler/maxbyteshandler.go (about)

     1  package handler
     2  
     3  import (
     4  	"net/http"
     5  
     6  	"github.com/lingyao2333/mo-zero/rest/internal"
     7  )
     8  
     9  // MaxBytesHandler returns a middleware that limit reading of http request body.
    10  func MaxBytesHandler(n int64) func(http.Handler) http.Handler {
    11  	if n <= 0 {
    12  		return func(next http.Handler) http.Handler {
    13  			return next
    14  		}
    15  	}
    16  
    17  	return func(next http.Handler) http.Handler {
    18  		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    19  			if r.ContentLength > n {
    20  				internal.Errorf(r, "request entity too large, limit is %d, but got %d, rejected with code %d",
    21  					n, r.ContentLength, http.StatusRequestEntityTooLarge)
    22  				w.WriteHeader(http.StatusRequestEntityTooLarge)
    23  			} else {
    24  				next.ServeHTTP(w, r)
    25  			}
    26  		})
    27  	}
    28  }