From 764ec5f38c03e0eac498dac62e068da4a68c8e7b Mon Sep 17 00:00:00 2001
From: nyne <me@nyne.dev>
Date: Mon, 26 May 2025 19:36:37 +0800
Subject: [PATCH] Add temporary token generation for secure file downloads

---
 server/api/file.go     | 42 ++++++++++++++++++++++++++++++++++++++++--
 server/service/file.go |  1 +
 server/utils/jwt.go    | 34 ++++++++++++++++++++++++++++++++++
 3 files changed, 75 insertions(+), 2 deletions(-)

diff --git a/server/api/file.go b/server/api/file.go
index a4d5194..ecf5701 100644
--- a/server/api/file.go
+++ b/server/api/file.go
@@ -1,9 +1,12 @@
 package api
 
 import (
+	"encoding/json"
 	"fmt"
+	"net/url"
 	"nysoure/server/model"
 	"nysoure/server/service"
+	"nysoure/server/utils"
 	"strconv"
 	"strings"
 
@@ -21,6 +24,7 @@ func AddFileRoutes(router fiber.Router) {
 		fileGroup.Get("/:id", getFile)
 		fileGroup.Put("/:id", updateFile)
 		fileGroup.Delete("/:id", deleteFile)
+		fileGroup.Get("/download/local", downloadLocalFile)
 		fileGroup.Get("/download/:id", downloadFile)
 	}
 }
@@ -204,6 +208,40 @@ func downloadFile(c fiber.Ctx) error {
 	if strings.HasPrefix(s, "http") {
 		return c.Redirect().Status(fiber.StatusFound).To(s)
 	}
-	c.Set("Content-Disposition", fmt.Sprintf("attachment; filename=\"%s\"", filename))
-	return c.SendFile(s)
+	data := map[string]string{
+		"path":     s,
+		"filename": filename,
+	}
+	j, _ := json.Marshal(data)
+	token, err := utils.GenerateTemporaryToken(string(j))
+	if err != nil {
+		return model.NewInternalServerError("Failed to generate download token")
+	}
+	return c.Redirect().Status(fiber.StatusFound).To(fmt.Sprintf("%s/api/files/download/local?token=%s", c.BaseURL(), token))
+}
+
+func downloadLocalFile(c fiber.Ctx) error {
+	token := c.Query("token")
+	if token == "" {
+		return model.NewRequestError("Download token is required")
+	}
+
+	data, err := utils.ParseTemporaryToken(token)
+	if err != nil {
+		return model.NewRequestError("Invalid or expired download token")
+	}
+
+	var fileData map[string]string
+	if err := json.Unmarshal([]byte(data), &fileData); err != nil {
+		return model.NewInternalServerError("Failed to parse download data")
+	}
+
+	path := fileData["path"]
+	filename := fileData["filename"]
+
+	c.Set("Content-Disposition", fmt.Sprintf("attachment; filename=\"%s\"", url.PathEscape(filename)))
+
+	return c.SendFile(path, fiber.SendFile{
+		ByteRange: true,
+	})
 }
diff --git a/server/service/file.go b/server/service/file.go
index 287ec22..cf602e6 100644
--- a/server/service/file.go
+++ b/server/service/file.go
@@ -403,6 +403,7 @@ func GetFile(fid string) (*model.FileView, error) {
 	return file.ToView(), nil
 }
 
+// DownloadFile handles the file download request. Return a presigned URL or a direct file path.
 func DownloadFile(ip, fid, cfToken string) (string, string, error) {
 	passed, err := verifyCfToken(cfToken)
 	if err != nil {
diff --git a/server/utils/jwt.go b/server/utils/jwt.go
index 3e01d18..fc335eb 100644
--- a/server/utils/jwt.go
+++ b/server/utils/jwt.go
@@ -60,3 +60,37 @@ func ParseToken(token string) (uint, error) {
 	}
 	return 0, errors.New("invalid token")
 }
+
+// GenerateTemporaryToken creates a JWT token that expires in 15 minutes
+func GenerateTemporaryToken(data string) (string, error) {
+	t := jwt.NewWithClaims(jwt.SigningMethodHS256,
+		jwt.MapClaims{
+			"data": data,
+			"exp":  time.Now().Add(15 * time.Minute).Unix(),
+		})
+	s, err := t.SignedString(key)
+	if err != nil {
+		return "", err
+	}
+	return s, nil
+}
+
+// ParseTemporaryToken parses a JWT token and returns the data if valid
+func ParseTemporaryToken(token string) (string, error) {
+	t, err := jwt.Parse(token, func(t *jwt.Token) (interface{}, error) {
+		return key, nil
+	})
+	if err != nil {
+		return "", err
+	}
+	if claims, ok := t.Claims.(jwt.MapClaims); ok && t.Valid {
+		data := claims["data"].(string)
+		expF := claims["exp"].(float64)
+		exp := time.Unix(int64(expF), 0)
+		if time.Now().After(exp) {
+			return "", errors.New("token expired")
+		}
+		return data, nil
+	}
+	return "", errors.New("invalid token")
+}