xinmi/internal/resumable/service_test.go

package resumable

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/Tencent/WeKnora/internal/types"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/require"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
)

func uploadContext(tenant uint64, user string) context.Context {
	ctx := context.WithValue(context.Background(), types.TenantIDContextKey, tenant)
	return context.WithValue(ctx, types.UserIDContextKey, user)
}
func uploadService(t *testing.T) *Service {
	t.Helper()
	db, e := gorm.Open(sqlite.Open(filepath.Join(t.TempDir(), "uploads.db")), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)})
	require.NoError(t, e)
	require.NoError(t, db.AutoMigrate(&Upload{}, &Chunk{}))
	require.NoError(t, db.Exec("CREATE TABLE tenants(id INTEGER PRIMARY KEY,storage_quota BIGINT,storage_used BIGINT,deleted_at DATETIME)").Error)
	require.NoError(t, db.Exec("INSERT INTO tenants(id,storage_quota,storage_used) VALUES(1,0,0),(2,0,0)").Error)
	sql, e := db.DB()
	require.NoError(t, e)
	sql.SetMaxOpenConns(1)
	t.Cleanup(func() { _ = sql.Close() })
	s, e := NewService(db, t.TempDir(), func(context.Context, string, string) error { return nil }, func(ctx context.Context, _ *Upload) (context.Context, error) { return ctx, nil }, func(_ context.Context, u *Upload, _ string) (int64, types.JSON, bool, error) {
		return u.Size, types.JSON(`{"ok":true}`), true, nil
	})
	require.NoError(t, e)
	return s
}
func sum(data []byte) string { h := sha256.Sum256(data); return hex.EncodeToString(h[:]) }
func create(t *testing.T, s *Service, data []byte) *Upload {
	t.Helper()
	u, e := s.Create(uploadContext(1, "alice"), CreateInput{Kind: "knowledge", TargetID: "kb", Name: "fixture.txt", Size: int64(len(data)), SHA256: sum(data)})
	require.NoError(t, e)
	return u
}

func TestUploadResumeChecksumOwnershipAndCrashRecovery(t *testing.T) {
	s := uploadService(t)
	ctx := uploadContext(1, "alice")
	data := []byte("abcdef")
	u := create(t, s, data)
	require.NoError(t, s.db.Model(&Upload{}).Where("id=?", u.ID).Update("chunk_size", 3).Error)
	_, e := s.Put(ctx, u.ID, 0, 3, sum([]byte("wrong")), bytes.NewReader(data[:3]))
	require.ErrorIs(t, e, ErrInvalid)
	row, e := s.Get(ctx, u.ID)
	require.NoError(t, e)
	require.Zero(t, row.Offset)
	_, e = s.Put(uploadContext(1, "bob"), u.ID, 0, 3, sum(data[:3]), bytes.NewReader(data[:3]))
	require.ErrorIs(t, e, gorm.ErrRecordNotFound)
	_, e = s.Get(uploadContext(2, "alice"), u.ID)
	require.ErrorIs(t, e, gorm.ErrRecordNotFound)
	row, e = s.Put(ctx, u.ID, 0, 3, sum(data[:3]), bytes.NewReader(data[:3]))
	require.NoError(t, e)
	require.EqualValues(t, 3, row.Offset)
	chunks, e := s.Chunks(ctx, u.ID, 0)
	require.NoError(t, e)
	require.Len(t, chunks, 1)
	require.Equal(t, sum(data[:3]), chunks[0].SHA256)
	_, e = s.Chunks(uploadContext(1, "bob"), u.ID, 0)
	require.ErrorIs(t, e, gorm.ErrRecordNotFound)
	_, e = s.Chunks(ctx, u.ID, 1)
	require.ErrorIs(t, e, ErrInvalid)
	// A response lost after commit may replay exactly the same chunk.
	row, e = s.Put(ctx, u.ID, 0, 3, sum(data[:3]), bytes.NewReader(data[:3]))
	require.NoError(t, e)
	require.EqualValues(t, 3, row.Offset)
	_, e = s.Put(ctx, u.ID, 0, 3, sum(data[3:]), bytes.NewReader(data[3:]))
	require.ErrorIs(t, e, ErrConflict)
	_, e = s.Finalize(ctx, u.ID)
	require.ErrorIs(t, e, ErrConflict)
	// Bytes written before a killed transaction must not survive the resume.
	f, e := os.OpenFile(s.path(u.ID), os.O_APPEND|os.O_WRONLY, 0600)
	require.NoError(t, e)
	_, e = f.WriteString("uncommitted")
	require.NoError(t, e)
	require.NoError(t, f.Close())
	restarted, e := NewService(s.db, s.root, s.authorize, s.workerContext, s.ingest)
	require.NoError(t, e)
	row, e = restarted.Put(ctx, u.ID, 3, 3, sum(data[3:]), bytes.NewReader(data[3:]))
	require.NoError(t, e)
	require.EqualValues(t, 6, row.Offset)
	stored, e := os.ReadFile(s.path(u.ID))
	require.NoError(t, e)
	require.Equal(t, data, stored)
	row, e = s.Finalize(ctx, u.ID)
	require.NoError(t, e)
	require.Equal(t, "queued", row.State)
	_, e = s.Finalize(ctx, u.ID)
	require.NoError(t, e)
	worked, e := s.ProcessNext(context.Background())
	require.NoError(t, e)
	require.True(t, worked)
	row, e = s.Get(ctx, u.ID)
	require.NoError(t, e)
	require.Equal(t, "completed", row.State)
	require.EqualValues(t, 6, row.IngestOffset)
}

func TestUploadBoundsQuotaCancelAndExpiry(t *testing.T) {
	s := uploadService(t)
	ctx := uploadContext(1, "alice")
	in := CreateInput{Kind: "knowledge", TargetID: "kb", Name: "large.txt", Size: MaxFileSize + 1}
	_, e := s.Create(ctx, in)
	require.ErrorIs(t, e, ErrInvalid)
	in.Size = MaxFileSize
	first, e := s.Create(ctx, in)
	require.NoError(t, e)
	require.Equal(t, MaxFileSize, first.Size)
	_, e = s.Create(ctx, in)
	require.NoError(t, e)
	in.Size = 1
	_, e = s.Create(ctx, in)
	require.ErrorIs(t, e, ErrQuota)
	require.NoError(t, s.Cancel(ctx, first.ID))
	small, e := s.Create(ctx, in)
	require.NoError(t, e)
	_, e = s.Put(ctx, small.ID, 0, ChunkSize+1, strings.Repeat("a", 64), strings.NewReader("x"))
	require.ErrorIs(t, e, ErrInvalid)
	_, e = s.Put(ctx, small.ID, 0, 1, sum([]byte("x")), strings.NewReader("xx"))
	require.ErrorIs(t, e, ErrInvalid)
	_, e = s.Put(ctx, small.ID, 0, 1, sum([]byte("x")), strings.NewReader("x"))
	require.NoError(t, e)
	require.NoError(t, s.db.Model(&Upload{}).Where("id=?", small.ID).Update("expires_at", time.Now().Add(-time.Hour)).Error)
	_, e = s.Finalize(ctx, small.ID)
	require.ErrorIs(t, e, ErrConflict)
	require.NoError(t, s.Cleanup(ctx))
	_, e = os.Stat(s.path(small.ID))
	require.True(t, os.IsNotExist(e))
	row, e := s.Get(ctx, small.ID)
	require.NoError(t, e)
	require.Equal(t, "expired", row.State)
	var chunks int64
	require.NoError(t, s.db.Model(&Chunk{}).Where("upload_id=?", small.ID).Count(&chunks).Error)
	require.Zero(t, chunks)
}

func TestUploadWorkerChecksWholeFileAndReauthorizes(t *testing.T) {
	for _, scenario := range []string{"corruption", "revoked"} {
		t.Run(scenario, func(t *testing.T) {
			s := uploadService(t)
			ctx := uploadContext(1, "alice")
			data := []byte("source")
			u := create(t, s, data)
			_, e := s.Put(ctx, u.ID, 0, int64(len(data)), sum(data), bytes.NewReader(data))
			require.NoError(t, e)
			_, e = s.Finalize(ctx, u.ID)
			require.NoError(t, e)
			called := false
			s.ingest = func(context.Context, *Upload, string) (int64, types.JSON, bool, error) {
				called = true
				return 6, types.JSON(`{}`), true, nil
			}
			if scenario == "corruption" {
				require.NoError(t, s.db.Model(&Upload{}).Where("id=?", u.ID).Update("sha256", "").Error)
				require.NoError(t, os.WriteFile(s.path(u.ID), []byte("edited"), 0600))
			} else {
				s.authorize = func(context.Context, string, string) error { return ErrForbidden }
			}
			worked, e := s.ProcessNext(context.Background())
			require.NoError(t, e)
			require.True(t, worked)
			var row Upload
			require.NoError(t, s.db.Where("id=?", u.ID).Take(&row).Error)
			require.Equal(t, "failed", row.State)
			require.False(t, called)
		})
	}
}

func TestUploadRevocationHidesEvidenceAndStopsWritesButAllowsCleanup(t *testing.T) {
	s := uploadService(t)
	ctx := uploadContext(1, "alice")
	u := create(t, s, []byte("source"))
	_, e := s.Put(ctx, u.ID, 0, 6, sum([]byte("source")), strings.NewReader("source"))
	require.NoError(t, e)
	s.authorize = func(context.Context, string, string) error { return ErrForbidden }
	_, e = s.Get(ctx, u.ID)
	require.ErrorIs(t, e, ErrForbidden)
	_, e = s.Put(ctx, u.ID, 0, 6, sum([]byte("source")), strings.NewReader("source"))
	require.ErrorIs(t, e, ErrForbidden)
	_, e = s.Finalize(ctx, u.ID)
	require.ErrorIs(t, e, ErrForbidden)
	rows, e := s.List(ctx, "", "", false)
	require.NoError(t, e)
	require.Empty(t, rows)
	require.NoError(t, s.Cancel(ctx, u.ID))
	var count int64
	require.NoError(t, s.db.Model(&Chunk{}).Where("upload_id=?", u.ID).Count(&count).Error)
	require.Zero(t, count)
	_, e = os.Stat(s.path(u.ID))
	require.True(t, os.IsNotExist(e))
}

func TestUploadWorkerDurableUnitsAndExpiredLease(t *testing.T) {
	s := uploadService(t)
	ctx := uploadContext(1, "alice")
	u := create(t, s, []byte("ab"))
	_, e := s.Put(ctx, u.ID, 0, 2, sum([]byte("ab")), strings.NewReader("ab"))
	require.NoError(t, e)
	_, e = s.Finalize(ctx, u.ID)
	require.NoError(t, e)
	s.ingest = func(_ context.Context, row *Upload, _ string) (int64, types.JSON, bool, error) {
		return row.IngestOffset + 1, types.JSON(`{"checkpoint":true}`), row.IngestOffset == 1, nil
	}
	worked, e := s.ProcessNext(context.Background())
	require.NoError(t, e)
	require.True(t, worked)
	row, e := s.Get(ctx, u.ID)
	require.NoError(t, e)
	require.EqualValues(t, 1, row.IngestOffset)
	require.Equal(t, "processing", row.State)
	require.NoError(t, s.db.Model(&Upload{}).Where("id=?", u.ID).Updates(map[string]any{"lease_until": time.Now().Add(time.Hour), "lease_token": "crashed-worker"}).Error)
	worked, e = s.ProcessNext(context.Background())
	require.NoError(t, e)
	require.False(t, worked)
	require.NoError(t, s.db.Model(&Upload{}).Where("id=?", u.ID).Update("lease_until", time.Now().Add(-time.Minute)).Error)
	worked, e = s.ProcessNext(context.Background())
	require.NoError(t, e)
	require.True(t, worked)
	row, e = s.Get(ctx, u.ID)
	require.NoError(t, e)
	require.Equal(t, "completed", row.State)
	require.EqualValues(t, 2, row.IngestOffset)
}

func TestUploadWorkerAndHTTPNeverExposeInternalFailures(t *testing.T) {
	const secret = "postgres://private:credential-sentinel@internal-host/database"
	s := uploadService(t)
	ctx := uploadContext(1, "alice")
	u := create(t, s, []byte("ok"))
	_, e := s.Put(ctx, u.ID, 0, 2, sum([]byte("ok")), strings.NewReader("ok"))
	require.NoError(t, e)
	_, e = s.Finalize(ctx, u.ID)
	require.NoError(t, e)
	s.ingest = func(context.Context, *Upload, string) (int64, types.JSON, bool, error) {
		return 0, nil, false, errors.New(secret)
	}
	worked, e := s.ProcessNext(context.Background())
	require.NoError(t, e)
	require.True(t, worked)
	row, e := s.Get(ctx, u.ID)
	require.NoError(t, e)
	require.Equal(t, "failed", row.State)
	require.NotContains(t, row.Error, "credential-sentinel")
	require.Contains(t, row.Error, "服务日志")
	r := gin.New()
	r.GET("/error", func(c *gin.Context) { respond(c, 200, nil, errors.New(secret)) })
	w := httptest.NewRecorder()
	r.ServeHTTP(w, httptest.NewRequest("GET", "/error", nil))
	require.Equal(t, 500, w.Code)
	require.NotContains(t, w.Body.String(), "credential-sentinel")
	// A retry preserves the completed transport and checkpoint.
	_, e = s.Finalize(ctx, u.ID)
	require.NoError(t, e)
	s.ingest = func(_ context.Context, row *Upload, _ string) (int64, types.JSON, bool, error) {
		return row.Size, types.JSON(`{}`), true, nil
	}
	worked, e = s.ProcessNext(context.Background())
	require.NoError(t, e)
	require.True(t, worked)
	row, e = s.Get(ctx, u.ID)
	require.NoError(t, e)
	require.Equal(t, "completed", row.State)
	require.Empty(t, row.Error)
}