// Copyright 2025 linux.do // Copyright 2026 Arctel.net // SPDX-License-Identifier: Apache-2.0 package upload import ( "bytes" "encoding/json" "image" "image/color" "image/png" "net/http" "net/http/httptest" "os" "testing" "github.com/Rain-kl/Wavelet/internal/common" "github.com/Rain-kl/Wavelet/internal/diskcache" "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, AccessMode: 1, } 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, AccessMode: 1, } _ = 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 only custom_type_xyz is present if len(resp.Data) != 1 || resp.Data[0] != "custom_type_xyz" { t.Errorf("expected only custom_type_xyz in types list, got: %v", resp.Data) } } func TestImageCompression(t *testing.T) { dbConn, _, cleanup := testhelper.SetupTestEnvironment(t) defer cleanup() cache := diskcache.GetGlobalCache() if err := cache.Clear(); err != nil { t.Fatalf("failed to clear disk cache before test: %v", err) } // Ensure uploads dir is cleaned up defer func() { _ = os.RemoveAll("uploads") }() defer func() { if err := cache.Clear(); err != nil { t.Errorf("failed to clear disk cache after test: %v", err) } }() // Create test user user := model.User{ ID: 555, Username: "compress_tester", IsActive: true, } dbConn.Create(&user) // Create a 1x1 pixel PNG image img := image.NewRGBA(image.Rect(0, 0, 1, 1)) img.Set(0, 0, color.RGBA{R: 255, G: 0, B: 0, A: 255}) var pngBuf bytes.Buffer if err := png.Encode(&pngBuf, img); err != nil { t.Fatalf("failed to encode test png: %v", err) } _ = os.MkdirAll("uploads", 0755) filePath := "uploads/test_image.png" if err := os.WriteFile(filePath, pngBuf.Bytes(), 0644); err != nil { t.Fatalf("failed to write test png: %v", err) } // Save upload record to DB uploadRecord := model.Upload{ ID: 3001, UserID: user.ID, FileName: "test_image.png", FilePath: filePath, FileSize: int64(pngBuf.Len()), MimeType: "image/png", Extension: "png", StorageDriver: "local", Type: "avatar", // Whitelisted by default Status: model.UploadStatusUsed, AccessMode: 1, } dbConn.Create(&uploadRecord) // Setup Router gin.SetMode(gin.TestMode) r := gin.New() r.GET("/f/:id", ServeFileByID) t.Run("serve original file without compress parameter", func(t *testing.T) { req, _ := http.NewRequest("GET", "/f/3001", nil) w := httptest.NewRecorder() r.ServeHTTP(w, req) if w.Code != http.StatusOK { t.Fatalf("expected status 200, got %d", w.Code) } // Content-Type should be image/png (default local serving type) if w.Header().Get("Content-Type") != "image/png" { t.Errorf("expected Content-Type image/png, got %s", w.Header().Get("Content-Type")) } if len(w.Body.Bytes()) != pngBuf.Len() { t.Errorf("expected body size %d, got %d", pngBuf.Len(), len(w.Body.Bytes())) } }) t.Run("serve compressed WebP file with medium quality", func(t *testing.T) { req, _ := http.NewRequest("GET", "/f/3001?quality=medium", nil) w := httptest.NewRecorder() r.ServeHTTP(w, req) if w.Code != http.StatusOK { t.Fatalf("expected status 200, got %d. Body: %s", w.Code, w.Body.String()) } // Content-Type should be image/webp if w.Header().Get("Content-Type") != "image/webp" { t.Errorf("expected Content-Type image/webp, got %s", w.Header().Get("Content-Type")) } cacheKey := imageCompressionCacheKey(&uploadRecord, imageQualityMedium) cachedBytes, err := cache.Get(cacheKey) if err != nil { t.Fatalf("disk cache Get(%q) returned error: %v", cacheKey, err) } if !bytes.Equal(cachedBytes, w.Body.Bytes()) { t.Errorf("cached compressed image differs from response") } if err := os.Remove(filePath); err != nil { t.Fatalf("failed to remove source image before cache-hit request: %v", err) } t.Cleanup(func() { if err := os.WriteFile(filePath, pngBuf.Bytes(), 0644); err != nil { t.Errorf("failed to restore source image: %v", err) } }) w2 := httptest.NewRecorder() r.ServeHTTP(w2, req) if w2.Code != http.StatusOK { t.Fatalf("expected status 200, got %d", w2.Code) } if !bytes.Equal(w2.Body.Bytes(), cachedBytes) { t.Errorf("cache-hit response differs from cached compressed image") } }) t.Run("serve compressed WebP file and check cache headers and 304 Not Modified", func(t *testing.T) { req, _ := http.NewRequest("GET", "/f/3001?quality=medium", nil) w := httptest.NewRecorder() r.ServeHTTP(w, req) if w.Code != http.StatusOK { t.Fatalf("expected status 200, got %d", w.Code) } etag := w.Header().Get("ETag") if etag == "" { t.Error("expected ETag header, got empty") } cacheControl := w.Header().Get("Cache-Control") if cacheControl != "public, max-age=31536000" { t.Errorf("expected Cache-Control 'public, max-age=31536000', got %q", cacheControl) } // Perform conditional GET request reqCond, _ := http.NewRequest("GET", "/f/3001?quality=medium", nil) reqCond.Header.Set("If-None-Match", etag) wCond := httptest.NewRecorder() r.ServeHTTP(wCond, reqCond) if wCond.Code != http.StatusNotModified { t.Errorf("expected status 304, got %d", wCond.Code) } }) t.Run("serve original file with origin quality", func(t *testing.T) { req, _ := http.NewRequest("GET", "/f/3001?quality=origin", nil) w := httptest.NewRecorder() r.ServeHTTP(w, req) if w.Code != http.StatusOK { t.Fatalf("expected status 200, got %d", w.Code) } if w.Header().Get("Content-Type") != "image/png" { t.Errorf("expected Content-Type image/png, got %s", w.Header().Get("Content-Type")) } if !bytes.Equal(w.Body.Bytes(), pngBuf.Bytes()) { t.Errorf("origin-quality response differs from original image") } }) } func TestNormalizeImageQuality(t *testing.T) { tests := []struct { name string quality string want string }{ {name: imageQualityLow, quality: imageQualityLow, want: imageQualityLow}, {name: imageQualityMedium, quality: imageQualityMedium, want: imageQualityMedium}, {name: imageQualityHigh, quality: imageQualityHigh, want: imageQualityHigh}, {name: "origin", quality: "origin", want: "origin"}, {name: "uppercase", quality: "LOW", want: imageQualityLow}, {name: "empty", quality: "", want: "origin"}, {name: "invalid", quality: "maximum", want: "origin"}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { if got := normalizeImageQuality(tt.quality); got != tt.want { t.Errorf("normalizeImageQuality(%q) = %q, want %q", tt.quality, got, tt.want) } }) } }