feat(api): add --api-body-limit flag that set the body limit for multipart/form-data requests

pkg/gotenberg/flags.go

@@ -168,33 +168,37 @@ func (f *ParsedFlags) MustDeprecatedDuration(deprecated string, newName string)
 	return f.MustDuration(newName)
 }
 
-// MustHumanReadableBytesString returns the human-readable bytes string of a
-// flag given by name.
+// MustHumanReadableBytes returns the human-readable bytes string of a flag
+// given by name.
 // It panics if an error occurs.
-func (f *ParsedFlags) MustHumanReadableBytesString(name string) string {
+func (f *ParsedFlags) MustHumanReadableBytes(name string) int64 {
 	val, err := f.GetString(name)
 	if err != nil {
 		panic(err)
 	}
 
-	_, err = bytes.Parse(val)
+	if val == "" {
+		return 0
+	}
+
+	b, err := bytes.Parse(val)
 	if err != nil {
 		panic(err)
 	}
 
-	return val
+	return b
 }
 
-// MustDeprecatedHumanReadableBytesString returns the human-readable bytes
-// string of a deprecated flag if it was explicitly set or the human-readable
-// bytes string of the new flag.
+// MustDeprecatedHumanReadableBytes returns the human-readable bytes of a
+// deprecated flag if it was explicitly set or the human-readable bytes string
+// of the new flag.
 // It panics if an error occurs.
-func (f *ParsedFlags) MustDeprecatedHumanReadableBytesString(deprecated string, newName string) string {
+func (f *ParsedFlags) MustDeprecatedHumanReadableBytes(deprecated string, newName string) int64 {
 	if f.Changed(deprecated) {
-		return f.MustHumanReadableBytesString(deprecated)
+		return f.MustHumanReadableBytes(deprecated)
 	}
 
-	return f.MustHumanReadableBytesString(newName)
+	return f.MustHumanReadableBytes(newName)
 }
 
 // MustRegexp returns the regular expression of a flag given by name.

pkg/gotenberg/flags_test.go

@@ -644,10 +644,11 @@ func TestParsedFlags_MustDeprecatedDuration(t *testing.T) {
 	}
 }
 
-func TestParsedFlags_MustHumanReadableBytesString(t *testing.T) {
+func TestParsedFlags_MustHumanReadableBytes(t *testing.T) {
 	fs := flag.NewFlagSet("tests", flag.ContinueOnError)
 	fs.String("foo", "1MB", "")
 	fs.String("bar", "1MB", "")
+	fs.String("qux", "", "")
 
 	err := fs.Parse([]string{"--foo=1GB", "--bar=foo"})
 	if err != nil {
@@ -671,6 +672,11 @@ func TestParsedFlags_MustHumanReadableBytesString(t *testing.T) {
 			name:        "bar",
 			expectPanic: true,
 		},
+		{
+			scenario:    "success: empty value",
+			name:        "qux",
+			expectPanic: false,
+		},
 	} {
 		t.Run(tc.scenario, func(t *testing.T) {
 			if tc.expectPanic {
@@ -689,31 +695,31 @@ func TestParsedFlags_MustHumanReadableBytesString(t *testing.T) {
 				}()
 			}
 
-			parsedFlags.MustHumanReadableBytesString(tc.name)
+			parsedFlags.MustHumanReadableBytes(tc.name)
 		})
 	}
 }
 
-func TestParsedFlags_MustDeprecatedHumanReadableBytesString(t *testing.T) {
+func TestParsedFlags_MustDeprecatedHumanReadableBytes(t *testing.T) {
 	for _, tc := range []struct {
 		scenario    string
 		rawFlags    []string
-		expectValue string
+		expectValue int64
 	}{
 		{
 			scenario:    "deprecated flag value",
 			rawFlags:    []string{"--foo=1MB"},
-			expectValue: "1MB",
+			expectValue: 1000000,
 		},
 		{
 			scenario:    "non-deprecated flag value",
 			rawFlags:    []string{"--bar=2MB"},
-			expectValue: "2MB",
+			expectValue: 2000000,
 		},
 		{
 			scenario:    "deprecated flag value > non-deprecated flag value",
 			rawFlags:    []string{"--foo=1MB", "--bar=2MB"},
-			expectValue: "1MB",
+			expectValue: 1000000,
 		},
 	} {
 		t.Run(tc.scenario, func(t *testing.T) {
@@ -728,9 +734,9 @@ func TestParsedFlags_MustDeprecatedHumanReadableBytesString(t *testing.T) {
 				t.Fatalf("expected no error but got: %v", err)
 			}
 
-			actual := parsedFlags.MustDeprecatedHumanReadableBytesString("foo", "bar")
+			actual := parsedFlags.MustDeprecatedHumanReadableBytes("foo", "bar")
 			if actual != tc.expectValue {
-				t.Errorf("expected '%s' but got '%s'", tc.expectValue, actual)
+				t.Errorf("expected %d but got %d", tc.expectValue, actual)
 			}
 		})
 	}

pkg/modules/api/api.go

@@ -32,6 +32,7 @@ type Api struct {
 	tlsCertFile               string
 	tlsKeyFile                string
 	startTimeout              time.Duration
+	bodyLimit                 int64
 	timeout                   time.Duration
 	rootPath                  string
 	traceHeader               string
@@ -174,6 +175,7 @@ func (a *Api) Descriptor() gotenberg.ModuleDescriptor {
 			fs.String("api-tls-key-file", "", "Path to the TLS/SSL key file - for HTTPS support")
 			fs.Duration("api-start-timeout", time.Duration(30)*time.Second, "Set the time limit for the API to start")
 			fs.Duration("api-timeout", time.Duration(30)*time.Second, "Set the time limit for requests")
+			fs.String("api-body-limit", "", "Set the body limit for multipart/form-data requests")
 			fs.String("api-root-path", "/", "Set the root path of the API - for service discovery via URL paths")
 			fs.String("api-trace-header", "Gotenberg-Trace", "Set the header name to use for identifying requests")
 			fs.Bool("api-enable-basic-auth", false, "Enable basic authentication - will look for the GOTENBERG_API_BASIC_AUTH_USERNAME and GOTENBERG_API_BASIC_AUTH_PASSWORD environment variables")
@@ -196,6 +198,7 @@ func (a *Api) Provision(ctx *gotenberg.Context) error {
 	a.tlsKeyFile = flags.MustString("api-tls-key-file")
 	a.startTimeout = flags.MustDuration("api-start-timeout")
 	a.timeout = flags.MustDuration("api-timeout")
+	a.bodyLimit = flags.MustHumanReadableBytes("api-body-limit")
 	a.rootPath = flags.MustString("api-root-path")
 	a.traceHeader = flags.MustString("api-trace-header")
 	a.downloadFromCfg = downloadFromConfig{
@@ -455,7 +458,7 @@ func (a *Api) Start() error {
 		}
 
 		if route.IsMultipart {
-			middlewares = append(middlewares, contextMiddleware(a.fs, a.timeout, a.downloadFromCfg))
+			middlewares = append(middlewares, contextMiddleware(a.fs, a.timeout, a.bodyLimit, a.downloadFromCfg))
 
 			for _, externalMultipartMiddleware := range externalMultipartMiddlewares {
 				middlewares = append(middlewares, externalMultipartMiddleware.Handler)

pkg/modules/api/context.go

@@ -13,6 +13,7 @@ import (
 	"os"
 	"path/filepath"
 	"strings"
+	"sync/atomic"
 	"time"
 
 	"github.com/google/uuid"
@@ -50,6 +51,28 @@ type Context struct {
 	context.Context
 }
 
+type trackingReader struct {
+	R            io.Reader
+	AddReadBytes func(n int64) error
+}
+
+func (t *trackingReader) Read(p []byte) (int, error) {
+	n, err := t.R.Read(p)
+	if n > 0 {
+		errAddRead := t.AddReadBytes(int64(n))
+		if errAddRead != nil {
+			return n, fmt.Errorf("add read bytes: %w", errAddRead)
+		}
+	}
+	if err != nil {
+		// It's a common practice in Go to return io.EOF unwrapped to signal
+		// the end of a data stream. Wrapping it can lead to unexpected
+		// behavior in standard library functions.
+		return n, err
+	}
+	return n, nil
+}
+
 type downloadFrom struct {
 	// Url is the URL to download a file from.
 	Url string `json:"url"`
@@ -65,9 +88,25 @@ func (o *osPathRename) Rename(oldpath, newpath string) error {
 }
 
 // newContext returns a [Context] by parsing a "multipart/form-data" request.
-func newContext(echoCtx echo.Context, logger *zap.Logger, fs *gotenberg.FileSystem, timeout time.Duration, downloadFromCfg downloadFromConfig, traceHeader, trace string) (*Context, context.CancelFunc, error) {
+func newContext(echoCtx echo.Context, logger *zap.Logger, fs *gotenberg.FileSystem, timeout time.Duration, bodyLimit int64, downloadFromCfg downloadFromConfig, traceHeader, trace string) (*Context, context.CancelFunc, error) {
 	processCtx, processCancel := context.WithTimeout(context.Background(), timeout)
 
+	// We want to make sure the multipart/form-data does not exceed a given
+	// limit. We consider: form fields (keys, values, files) and files
+	// downloaded remotely ("download from" feature).
+	var totalBytesRead atomic.Int64
+
+	addReadBytes := func(n int64) error {
+		newTotal := totalBytesRead.Add(n)
+		if bodyLimit != 0 && newTotal > bodyLimit {
+			return WrapError(
+				fmt.Errorf("body limit reached (> %d)", bodyLimit),
+				NewSentinelHttpError(http.StatusRequestEntityTooLarge, http.StatusText(http.StatusRequestEntityTooLarge)),
+			)
+		}
+		return nil
+	}
+
 	ctx := &Context{
 		outputPaths: make([]string, 0),
 		cancelled:   false,
@@ -129,6 +168,19 @@ func newContext(echoCtx echo.Context, logger *zap.Logger, fs *gotenberg.FileSyst
 		return nil, cancel, fmt.Errorf("get multipart form: %w", err)
 	}
 
+	// This will ensure we do not exceed the body limit.
+	var formValuesSize int64
+	for key, valArray := range form.Value {
+		formValuesSize += int64(len(key))
+		for _, val := range valArray {
+			formValuesSize += int64(len(val))
+		}
+	}
+	err = addReadBytes(formValuesSize)
+	if err != nil {
+		return nil, cancel, fmt.Errorf("add read bytes: %w", err)
+	}
+
 	dirPath, err := fs.MkdirAll()
 	if err != nil {
 		return nil, cancel, fmt.Errorf("create working directory: %w", err)
@@ -262,9 +314,12 @@ func newContext(echoCtx echo.Context, logger *zap.Logger, fs *gotenberg.FileSyst
 					}
 				}()
 
-				_, err = io.Copy(out, resp.Body)
+				// This will ensure we do not exceed the body limit.
+				reader := &trackingReader{R: resp.Body, AddReadBytes: addReadBytes}
+
+				_, err = io.Copy(out, reader)
 				if err != nil {
-					return fmt.Errorf("copy downloaded file from '%s' to local file: %v", dl.Url, err)
+					return fmt.Errorf("copy downloaded file from '%s' to local file: %w", dl.Url, err)
 				}
 
 				ctx.files[filename] = path
@@ -292,6 +347,9 @@ func newContext(echoCtx echo.Context, logger *zap.Logger, fs *gotenberg.FileSyst
 			}
 		}()
 
+		// This will ensure we do not exceed the body limit.
+		reader := &trackingReader{R: in, AddReadBytes: addReadBytes}
+
 		// Avoid directory traversal and make sure filename characters are
 		// normalized.
 		// See: https://github.com/gotenberg/gotenberg/issues/662.
@@ -309,7 +367,7 @@ func newContext(echoCtx echo.Context, logger *zap.Logger, fs *gotenberg.FileSyst
 			}
 		}()
 
-		_, err = io.Copy(out, in)
+		_, err = io.Copy(out, reader)
 		if err != nil {
 			return fmt.Errorf("copy multipart file to local file: %w", err)
 		}
@@ -331,6 +389,7 @@ func newContext(echoCtx echo.Context, logger *zap.Logger, fs *gotenberg.FileSyst
 
 	ctx.Log().Debug(fmt.Sprintf("form fields: %+v", ctx.values))
 	ctx.Log().Debug(fmt.Sprintf("form files: %+v", ctx.files))
+	ctx.Log().Debug(fmt.Sprintf("total bytes: %d", totalBytesRead.Load()))
 
 	return ctx, cancel, err
 }

pkg/modules/api/context_test.go

@@ -95,6 +95,7 @@ func TestNewContext(t *testing.T) {
 	for _, tc := range []struct {
 		scenario         string
 		request          *http.Request
+		bodyLimit        int64
 		downloadFromCfg  downloadFromConfig
 		downloadFromSrv  *echo.Echo
 		expectContext    *Context
@@ -143,6 +144,95 @@ func TestNewContext(t *testing.T) {
 			expectHttpError:  true,
 			expectHttpStatus: http.StatusBadRequest,
 		},
+		{
+			scenario: "request entity too large: form values",
+			request: func() *http.Request {
+				body := &bytes.Buffer{}
+				writer := multipart.NewWriter(body)
+				defer func() {
+					err := writer.Close()
+					if err != nil {
+						t.Fatalf("expected no error but got: %v", err)
+					}
+				}()
+				err := writer.WriteField("key", "value")
+				if err != nil {
+					t.Fatalf("expected no error but got: %v", err)
+				}
+				req := httptest.NewRequest(http.MethodPost, "/", body)
+				req.Header.Set(echo.HeaderContentType, writer.FormDataContentType())
+				return req
+			}(),
+			bodyLimit:        1,
+			downloadFromCfg:  defaultDownloadFromCfg,
+			expectError:      true,
+			expectHttpError:  true,
+			expectHttpStatus: http.StatusRequestEntityTooLarge,
+		},
+		{
+			scenario: "request entity too large: downloadFrom",
+			request: func() *http.Request {
+				body := &bytes.Buffer{}
+				writer := multipart.NewWriter(body)
+				defer func() {
+					err := writer.Close()
+					if err != nil {
+						t.Fatalf("expected no error but got: %v", err)
+					}
+				}()
+				err := writer.WriteField("downloadFrom", `[{"url":"http://localhost:80/"}]`)
+				if err != nil {
+					t.Fatalf("expected no error but got: %v", err)
+				}
+				req := httptest.NewRequest(http.MethodPost, "/", body)
+				req.Header.Set(echo.HeaderContentType, writer.FormDataContentType())
+				return req
+			}(),
+			bodyLimit: 45, // form values = 44 bytes.
+			downloadFromSrv: func() *echo.Echo {
+				srv := echo.New()
+				srv.HideBanner = true
+				srv.GET("/", func(c echo.Context) error {
+					c.Response().Header().Set(echo.HeaderContentDisposition, `attachment; filename="bar.txt"`)
+					c.Response().Header().Set(echo.HeaderContentType, "text/plain")
+					return c.String(http.StatusOK, http.StatusText(http.StatusOK))
+				})
+				return srv
+			}(),
+			downloadFromCfg:  defaultDownloadFromCfg,
+			expectError:      true,
+			expectHttpError:  true,
+			expectHttpStatus: http.StatusRequestEntityTooLarge,
+		},
+		{
+			scenario: "request entity too large: form files",
+			request: func() *http.Request {
+				body := &bytes.Buffer{}
+				writer := multipart.NewWriter(body)
+				defer func() {
+					err := writer.Close()
+					if err != nil {
+						t.Fatalf("expected no error but got: %v", err)
+					}
+				}()
+				part, err := writer.CreateFormFile("foo.txt", "foo.txt")
+				if err != nil {
+					t.Fatalf("expected no error but got: %v", err)
+				}
+				_, err = part.Write([]byte("foo"))
+				if err != nil {
+					t.Fatalf("expected no error but got: %v", err)
+				}
+				req := httptest.NewRequest(http.MethodPost, "/", body)
+				req.Header.Set(echo.HeaderContentType, writer.FormDataContentType())
+				return req
+			}(),
+			bodyLimit:        1,
+			downloadFromCfg:  defaultDownloadFromCfg,
+			expectError:      true,
+			expectHttpError:  true,
+			expectHttpStatus: http.StatusRequestEntityTooLarge,
+		},
 		{
 			scenario: "invalid downloadFrom form field: cannot unmarshal",
 			request: func() *http.Request {
@@ -458,7 +548,7 @@ func TestNewContext(t *testing.T) {
 			}
 
 			handler := func(c echo.Context) error {
-				ctx, cancel, err := newContext(c, zap.NewNop(), gotenberg.NewFileSystem(), time.Duration(10)*time.Second, tc.downloadFromCfg, "Gotenberg-Trace", "123")
+				ctx, cancel, err := newContext(c, zap.NewNop(), gotenberg.NewFileSystem(), time.Duration(10)*time.Second, tc.bodyLimit, tc.downloadFromCfg, "Gotenberg-Trace", "123")
 				defer cancel()
 				// Context already cancelled.
 				defer cancel()

pkg/modules/api/middlewares.go

@@ -236,7 +236,7 @@ func basicAuthMiddleware(username, password string) echo.MiddlewareFunc {
 //
 //	ctx := c.Get("context").(*api.Context)
 //	cancel := c.Get("cancel").(context.CancelFunc)
-func contextMiddleware(fs *gotenberg.FileSystem, timeout time.Duration, downloadFromCfg downloadFromConfig) echo.MiddlewareFunc {
+func contextMiddleware(fs *gotenberg.FileSystem, timeout time.Duration, bodyLimit int64, downloadFromCfg downloadFromConfig) echo.MiddlewareFunc {
 	return func(next echo.HandlerFunc) echo.HandlerFunc {
 		return func(c echo.Context) error {
 			logger := c.Get("logger").(*zap.Logger)
@@ -245,7 +245,7 @@ func contextMiddleware(fs *gotenberg.FileSystem, timeout time.Duration, download
 
 			// We create a context with a timeout so that underlying processes are
 			// able to stop early and handle correctly a timeout scenario.
-			ctx, cancel, err := newContext(c, logger, fs, timeout, downloadFromCfg, traceHeader, trace)
+			ctx, cancel, err := newContext(c, logger, fs, timeout, bodyLimit, downloadFromCfg, traceHeader, trace)
 			if err != nil {
 				cancel()
 

pkg/modules/api/middlewares_test.go

@@ -462,7 +462,7 @@ func TestContextMiddleware(t *testing.T) {
 		c.Set("trace", "foo")
 		c.Set("startTime", time.Now())
 
-		err := contextMiddleware(gotenberg.NewFileSystem(), time.Duration(10)*time.Second, downloadFromConfig{})(tc.next)(c)
+		err := contextMiddleware(gotenberg.NewFileSystem(), time.Duration(10)*time.Second, 0, downloadFromConfig{})(tc.next)(c)
 
 		if tc.expectErr && err == nil {
 			t.Errorf("test %d: expected error but got: %v", i, err)