pocketbase/apis/middlewares_body_limit.go

package apis

import (
	"io"
	"net/http"

	"github.com/pocketbase/pocketbase/core"
	"github.com/pocketbase/pocketbase/tools/hook"
	"github.com/pocketbase/pocketbase/tools/router"
)

var ErrRequestEntityTooLarge = router.NewApiError(http.StatusRequestEntityTooLarge, "Request entity too large", nil)

const DefaultMaxBodySize int64 = 32 << 20 // @todo consider replacing with router.DefaultMaxMemory

const (
	DefaultBodyLimitMiddlewareId       = "pbBodyLimit"
	DefaultBodyLimitMiddlewarePriority = DefaultRateLimitMiddlewarePriority + 10
)

// BodyLimit returns a middleware handler that changes the default request body size limit.
//
// If limitBytes <= 0, no limit is applied.
//
// Otherwise, if the request body size exceeds the configured limitBytes,
// it sends 413 error response.
func BodyLimit(limitBytes int64) *hook.Handler[*core.RequestEvent] {
	return &hook.Handler[*core.RequestEvent]{
		Id:       DefaultBodyLimitMiddlewareId,
		Priority: DefaultBodyLimitMiddlewarePriority,
		Func: func(e *core.RequestEvent) error {
			err := applyBodyLimit(e, limitBytes)
			if err != nil {
				return err
			}

			return e.Next()
		},
	}
}

func dynamicCollectionBodyLimit(collectionPathParam string) *hook.Handler[*core.RequestEvent] {
	if collectionPathParam == "" {
		collectionPathParam = "collection"
	}

	return &hook.Handler[*core.RequestEvent]{
		Id:       DefaultBodyLimitMiddlewareId,
		Priority: DefaultBodyLimitMiddlewarePriority,
		Func: func(e *core.RequestEvent) error {
			collection, err := e.App.FindCachedCollectionByNameOrId(e.Request.PathValue(collectionPathParam))
			if err != nil {
				return e.NotFoundError("Missing or invalid collection context.", err)
			}

			limitBytes := DefaultMaxBodySize
			if !collection.IsView() {
				for _, f := range collection.Fields {
					if calc, ok := f.(core.MaxBodySizeCalculator); ok {
						limitBytes += calc.CalculateMaxBodySize()
					}
				}
			}

			err = applyBodyLimit(e, limitBytes)
			if err != nil {
				return err
			}

			return e.Next()
		},
	}
}

func applyBodyLimit(e *core.RequestEvent, limitBytes int64) error {
	// no limit
	if limitBytes <= 0 {
		return nil
	}

	// optimistically check the submitted request content length
	if e.Request.ContentLength > limitBytes {
		return ErrRequestEntityTooLarge
	}

	// replace the request body
	e.Request.Body = newLimitedReader(e.Request.Body, limitBytes)

	return nil
}

func newLimitedReader(body io.ReadCloser, limitBytes int64) *limitedReader {
	return &limitedReader{
		ReadCloser: body,
		limit:      limitBytes,
		remaining:  limitBytes,
	}
}

type limitedReader struct {
	io.ReadCloser
	limit     int64
	remaining int64
}

func (r *limitedReader) Read(b []byte) (int, error) {
	if r.remaining <= 0 {
		return 0, ErrRequestEntityTooLarge
	}

	if int64(len(b)) > r.remaining {
		b = b[0:r.remaining]
	}

	n, err := r.ReadCloser.Read(b)

	r.remaining -= int64(n)

	return n, err
}

// explicit casts to ensure that the main struct methods will be invoked
// (extra precautions in case of nested interface wrapping erasure)
// ---

func (r *limitedReader) Reread() {
	rereader, ok := r.ReadCloser.(router.Rereader)
	if ok {
		rereader.Reread()
		r.remaining = r.limit
	}
}

func (r *limitedReader) Close() error {
	closer, ok := r.ReadCloser.(io.Closer)
	if ok {
		return closer.Close()
	}
	return nil
}