文件管理权限控制

2026-10-06 15:46:37 +08:00 · 2026-06-11 13:58:45 +08:00
parent 312bc7d4d5
commit 5e0de01c4b
18 changed files with 993 additions and 337 deletions
@@ -0,0 +1,211 @@
+// Copyright 2025 linux.do
+// Copyright 2026 Arctel.net
+// SPDX-License-Identifier: Apache-2.0
+
+package upload
+
+import (
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"testing"
+
+	"github.com/Rain-kl/Wavelet/internal/common"
+	"github.com/Rain-kl/Wavelet/internal/model"
+	"github.com/Rain-kl/Wavelet/internal/testhelper"
+	"github.com/gin-contrib/sessions"
+	"github.com/gin-contrib/sessions/cookie"
+	"github.com/gin-gonic/gin"
+)
+
+func TestServeFileByIDAccessControl(t *testing.T) {
+	dbConn, _, cleanup := testhelper.SetupTestEnvironment(t)
+	defer cleanup()
+
+	// Ensure uploads dir is cleaned up
+	defer func() { _ = os.RemoveAll("uploads") }()
+
+	// Create a user in DB
+	user := model.User{
+		ID:       12345,
+		Username: "file_test_user",
+		IsActive: true,
+	}
+	if err := dbConn.Create(&user).Error; err != nil {
+		t.Fatalf("failed to create user: %v", err)
+	}
+
+	// Create an access token for this user
+	tokenStr := "test-secret-token-123"
+	tokenHash := model.HashToken(tokenStr)
+	tokenRecord := model.AccessToken{
+		UserID:    user.ID,
+		Name:      "test_token",
+		TokenHash: tokenHash,
+	}
+	if err := dbConn.Create(&tokenRecord).Error; err != nil {
+		t.Fatalf("failed to create token: %v", err)
+	}
+
+	// Create two files: one in whitelist (avatar), one not in whitelist (attachment)
+	avatarFile := model.Upload{
+		ID:            8001,
+		UserID:        user.ID,
+		FileName:      "avatar.png",
+		FilePath:      "uploads/avatar.png",
+		FileSize:      5,
+		MimeType:      "image/png",
+		Extension:     "png",
+		StorageDriver: "local",
+		Type:          "avatar",
+		Status:        model.UploadStatusUsed,
+	}
+	attachmentFile := model.Upload{
+		ID:            8002,
+		UserID:        user.ID,
+		FileName:      "doc.pdf",
+		FilePath:      "uploads/doc.pdf",
+		FileSize:      5,
+		MimeType:      "application/pdf",
+		Extension:     "pdf",
+		StorageDriver: "local",
+		Type:          "attachment",
+		Status:        model.UploadStatusUsed,
+	}
+
+	_ = os.MkdirAll("uploads", 0755)
+	_ = os.WriteFile(avatarFile.FilePath, []byte("image"), 0644)
+	_ = os.WriteFile(attachmentFile.FilePath, []byte("bytes"), 0644)
+
+	dbConn.Create(&avatarFile)
+	dbConn.Create(&attachmentFile)
+
+	// Set up router
+	gin.SetMode(gin.TestMode)
+	r := gin.New()
+	store := cookie.NewStore([]byte("secret"))
+	r.Use(sessions.Sessions("test_session", store))
+	r.GET("/f/:id", ServeFileByID)
+
+	t.Run("whitelisted file type (avatar) accessed without authentication", func(t *testing.T) {
+		req, _ := http.NewRequest("GET", "/f/8001", nil)
+		w := httptest.NewRecorder()
+		r.ServeHTTP(w, req)
+
+		if w.Code != http.StatusOK {
+			t.Errorf("expected 200, got %d. Body: %s", w.Code, w.Body.String())
+		}
+		if w.Body.String() != "image" {
+			t.Errorf("expected 'image', got %q", w.Body.String())
+		}
+	})
+
+	t.Run("non-whitelisted file type (attachment) accessed without authentication returns 401", func(t *testing.T) {
+		req, _ := http.NewRequest("GET", "/f/8002", nil)
+		w := httptest.NewRecorder()
+		r.ServeHTTP(w, req)
+
+		if w.Code != http.StatusUnauthorized {
+			t.Errorf("expected 401, got %d. Body: %s", w.Code, w.Body.String())
+		}
+
+		var body map[string]any
+		if err := json.Unmarshal(w.Body.Bytes(), &body); err != nil {
+			t.Fatalf("failed to parse JSON: %v", err)
+		}
+		if body["error_msg"] != common.UnAuthorized {
+			t.Errorf("expected error_msg %q, got %v", common.UnAuthorized, body["error_msg"])
+		}
+	})
+
+	t.Run("non-whitelisted file type (attachment) accessed with valid token succeeds", func(t *testing.T) {
+		req, _ := http.NewRequest("GET", "/f/8002", nil)
+		req.Header.Set("X-Access-Token", tokenStr)
+		w := httptest.NewRecorder()
+		r.ServeHTTP(w, req)
+
+		if w.Code != http.StatusOK {
+			t.Errorf("expected 200, got %d. Body: %s", w.Code, w.Body.String())
+		}
+		if w.Body.String() != "bytes" {
+			t.Errorf("expected 'bytes', got %q", w.Body.String())
+		}
+	})
+
+	t.Run("accessing non-existent file returns 404", func(t *testing.T) {
+		req, _ := http.NewRequest("GET", "/f/9999", nil)
+		w := httptest.NewRecorder()
+		r.ServeHTTP(w, req)
+
+		if w.Code != http.StatusNotFound {
+			t.Errorf("expected 404, got %d", w.Code)
+		}
+	})
+}
+
+func TestGetDistinctUploadTypes(t *testing.T) {
+	dbConn, _, cleanup := testhelper.SetupTestEnvironment(t)
+	defer cleanup()
+
+	// Seed some uploads with new custom types
+	user := model.User{ID: 2222, Username: "test_user_2"}
+	dbConn.Create(&user)
+
+	customUpload := model.Upload{
+		ID:            9001,
+		UserID:        user.ID,
+		FileName:      "custom.txt",
+		FilePath:      "uploads/custom.txt",
+		FileSize:      10,
+		MimeType:      "text/plain",
+		Extension:     "txt",
+		StorageDriver: "local",
+		Type:          "custom_type_xyz",
+		Status:        model.UploadStatusUsed,
+	}
+	dbConn.Create(&customUpload)
+
+	gin.SetMode(gin.TestMode)
+	r := gin.New()
+	r.GET("/api/v1/admin/uploads/types", GetDistinctUploadTypes)
+
+	req, _ := http.NewRequest("GET", "/api/v1/admin/uploads/types", nil)
+	w := httptest.NewRecorder()
+	r.ServeHTTP(w, req)
+
+	if w.Code != http.StatusOK {
+		t.Fatalf("expected 200, got %d", w.Code)
+	}
+
+	var resp struct {
+		ErrorMsg string   `json:"error_msg"`
+		Data     []string `json:"data"`
+	}
+	if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
+		t.Fatalf("failed to parse JSON: %v", err)
+	}
+
+	if resp.ErrorMsg != "" {
+		t.Fatalf("unexpected error: %s", resp.ErrorMsg)
+	}
+
+	// Verify that custom_type_xyz and default types are present
+	hasCustom := false
+	hasAvatar := false
+	for _, typeName := range resp.Data {
+		if typeName == "custom_type_xyz" {
+			hasCustom = true
+		}
+		if typeName == "avatar" {
+			hasAvatar = true
+		}
+	}
+
+	if !hasCustom {
+		t.Errorf("expected custom_type_xyz to be in types, got: %v", resp.Data)
+	}
+	if !hasAvatar {
+		t.Errorf("expected avatar to be in types, got: %v", resp.Data)
+	}
+}