From 1ba891eef05476f38c4a8d13a60c06944816ed78 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?O=CC=81scar=20M=2E=20Lage?= <info@oscarmlage.com>
Date: Thu, 5 Dec 2024 18:46:36 +0100
Subject: [PATCH] Imp: Limit somehow the upload file in size and mimetype

---
 internal/web/controller.go | 12 +++++++++++-
 internal/web/utils.go      | 27 +++++++++++++++++++++++++++
 2 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/internal/web/controller.go b/internal/web/controller.go
index 8cd0122..064ba5e 100644
--- a/internal/web/controller.go
+++ b/internal/web/controller.go
@@ -25,13 +25,23 @@ func handleUpload(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	file, _, err := r.FormFile("file")
+	file, fileHeader, err := r.FormFile("file")
 	if err != nil {
 		http.Error(w, "Error reading file", http.StatusBadRequest)
 		return
 	}
 	defer file.Close()
 
+	isValid, err := isTextFileAndSizeOk(file, fileHeader.Size)
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusBadRequest)
+		return
+	}
+	if !isValid {
+		http.Error(w, "Invalid file type or size", http.StatusBadRequest)
+		return
+	}
+
 	commandCounts, categories, pipeRedirectionCounts, commonPatterns := ProcessHistory(file)
 
 	limit := os.Getenv("TOP_N_COMMANDS")
diff --git a/internal/web/utils.go b/internal/web/utils.go
index 0dccc29..1ba8463 100644
--- a/internal/web/utils.go
+++ b/internal/web/utils.go
@@ -5,7 +5,9 @@ import (
 	"encoding/json"
 	"fmt"
 	"html/template"
+	"io"
 	"log"
+	"mime/multipart"
 	"net/http"
 	"os"
 	"path/filepath"
@@ -166,3 +168,28 @@ func loadTemplatesFromDir(dir string) ([]string, error) {
 
 	return templates, err
 }
+
+func isTextFileAndSizeOk(file multipart.File, size int64) (bool, error) {
+	if size > 1*1024*1024 {
+		return false, fmt.Errorf("File size exceeds 1MB")
+	}
+
+	buffer := make([]byte, 512)
+	_, err := file.Read(buffer)
+	if err != nil && err != io.EOF {
+		return false, fmt.Errorf("Error reading file: %v", err)
+	}
+
+	mimeType := http.DetectContentType(buffer)
+
+	if !strings.HasPrefix(mimeType, "text/") {
+		return false, fmt.Errorf("File is not a text file (mimetype: %s)", mimeType)
+	}
+
+	_, err = file.Seek(0, 0)
+	if err != nil {
+		return false, fmt.Errorf("Error seeking file: %v", err)
+	}
+
+	return true, nil
+}