OpenFlare/backend/plugins/domain/upload/filesrv/file_server_test.go

// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0

package filesrv

import (
	"Wavelet/core/contracts"
	"Wavelet/pkg/response"
	"Wavelet/pkg/testhelper"
	"Wavelet/plugins/domain/upload/cache"
	"Wavelet/plugins/domain/upload/models"
	"Wavelet/plugins/domain/upload/shared"
	"bytes"
	"context"
	"crypto/sha256"
	"fmt"
	"image"
	"image/color"
	"image/png"
	"io"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"sync"
	"testing"
	"time"

	"github.com/gin-contrib/sessions"
	"github.com/gin-contrib/sessions/cookie"
	"github.com/gin-gonic/gin"

	uploadutil "Wavelet/plugins/domain/upload/util"
)

func init() {
	testhelper.RegisterCleanup(cache.ResetUploadMetaCacheForTest)
}

type localTestStorageService struct {
	mu   sync.RWMutex
	root string
}

func (s *localTestStorageService) Put(_ context.Context, key string, body io.Reader, _ int64, _ string) (contracts.StoragePutResult, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	path := filepath.Join(s.root, key)
	_ = os.MkdirAll(filepath.Dir(path), 0o755)
	f, err := os.Create(path)
	if err != nil {
		return contracts.StoragePutResult{}, err
	}
	defer f.Close()
	_, err = io.Copy(f, body)
	return contracts.StoragePutResult{Key: key, Bucket: "local"}, err
}

func (s *localTestStorageService) Get(ctx context.Context, key string) (*contracts.StorageObject, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	s.mu.RLock()
	defer s.mu.RUnlock()
	path := filepath.Join(s.root, key)
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	info, _ := f.Stat()
	return &contracts.StorageObject{
		Key:           key,
		Body:          f,
		ContentLength: info.Size(),
		ContentType:   "image/png",
	}, nil
}

func (s *localTestStorageService) Delete(_ context.Context, key string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	return os.Remove(filepath.Join(s.root, key))
}

func (s *localTestStorageService) Ingest(_ context.Context, _ io.Reader, _ contracts.IngestOptions) (*contracts.IngestResult, error) {
	return nil, nil
}

func TestServeFileByIDAccessControl(t *testing.T) {
	dbConn, cleanup := shared.SetupTestEnv(t)
	defer cleanup()
	cache.ResetAccessCaches()

	tempDir := t.TempDir()
	storageSvc := &localTestStorageService{root: tempDir}
	shared.SetStorageService(storageSvc)

	// Create a user in DB
	user := contracts.UserDTO{
		ID:       12345,
		Username: "file_test_user",
		IsActive: true,
	}
	if err := dbConn.Table("w_users").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 := fmt.Sprintf("%x", sha256.Sum256([]byte(tokenStr)))
	tokenRecord := map[string]any{
		"user_id":      user.ID,
		"name":         "test_token",
		"token_hash":   tokenHash,
		"masked_token": "test-***",
	}
	if err := dbConn.Table("w_access_tokens").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 := models.Upload{
		ID:         8001,
		UserID:     user.ID,
		FileName:   "avatar.png",
		FilePath:   "avatar.png",
		FileSize:   5,
		MimeType:   "image/png",
		Extension:  "png",
		Type:       "avatar",
		Status:     models.UploadStatusUsed,
		AccessMode: 1,
	}
	attachmentFile := models.Upload{
		ID:         8002,
		UserID:     user.ID,
		FileName:   "doc.pdf",
		FilePath:   "doc.pdf",
		FileSize:   5,
		MimeType:   "application/pdf",
		Extension:  "pdf",
		Type:       "attachment",
		Status:     models.UploadStatusUsed,
		AccessMode: 1,
	}

	if err := os.WriteFile(filepath.Join(tempDir, "avatar.png"), []byte("image"), 0o644); err != nil {
		t.Fatalf("failed to write avatar file: %v", err)
	}
	if err := os.WriteFile(filepath.Join(tempDir, "doc.pdf"), []byte("bytes"), 0o644); err != nil {
		t.Fatalf("failed to write attachment file: %v", err)
	}

	dbConn.Create(&avatarFile)
	dbConn.Create(&attachmentFile)

	gin.SetMode(gin.TestMode)
	r := gin.New()
	r.Use(response.ErrorHandlerMiddleware())
	store := cookie.NewStore([]byte("secret"))
	r.Use(sessions.Sessions("wavelet_session_id", store))
	r.GET("/f/:id", ServeFileByID)

	t.Run("public access allowed for whitelist type (avatar)", func(t *testing.T) {
		req, _ := http.NewRequest("GET", "/f/8001", nil)
		w := httptest.NewRecorder()
		r.ServeHTTP(w, req)

		if w.Code != http.StatusOK {
			t.Fatalf("expected status 200 for public file, got %d", w.Code)
		}
	})

	t.Run("public access rejected for non-whitelist type (attachment)", func(t *testing.T) {
		req, _ := http.NewRequest("GET", "/f/8002", nil)
		w := httptest.NewRecorder()
		r.ServeHTTP(w, req)

		if w.Code != http.StatusUnauthorized {
			t.Fatalf("expected status 401 for private file without auth, got %d", w.Code)
		}
	})

	t.Run("authenticated access allowed for non-whitelist type (attachment)", func(t *testing.T) {
		req, _ := http.NewRequest("GET", "/f/8002", nil)
		req.Header.Set("Authorization", "Bearer "+tokenStr)
		w := httptest.NewRecorder()
		r.ServeHTTP(w, req)

		if w.Code != http.StatusOK {
			t.Fatalf("expected status 200 for authenticated request, got %d", w.Code)
		}
	})

	t.Run("non-existent file returns 404", func(t *testing.T) {
		req, _ := http.NewRequest("GET", "/f/99999", nil)
		w := httptest.NewRecorder()
		r.ServeHTTP(w, req)

		if w.Code != http.StatusNotFound {
			t.Fatalf("expected status 404 for non-existent file, got %d", w.Code)
		}
	})

	t.Run("invalid id format returns 400", func(t *testing.T) {
		req, _ := http.NewRequest("GET", "/f/invalid_id", nil)
		w := httptest.NewRecorder()
		r.ServeHTTP(w, req)

		if w.Code != http.StatusBadRequest {
			t.Fatalf("expected status 400 for invalid id format, got %d", w.Code)
		}
	})
}

func TestServeFileByIDImageCompression(t *testing.T) {
	dbConn, cleanup := shared.SetupTestEnv(t)
	defer cleanup()
	cache.ResetAccessCaches()

	tempDir := t.TempDir()
	storageSvc := &localTestStorageService{root: tempDir}
	shared.SetStorageService(storageSvc)

	// Create test user
	user := contracts.UserDTO{
		ID:       54321,
		Username: "compress_test_user",
		IsActive: true,
	}
	dbConn.Table("w_users").Create(&user)

	// Create a small 1x1 test 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)
	}

	filePath := filepath.Join(tempDir, "test_image.png")
	if err := os.WriteFile(filePath, pngBuf.Bytes(), 0o644); err != nil {
		t.Fatalf("failed to write test png: %v", err)
	}

	// Save upload record to DB
	uploadRecord := models.Upload{
		ID:         3001,
		UserID:     user.ID,
		FileName:   "test_image.png",
		FilePath:   "test_image.png",
		FileSize:   int64(pngBuf.Len()),
		MimeType:   "image/png",
		Extension:  "png",
		Type:       "avatar", // Whitelisted by default
		Status:     models.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)
		}
		if w.Header().Get("Content-Type") != "image/png" {
			t.Errorf("expected Content-Type image/png, got %s", w.Header().Get("Content-Type"))
		}
		if w.Header().Get("X-Cache") != "" {
			t.Errorf("expected no X-Cache header for original file, got %s", w.Header().Get("X-Cache"))
		}
		if w.Header().Get("ETag") == "" {
			t.Errorf("expected ETag header for original file")
		}
	})

	t.Run("first request with quality=medium produces cache MISS and converts to WebP", 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)
		}
		if w.Header().Get("Content-Type") != "image/webp" {
			t.Errorf("expected Content-Type image/webp, got %s", w.Header().Get("Content-Type"))
		}
		if w.Header().Get("X-Cache") != "MISS" {
			t.Errorf("expected X-Cache MISS on first compress request, got %s", w.Header().Get("X-Cache"))
		}
		if w.Header().Get("ETag") == "" {
			t.Errorf("expected ETag header")
		}
		if len(w.Body.Bytes()) == 0 {
			t.Errorf("expected non-empty body")
		}
	})

	t.Run("second request with quality=medium produces cache HIT", 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)
		}
		if w.Header().Get("Content-Type") != "image/webp" {
			t.Errorf("expected Content-Type image/webp, got %s", w.Header().Get("Content-Type"))
		}
		if w.Header().Get("X-Cache") != "HIT" {
			t.Errorf("expected X-Cache HIT on second compress request, got %s", w.Header().Get("X-Cache"))
		}
	})

	t.Run("request with quality=origin behaves like original request", 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 w.Header().Get("X-Cache") != "" {
			t.Errorf("expected no X-Cache header for origin quality, got %s", w.Header().Get("X-Cache"))
		}
	})

	t.Run("conditional GET with matching If-None-Match returns 304", func(t *testing.T) {
		req, _ := http.NewRequest("GET", "/f/3001?quality=medium", nil)
		w := httptest.NewRecorder()
		r.ServeHTTP(w, req)

		etag := w.Header().Get("ETag")
		if etag == "" {
			t.Fatalf("expected ETag header from initial request")
		}

		// Second request with If-None-Match
		req2, _ := http.NewRequest("GET", "/f/3001?quality=medium", nil)
		req2.Header.Set("If-None-Match", etag)
		w2 := httptest.NewRecorder()
		r.ServeHTTP(w2, req2)

		if w2.Code != http.StatusNotModified {
			t.Fatalf("expected status 304 Not Modified, got %d", w2.Code)
		}
		if w2.Body.Len() != 0 {
			t.Errorf("expected empty body on 304 response, got %d bytes", w2.Body.Len())
		}
	})
}

// The singleflight body generates once on behalf of every concurrent requester
// sharing a cache key, so the caller that happens to arrive first must not be
// able to fail the others by disconnecting.
func TestEnsureCompressedImageCacheSurvivesCallerCancellation(t *testing.T) {
	tempDir := t.TempDir()
	shared.SetStorageService(&localTestStorageService{root: tempDir})

	img := image.NewRGBA(image.Rect(0, 0, 1, 1))
	img.Set(0, 0, color.RGBA{R: 0, G: 255, 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)
	}
	const filePath = "cancel_probe.png"
	if err := os.WriteFile(filepath.Join(tempDir, filePath), pngBuf.Bytes(), 0o600); err != nil {
		t.Fatalf("failed to write test png: %v", err)
	}

	upload := &models.Upload{
		ID:        990001,
		FilePath:  filePath,
		FileSize:  int64(pngBuf.Len()),
		MimeType:  "image/png",
		Extension: "png",
		// Unique per run so the persistent disk cache can never serve this key.
		Hash:      fmt.Sprintf("cancel-probe-%d", time.Now().UnixNano()),
		UpdatedAt: time.Now(),
	}

	ctx, cancel := context.WithCancel(context.Background())
	cancel()

	webpBytes, cached, err := EnsureCompressedImageCache(ctx, upload, "medium")
	if err != nil {
		t.Fatalf("compressed image generation failed with a canceled caller context: %v", err)
	}
	if cached {
		t.Errorf("expected a freshly generated image, got a cache hit")
	}
	if len(webpBytes) == 0 {
		t.Errorf("expected non-empty webp bytes")
	}
}

func TestNormalizeImageQuality(t *testing.T) {
	tests := []struct {
		name    string
		quality string
		want    string
	}{
		{name: "empty quality returns origin", quality: "", want: shared.ImageQualityOrigin},
		{name: "origin returns origin", quality: "origin", want: shared.ImageQualityOrigin},
		{name: "ORIGIN case-insensitive returns origin", quality: "ORIGIN", want: shared.ImageQualityOrigin},
		{name: "low returns low", quality: "low", want: shared.ImageQualityLow},
		{name: "LOW returns low", quality: "LOW", want: shared.ImageQualityLow},
		{name: "medium returns medium", quality: "medium", want: shared.ImageQualityMedium},
		{name: "high returns high", quality: "high", want: shared.ImageQualityHigh},
		{name: "unknown quality defaults to origin", quality: "ultra_hd", want: shared.ImageQualityOrigin},
		{name: "whitespace padded quality is trimmed", quality: "  medium  ", want: shared.ImageQualityMedium},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := uploadutil.NormalizeImageQuality(tt.quality); got != tt.want {
				t.Errorf("NormalizeImageQuality(%q) = %q, want %q", tt.quality, got, tt.want)
			}
		})
	}
}