OpenFlare/internal/apps/upload/file_server_test.go

// 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 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)
	}
}