github.com/grafana/pyroscope@v1.18.0/pkg/util/body/limit.go (about)

     1  package body
     2  
     3  import (
     4  	"context"
     5  	"net/http"
     6  
     7  	"github.com/grafana/dskit/tenant"
     8  )
     9  
    10  type Limits interface {
    11  	IngestionBodyLimitBytes(tenantID string) int64
    12  }
    13  
    14  func getMaxBodySize(ctx context.Context, limits Limits) int64 {
    15  	tenantID, err := tenant.TenantID(ctx)
    16  	if err != nil {
    17  		return 0
    18  	}
    19  
    20  	return limits.IngestionBodyLimitBytes(tenantID)
    21  }
    22  
    23  func NewSizeLimitHandler(limits Limits) func(h http.Handler) http.Handler {
    24  	return func(h http.Handler) http.Handler {
    25  		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    26  			maxBodySize := getMaxBodySize(r.Context(), limits)
    27  
    28  			if maxBodySize > 0 {
    29  				r.Body = http.MaxBytesReader(w, r.Body, maxBodySize)
    30  			}
    31  
    32  			h.ServeHTTP(w, r)
    33  		})
    34  	}
    35  }