Add configurable MAX_UPLOAD_SIZE limit with clear error responses

main.go

@@ -15,6 +15,7 @@ import (
 	"os"
 	"os/signal"
 	"path/filepath"
+	"strconv"
 	"strings"
 	"syscall"
 	"time"
@@ -25,7 +26,11 @@ var uploadDir string
 const idLength = 8
 const idCharset = "abcdefghijklmnopqrstuvwxyz0123456789"
 
+// defaultMaxUploadSize is used when MAX_UPLOAD_SIZE is not set.
+const defaultMaxUploadSize int64 = 32 << 20 // 32MB
+
 var indexHTML *template.Template
+var maxUploadSize int64 = defaultMaxUploadSize
 
 func main() {
 	uploadDir = os.Getenv("UPLOAD_DIR")
@@ -44,6 +49,14 @@ func main() {
 		fileExpires = expires
 	}
 
+	if maxUploadSizeEnv := os.Getenv("MAX_UPLOAD_SIZE"); maxUploadSizeEnv != "" {
+		size, err := strconv.ParseInt(maxUploadSizeEnv, 10, 64)
+		if err != nil || size <= 0 {
+			log.Fatalf("MAX_UPLOAD_SIZE %s is invalid: %v", maxUploadSizeEnv, err)
+		}
+		maxUploadSize = size
+	}
+
 	port := os.Getenv("PORT")
 	if port == "" {
 		port = "8080"
@@ -57,6 +70,7 @@ func main() {
 
 	log.Printf("Starting server on port %s, upload directory: %s", port, uploadDir)
 	log.Printf("Files expires after: %v (DRYRUN_CLEAN=%v)", fileExpires, dryRun)
+	log.Printf("Max upload size: %d bytes", maxUploadSize)
 	if fileExpires != 0 {
 		go cleanupProcedure(fileExpires, dryRun)
 	} else {
@@ -169,8 +183,16 @@ func handleUpload(w http.ResponseWriter, r *http.Request, start time.Time) {
 	}
 	defer file.Close()
 
-	written, err := io.Copy(file, r.Body)
+	limitedBody := http.MaxBytesReader(w, r.Body, maxUploadSize)
+	written, err := io.Copy(file, limitedBody)
 	if err != nil {
+		file.Close()
+		os.Remove(targetPath)
+		if isMaxBytesError(err) {
+			http.Error(w, fmt.Sprintf("File exceeds maximum upload size of %d bytes", maxUploadSize), http.StatusRequestEntityTooLarge)
+			logRequest(r, 413, time.Since(start))
+			return
+		}
 		http.Error(w, "Failed to write file", http.StatusInternalServerError)
 		logRequest(r, 500, time.Since(start))
 		return
@@ -187,16 +209,34 @@ func handleUpload(w http.ResponseWriter, r *http.Request, start time.Time) {
 	logRequest(r, 200, time.Since(start))
 }
 
+func writeMultipartError(w http.ResponseWriter, status int, message string) {
+	w.Header().Set("Content-Type", "application/json")
+	w.WriteHeader(status)
+	json.NewEncoder(w).Encode(map[string]string{"error": message})
+}
+
 func handleMultipartUpload(w http.ResponseWriter, r *http.Request, start time.Time) {
-	if err := r.ParseMultipartForm(32 << 20); err != nil {
-		http.Error(w, "Failed to parse multipart form", http.StatusBadRequest)
+	r.Body = http.MaxBytesReader(w, r.Body, maxUploadSize)
+
+	memoryLimit := maxUploadSize
+	if memoryLimit > defaultMaxUploadSize {
+		memoryLimit = defaultMaxUploadSize
+	}
+
+	if err := r.ParseMultipartForm(memoryLimit); err != nil {
+		if isMaxBytesError(err) {
+			writeMultipartError(w, http.StatusRequestEntityTooLarge, fmt.Sprintf("Upload exceeds maximum size of %d bytes", maxUploadSize))
+			logRequest(r, 413, time.Since(start))
+			return
+		}
+		writeMultipartError(w, http.StatusBadRequest, "Failed to parse multipart form")
 		logRequest(r, 400, time.Since(start))
 		return
 	}
 
 	files := r.MultipartForm.File["files"]
 	if len(files) == 0 {
-		http.Error(w, "No files provided", http.StatusBadRequest)
+		writeMultipartError(w, http.StatusBadRequest, "No files provided")
 		logRequest(r, 400, time.Since(start))
 		return
 	}
@@ -231,6 +271,7 @@ func handleMultipartUpload(w http.ResponseWriter, r *http.Request, start time.Ti
 		src.Close()
 
 		if err != nil {
+			os.Remove(targetPath)
 			continue
 		}
 
@@ -366,6 +407,13 @@ func getProtocol(r *http.Request) string {
 	return "http"
 }
 
+// isMaxBytesError reports whether err was caused by exceeding the limit
+// set on an http.MaxBytesReader.
+func isMaxBytesError(err error) bool {
+	var maxBytesErr *http.MaxBytesError
+	return errors.As(err, &maxBytesErr)
+}
+
 func detectContentType(filename string) string {
 	ext := filepath.Ext(filename)
 	mimeType := mime.TypeByExtension(ext)