feat(handlers): add download endpoint for unsupported file types

- Added the /download endpoint to force file download. - Implemented a handler that validates the path and searches for the corresponding file in the "uploads" directory. - For unsupported files, an HTML page is displayed with a download link. - Improved validation of directories and files, with proper error handling. - Files are now served with the "Content-Disposition: attachment" header to force download.
ImnotEdMateo · Jan 26, 2025 · 44c00fe · 44c00fe
1 parent b3edf4b
commit 44c00fe
Show file tree

Hide file tree

Showing 4 changed files with 171 additions and 46 deletions.
diff --git a/handlers/download.go b/handlers/download.go
@@ -0,0 +1,41 @@
+package handlers
+
+import (
+	"fmt"
+	"net/http"
+	"os"
+	"path/filepath"
+)
+
+func DownloadHandler(w http.ResponseWriter, r *http.Request) {
+	path := r.URL.Path[len("/download/"):]
+	if path == "" {
+		http.NotFound(w, r)
+		return
+	}
+
+	dirPath := filepath.Join("uploads", path)
+
+	dirInfo, err := os.Stat(dirPath)
+	if os.IsNotExist(err) || !dirInfo.IsDir() {
+		http.NotFound(w, r)
+		return
+	}
+
+	files, err := os.ReadDir(dirPath)
+	if err != nil || len(files) == 0 {
+		http.NotFound(w, r)
+		return
+	}
+
+	fileInfo, err := files[0].Info()
+	if err != nil || !fileInfo.Mode().IsRegular() {
+		http.NotFound(w, r)
+		return
+	}
+
+	filePath := filepath.Join(dirPath, fileInfo.Name())
+	w.Header().Set("Content-Disposition", fmt.Sprintf(`attachment; filename="%s"`, fileInfo.Name()))
+	w.Header().Set("Content-Type", "application/octet-stream")
+	http.ServeFile(w, r, filePath)
+}
diff --git a/handlers/files.go b/handlers/files.go
@@ -2,6 +2,7 @@ package handlers
 
 import (
 	"fmt"
+	"html/template"
 	"mime"
 	"net/http"
 	"os"
@@ -12,79 +13,161 @@ import (
 	"github.com/imnotedmateo/usb/config"
 )
 
-func FileOrPageHandler(w http.ResponseWriter, r *http.Request) {
-	// Get the directory path from the URL (removes the leading and trailing slashes)
-	path := r.URL.Path[1:]
-	if len(path) > 0 && path[len(path)-1] == '/' {
-		path = path[:len(path)-1]
-	}
+const downloadTemplateStr = `
+<!DOCTYPE html>
+<html lang="es">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1">
+    <title>{{.FileName}} - USB by edmateo.site</title>
+    <link href="/static/index.css" rel="stylesheet">
+    <link href="/static/themes/{{.Theme}}" rel="stylesheet">
+    <link rel="icon" type="image/x-icon" href="/static/assets/favicon.ico">
+</head>
+<body>
+    <div class="wrapper">
+        <h1><span class="initial">U</span>pload <span class="initial">S</span>erver for <span class="initial">B</span>ullshit</h1>
+        <h2>File: {{.FileName}}</h2>
+        <p>This file cannot be displayed in the browser. You can download it using the link below:</p>
+        <div class="form-shit">
+          <form action="/download/{{.Path}}" method="get">
+            <input type="submit" value="DOWNLOAD">
+          </form>
+        </div>
+        <p>
+            <b>
+                <a href="https://github.com/ImnotEdMateo/ubs">SOURCE CODE</a>
+            </b>
+        </p>
+    </div>
+</body>
+</html>
+`
 
-	// Show the main page if there is no path
+func FileOrPageHandler(w http.ResponseWriter, r *http.Request) {
+	path := getSanitizedPath(r)
 	if path == "" {
 		WebPageHandler(w, r)
 		return
 	}
 
-	// Determine the valid pattern for the path
-	var validPathPattern string
-	if config.RandomPath == "GUID" {
-		validPathPattern = `^[a-f0-9\-]{36}$`
-	} else {
-		numChars, _ := strconv.Atoi(config.RandomPath)
-		validPathPattern = fmt.Sprintf(`^[a-zA-Z0-9]{%d}$`, numChars)
-	}
-
-	// Validate the path against the pattern
-	matched, err := regexp.MatchString(validPathPattern, path)
-	if err != nil || !matched {
+	if err := validatePath(path); err != nil {
 		http.NotFound(w, r)
 		return
 	}
 
-	// Build the full path to the directory
 	dirPath := filepath.Join("uploads", path)
-
-	// Check if the directory exists
-	dirInfo, err := os.Stat(dirPath)
-	if os.IsNotExist(err) || !dirInfo.IsDir() {
+	if err := validateDirectory(dirPath); err != nil {
 		http.NotFound(w, r)
 		return
 	}
 
-	// Look for files inside the directory
-	files, err := os.ReadDir(dirPath)
-	if err != nil || len(files) == 0 {
+	fileName, err := getFirstFileName(dirPath)
+	if err != nil {
 		http.NotFound(w, r)
 		return
 	}
 
-	// Take the first file inside the directory
-	fileName := files[0].Name()
-	filePath := filepath.Join(dirPath, fileName)
-
-	// Check if the file exists and is a regular file
-	fileInfo, err := os.Stat(filePath)
-	if err != nil || fileInfo.IsDir() {
-		http.NotFound(w, r)
+	mimeType := getMimeType(fileName)
+	if !isBrowserSupported(mimeType) {
+		renderDownloadTemplate(w, fileName, path)
 		return
 	}
 
-	// Get the MIME type of the file
+	serveFileInline(w, r, filepath.Join(dirPath, fileName), fileName, mimeType)
+}
+
+func getSanitizedPath(r *http.Request) string {
+	path := r.URL.Path[1:]
+	if len(path) > 0 && path[len(path)-1] == '/' {
+		path = path[:len(path)-1]
+	}
+	return path
+}
+
+func validatePath(path string) error {
+	pattern := getPathPattern()
+	matched, err := regexp.MatchString(pattern, path)
+	if err != nil || !matched {
+		return fmt.Errorf("invalid path")
+	}
+	return nil
+}
+
+func getPathPattern() string {
+	if config.RandomPath == "GUID" {
+		return `^[a-f0-9\-]{36}$`
+	}
+	numChars, _ := strconv.Atoi(config.RandomPath)
+	return fmt.Sprintf(`^[a-zA-Z0-9]{%d}$`, numChars)
+}
+
+func validateDirectory(dirPath string) error {
+	dirInfo, err := os.Stat(dirPath)
+	if err != nil || !dirInfo.IsDir() {
+		return fmt.Errorf("invalid directory")
+	}
+	return nil
+}
+
+func getFirstFileName(dirPath string) (string, error) {
+	files, err := os.ReadDir(dirPath)
+	if err != nil || len(files) == 0 {
+		return "", fmt.Errorf("no files found")
+	}
+	return files[0].Name(), nil
+}
+
+func getMimeType(fileName string) string {
 	ext := filepath.Ext(fileName)
 	mimeType := mime.TypeByExtension(ext)
 	if mimeType == "" {
-		// Detect the MIME type if not found by extension
-		mimeType = "application/octet-stream"
+		return "application/octet-stream"
 	}
+	return mimeType
+}
+
+func renderDownloadTemplate(w http.ResponseWriter, fileName, path string) {
+    tmpl, err := template.New("download").Parse(downloadTemplateStr)
+    if err != nil {
+        http.Error(w, "Error rendering template", http.StatusInternalServerError)
+        return
+    }
+
+    data := struct {
+        FileName string
+        Path     string
+        Theme    string
+    }{
+        FileName: fileName,
+        Path:     path,
+        Theme:    config.Theme,
+    }
 
-	// Set headers
+    w.Header().Set("Content-Type", "text/html; charset=utf-8")
+    tmpl.Execute(w, data)
+}
+
+func serveFileInline(w http.ResponseWriter, r *http.Request, filePath, fileName, mimeType string) {
 	w.Header().Set("Content-Disposition", fmt.Sprintf(`inline; filename="%s"`, fileName))
 	w.Header().Set("Content-Type", mimeType)
+	http.ServeFile(w, r, filePath)
+}
 
-	// If the file is not compatible with the browser, force download
-	if mimeType == "application/octet-stream" {
-		w.Header().Set("Content-Disposition", fmt.Sprintf(`attachment; filename="%s"`, fileName))
+func isBrowserSupported(mimeType string) bool {
+	supportedTypes := map[string]bool{
+		"text/html":              true,
+		"text/plain":             true,
+		"text/css":               true,
+		"application/javascript": true,
+		"image/jpeg":             true,
+		"image/png":              true,
+		"image/gif":              true,
+		"image/svg+xml":          true,
+		"application/pdf":        true,
+		"video/mp4":              true,
+		"audio/mpeg":             true,
+		"audio/wav":              true,
 	}
-
-	http.ServeFile(w, r, filePath)
+	return supportedTypes[mimeType]
 }
diff --git a/main.go b/main.go
@@ -22,6 +22,7 @@ func main() {
 
   http.HandleFunc("/", handlers.FileOrPageHandler)
   http.HandleFunc("/upload", handlers.UploadHandler)
+  http.HandleFunc("/download/", handlers.DownloadHandler)
 
   if err := http.ListenAndServe(":"+port, nil); err != nil {
     log.Fatalf("Error starting server: %v", err)

diff --git a/static/index.css b/static/index.css
@@ -8,7 +8,7 @@ body {
   margin: 0;
 }
 
-h1, p, span {
+h1, h2, p, span {
   margin: 45px 0;
   text-align: center;
 }
@@ -20,7 +20,7 @@ h1, p, span {
     margin: 0 10px;
   }
 
-  h1, p, span {
+  h1, h2, p, span {
     margin: 50px 0;
   }
 }