new: option to set max upload size

ybizeul · Aug 23, 2023 · a45bdcf · a45bdcf
1 parent 62a525b
commit a45bdcf
Show file tree

Hide file tree

Showing 4 changed files with 48 additions and 6 deletions.
diff --git a/cmd/ybfeed/main.go b/cmd/ybfeed/main.go
@@ -13,6 +13,7 @@ import (
 var HTTP_PORT int
 var DEBUG bool
 var dataDir string
+var maxBodySize int
 
 var logLevel slog.LevelVar
 
@@ -46,14 +47,22 @@ func main() {
 				Usage:       "Data directory path",
 				Destination: &dataDir,
 			},
+			&cli.IntFlag{
+				Name:        "max-upload-size",
+				Aliases:     []string{"m"},
+				Value:       5,
+				EnvVars:     []string{"YBF_MAX_UPLOAD_SIZE"},
+				Usage:       "Max upload size in MB",
+				Destination: &maxBodySize,
+			},
 		},
 		Action: func(cCtx *cli.Context) error {
 			logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: &logLevel}))
 			slog.SetDefault(logger)
 			if DEBUG {
 				logLevel.Set(slog.LevelDebug)
 			}
-			slog.Info("Debugging", slog.Bool("status", DEBUG))
+			slog.Debug("Running in DEBUG mode")
 			run()
 			return nil
 		},
@@ -71,12 +80,16 @@ func run() {
 	// Start HTTP Server
 	r := http.NewServeMux()
 
-	api := handlers.ApiHandler{BasePath: dataDir, Version: version}
+	api := handlers.ApiHandler{
+		BasePath:    dataDir,
+		Version:     version,
+		MaxBodySize: maxBodySize * 1024 * 1024,
+	}
 
 	r.HandleFunc("/api/", api.ApiHandleFunc)
 	r.HandleFunc("/", handlers.RootHandlerFunc)
 
-	slog.Info("ybFeed starting", slog.String("version", version), slog.String("data_dir", dataDir), slog.Int("port", HTTP_PORT))
+	slog.Info("ybFeed starting", slog.String("version", version), slog.String("data_dir", dataDir), slog.Int("port", HTTP_PORT), slog.Int("max-upload-size", maxBodySize))
 	err := http.ListenAndServe(fmt.Sprintf(":%d", HTTP_PORT), r)
 	if err != nil {
 		slog.Error("Unable to start HTTP server", slog.String("error", err.Error()))

diff --git a/internal/feed/feed.go b/internal/feed/feed.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"io"
 	"io/fs"
+	"net/http"
 	"os"
 	"path"
 	"path/filepath"
@@ -257,6 +258,13 @@ func (feed *Feed) AddItem(contentType string, f io.ReadCloser) error {
 
 	content, err := io.ReadAll(f)
 	if err != nil {
+		_, ok := err.(*http.MaxBytesError)
+		if ok {
+			return &FeedError{
+				Code:    413,
+				Message: "Max body size exceeded",
+			}
+		}
 		return &FeedError{
 			Code:    500,
 			Message: "Unable to read stream",

diff --git a/internal/handlers/handlers.go b/internal/handlers/handlers.go
@@ -57,8 +57,9 @@ func RootHandlerFunc(w http.ResponseWriter, r *http.Request) {
 
 // Handle requests to /api
 type ApiHandler struct {
-	BasePath string
-	Version  string
+	BasePath    string
+	Version     string
+	MaxBodySize int
 }
 
 func NewApiHandler(basePath string) *ApiHandler {
@@ -213,7 +214,8 @@ func (api *ApiHandler) feedPostHandlerFunc(w http.ResponseWriter, r *http.Reques
 
 	contentType := r.Header.Get("Content-type")
 
-	err = f.AddItem(contentType, r.Body)
+	err = f.AddItem(contentType, http.MaxBytesReader(w, r.Body, int64(api.MaxBodySize)))
+	//err = f.AddItem(contentType, r.Body)
 
 	if err != nil {
 		yberr := err.(*feed.FeedError)

diff --git a/internal/handlers/handlers_test.go b/internal/handlers/handlers_test.go
@@ -357,6 +357,25 @@ func TestAddAndRemoveContent(t *testing.T) {
 	}
 }
 
+func TestAddContentTooBig(t *testing.T) {
+	b := bytes.NewBuffer(make([]byte, 6*1024*1024))
+
+	res := APITestRequest{
+		method:         http.MethodPost,
+		body:           b,
+		cookieAuthType: AuthTypeAuth,
+		contentType:    "image/png",
+	}.performRequest()
+
+	if res.StatusCode != 413 {
+		b, err := io.ReadAll(res.Body)
+		if err != nil {
+			t.Errorf("Expect code 413 but got %d (%s)", res.StatusCode, err.Error())
+		}
+		t.Errorf("Expect code 413 but got %d (%s)", res.StatusCode, string(b))
+	}
+}
+
 func TestAddContentWrongContentType(t *testing.T) {
 	res := APITestRequest{
 		method:         http.MethodPost,