From 180f32902b3e37ee923bac1f882822ab525ccf59 Mon Sep 17 00:00:00 2001
From: skidoodle <contact@albert.lol>
Date: Thu, 22 Jan 2026 05:55:24 +0100
Subject: [PATCH] fix: patch flaws and refactor routes

Signed-off-by: skidoodle <contact@albert.lol>
---
 internal/app/server.go | 41 +++++++++++++++++++++++------------------
 internal/app/upload.go | 35 +++++++++++++++++++++++++++--------
 2 files changed, 50 insertions(+), 26 deletions(-)

diff --git a/internal/app/server.go b/internal/app/server.go
index 0b8463f..2bbaa8c 100644
--- a/internal/app/server.go
+++ b/internal/app/server.go
@@ -11,23 +11,7 @@ import (
 func (app *App) Routes() *http.ServeMux {
 	mux := http.NewServeMux()
 
-	fileServer := http.FileServer(http.FS(app.Assets))
-
-	staticHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		if r.URL.Path == "" || strings.HasSuffix(r.URL.Path, "/") {
-			http.NotFound(w, r)
-			return
-		}
-
-		if strings.HasSuffix(r.URL.Path, ".html") {
-			http.NotFound(w, r)
-			return
-		}
-
-		fileServer.ServeHTTP(w, r)
-	})
-
-	mux.Handle("GET /static/", http.StripPrefix("/static/", staticHandler))
+	mux.Handle("GET /static/", http.StripPrefix("/static/", app.handleStatic()))
 	mux.HandleFunc("GET /{$}", app.HandleHome)
 	mux.HandleFunc("POST /{$}", app.HandleUpload)
 	mux.HandleFunc("POST /upload/chunk", app.HandleChunk)
@@ -37,6 +21,18 @@ func (app *App) Routes() *http.ServeMux {
 	return mux
 }
 
+func (app *App) handleStatic() http.Handler {
+	fs := http.FileServer(http.FS(app.Assets))
+
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path == "" || strings.HasSuffix(r.URL.Path, "/") || strings.HasSuffix(r.URL.Path, ".html") {
+			http.NotFound(w, r)
+			return
+		}
+		fs.ServeHTTP(w, r)
+	})
+}
+
 func (app *App) HandleHome(writer http.ResponseWriter, request *http.Request) {
 	err := app.Tmpl.ExecuteTemplate(writer, "layout", map[string]any{
 		"MaxMB":   app.Conf.MaxMB,
@@ -52,7 +48,16 @@ func (app *App) HandleHome(writer http.ResponseWriter, request *http.Request) {
 func (app *App) RespondWithLink(writer http.ResponseWriter, request *http.Request, key []byte, originalName string) {
 	keySlug := base64.RawURLEncoding.EncodeToString(key)
 	ext := filepath.Ext(originalName)
-	link := fmt.Sprintf("%s/%s%s", request.Host, keySlug, ext)
+
+	const unsafeChars = "\"<> \\/:;?@[]^`{}|~"
+	safeExt := strings.Map(func(r rune) rune {
+		if strings.ContainsRune(unsafeChars, r) {
+			return -1
+		}
+		return r
+	}, ext)
+
+	link := fmt.Sprintf("%s/%s%s", request.Host, keySlug, safeExt)
 
 	if request.Header.Get("X-Requested-With") == "XMLHttpRequest" {
 		html := `
diff --git a/internal/app/upload.go b/internal/app/upload.go
index 5460c34..998c523 100644
--- a/internal/app/upload.go
+++ b/internal/app/upload.go
@@ -84,20 +84,16 @@ func (app *App) HandleUpload(writer http.ResponseWriter, request *http.Request)
 		errChan <- err
 	}()
 
-	defer func() {
-		if closeErr := pr.Close(); closeErr != nil {
-			app.Logger.Error("Failed to close pipe reader", "err", closeErr)
-		}
-	}()
-
 	streamer, err := crypto.NewGCMStreamer(ephemeralKey)
 	if err != nil {
+		_ = pr.Close()
 		app.Logger.Error("Failed to create streamer", "err", err)
 		app.SendError(writer, request, http.StatusInternalServerError)
 		return
 	}
 
 	if err := streamer.EncryptStream(tmp, pr); err != nil {
+		_ = pr.Close()
 		app.Logger.Error("Failed to encrypt stream", "err", err)
 		app.SendError(writer, request, http.StatusInternalServerError)
 		return
@@ -128,7 +124,8 @@ func (app *App) HandleUpload(writer http.ResponseWriter, request *http.Request)
 }
 
 func (app *App) HandleChunk(writer http.ResponseWriter, request *http.Request) {
-	request.Body = http.MaxBytesReader(writer, request.Body, MaxRequestOverhead)
+	const MaxChunkBody = UploadChunkSize + (1 << 20)
+	request.Body = http.MaxBytesReader(writer, request.Body, MaxChunkBody)
 
 	uid := request.FormValue("upload_id")
 	idx, err := strconv.Atoi(request.FormValue("index"))
@@ -146,7 +143,7 @@ func (app *App) HandleChunk(writer http.ResponseWriter, request *http.Request) {
 
 	file, _, err := request.FormFile("chunk")
 	if err != nil {
-		if err.Error() == "http: request body too large" {
+		if strings.Contains(err.Error(), "request body too large") {
 			app.SendError(writer, request, http.StatusRequestEntityTooLarge)
 			return
 		}
@@ -186,6 +183,28 @@ func (app *App) HandleFinish(writer http.ResponseWriter, request *http.Request)
 		}
 	}()
 
+	var totalSize int64
+	for i := range total {
+		info, err := os.Stat(filepath.Join(app.Conf.StorageDir, TempDirName, uid, strconv.Itoa(i)))
+		if err != nil {
+			app.Logger.Error("Missing chunk", "index", i, "err", err)
+			app.SendError(writer, request, http.StatusBadRequest)
+			return
+		}
+		chunkContentSize := info.Size() - crypto.KeySize
+		if chunkContentSize < 0 {
+			app.SendError(writer, request, http.StatusBadRequest)
+			return
+		}
+		totalSize += chunkContentSize
+	}
+
+	if totalSize > (app.Conf.MaxMB * MegaByte) {
+		app.Logger.Warn("Upload exceeded quota", "uid", uid, "size", totalSize)
+		app.SendError(writer, request, http.StatusRequestEntityTooLarge)
+		return
+	}
+
 	hasher := sha256.New()
 	for i := range total {
 		rc, err := app.openChunkDecryptor(uid, i)