xinmi/internal/resumable/service.go

package resumable

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/Tencent/WeKnora/internal/types"
	"github.com/google/uuid"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

type Service struct {
	db            *gorm.DB
	root          string
	authorize     AuthorizeFunc
	workerContext ContextFunc
	ingest        ImportFunc
	once          sync.Once
}

func NewService(db *gorm.DB, root string, authorize AuthorizeFunc, workerContext ContextFunc, ingest ImportFunc) (*Service, error) {
	if root == "" {
		root = filepath.Join("data", "resumable-uploads")
	}
	root, err := filepath.Abs(root)
	if err != nil {
		return nil, err
	}
	if err = os.MkdirAll(root, 0700); err != nil {
		return nil, err
	}
	return &Service{db: db, root: root, authorize: authorize, workerContext: workerContext, ingest: ingest}, nil
}

func identity(ctx context.Context) (uint64, string, error) {
	t, ok := types.TenantIDFromContext(ctx)
	u, userOK := types.UserIDFromContext(ctx)
	if !ok || t == 0 || !userOK || u == "" || types.IsSyntheticUserID(u) {
		return 0, "", ErrForbidden
	}
	return t, u, nil
}
func (s *Service) scoped(ctx context.Context) (*gorm.DB, error) {
	t, u, e := identity(ctx)
	if e != nil {
		return nil, e
	}
	return s.db.WithContext(ctx).Where("tenant_id=? AND owner_id=?", t, u), nil
}
func validHash(v string) bool {
	b, e := hex.DecodeString(v)
	return e == nil && len(b) == 32 && strings.ToLower(v) == v
}
func (s *Service) path(id string) string         { return filepath.Join(s.root, id+".upload") }
func (s *Service) materialPath(id string) string { return filepath.Join(s.root, id+".material") }

// FilePath is the durable raw upload bytes on disk.
func (s *Service) FilePath(id string) string { return s.path(id) }

// SaveMaterial stores parsed digital-human source text beside the upload file
// so result jsonb only needs small metadata.
func (s *Service) SaveMaterial(id string, source types.DigitalHumanSource) error {
	if id == "" || strings.TrimSpace(source.Name) == "" || source.Text == "" {
		return ErrInvalid
	}
	data, err := json.Marshal(source)
	if err != nil {
		return err
	}
	tmp := s.materialPath(id) + ".tmp"
	if err = os.WriteFile(tmp, data, 0600); err != nil {
		return err
	}
	return os.Rename(tmp, s.materialPath(id))
}

// LoadMaterial reads the file-backed digital-human source for an upload.
func (s *Service) LoadMaterial(id string) (types.DigitalHumanSource, error) {
	data, err := os.ReadFile(s.materialPath(id))
	if err != nil {
		return types.DigitalHumanSource{}, err
	}
	var source types.DigitalHumanSource
	if err = json.Unmarshal(data, &source); err != nil {
		return types.DigitalHumanSource{}, err
	}
	if strings.TrimSpace(source.Name) == "" || source.Text == "" {
		return types.DigitalHumanSource{}, ErrInvalid
	}
	return source, nil
}

func (s *Service) removeStorage(id string) error {
	for _, p := range []string{s.path(id), s.materialPath(id), s.materialPath(id) + ".tmp"} {
		if e := os.Remove(p); e != nil && !errors.Is(e, os.ErrNotExist) {
			return e
		}
	}
	return nil
}

// listSelect projects columns for List. Summary mode never loads full result
// jsonb payloads — only a tiny selection object when present.
func (s *Service) listSelect(summary bool) string {
	cols := `id, tenant_id, owner_id, kind, target_id, name, size, "offset", chunk_size, sha256, fingerprint, state, error, options, ingest_offset, part, expires_at, lease_until, lease_token, created_at, updated_at`
	if !summary {
		return cols + ", result"
	}
	switch s.db.Dialector.Name() {
	case "postgres":
		return cols + `, CASE WHEN result->>'selection' IN ('full','sampled') THEN jsonb_build_object('selection', result->'selection') ELSE NULL END AS result`
	default:
		return cols + `, CASE WHEN json_extract(result, '$.selection') IN ('full','sampled') THEN json_object('selection', json_extract(result, '$.selection')) ELSE NULL END AS result`
	}
}

func (s *Service) Create(ctx context.Context, in CreateInput) (*Upload, error) {
	t, u, e := identity(ctx)
	if e != nil {
		return nil, e
	}
	if (in.Kind != "knowledge" && in.Kind != "digital-human") || len(in.TargetID) > 36 || in.TargetID == "" || in.Size < 1 || in.Size > MaxFileSize || len(in.Name) > 500 || strings.TrimSpace(in.Name) == "" || strings.ContainsAny(in.Name, "\x00\r\n") || len(in.Fingerprint) > 128 || (in.SHA256 != "" && !validHash(in.SHA256)) || len(in.Options) > 32<<10 {
		return nil, ErrInvalid
	}
	if s.authorize == nil {
		return nil, ErrForbidden
	}
	if e = s.authorize(ctx, in.Kind, in.TargetID); e != nil {
		return nil, e
	}
	if len(in.Options) == 0 {
		in.Options = types.JSON(`{}`)
	}
	row := &Upload{ID: uuid.NewString(), TenantID: t, OwnerID: u, Kind: in.Kind, TargetID: in.TargetID, Name: filepath.Base(in.Name), Size: in.Size, ChunkSize: ChunkSize, SHA256: in.SHA256, Fingerprint: in.Fingerprint, Options: in.Options, Result: types.JSON(`{}`), State: "uploading", ExpiresAt: time.Now().Add(7 * 24 * time.Hour)}
	e = s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		// Serialize reservations across every user in this tenant.
		var tenant types.Tenant
		if e := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=?", t).Take(&tenant).Error; e != nil {
			return e
		}
		var totals struct {
			TenantBytes int64
			UserBytes   int64
			Active      int64
		}
		e := tx.Model(&Upload{}).Select("COALESCE(SUM(size),0) AS tenant_bytes, COALESCE(SUM(CASE WHEN owner_id=? THEN size ELSE 0 END),0) AS user_bytes, COALESCE(SUM(CASE WHEN owner_id=? AND state IN ('uploading','queued','verifying','processing') THEN 1 ELSE 0 END),0) AS active", u, u).Where("tenant_id=? AND state NOT IN ('cancelled','expired')", t).Scan(&totals).Error
		if e != nil {
			return e
		}
		if totals.Active >= MaxActive || totals.UserBytes+in.Size > UserQuota || totals.TenantBytes+in.Size > TenantQuota {
			return ErrQuota
		}
		if tenant.StorageQuota > 0 && tenant.StorageUsed+totals.TenantBytes+in.Size > tenant.StorageQuota {
			return ErrQuota
		}
		return tx.Create(row).Error
	})
	return row, e
}

func (s *Service) Get(ctx context.Context, id string) (*Upload, error) {
	q, e := s.scoped(ctx)
	if e != nil {
		return nil, e
	}
	var row Upload
	e = q.Where("id=?", id).Take(&row).Error
	if e == nil {
		if s.authorize == nil {
			return nil, ErrForbidden
		}
		if e = s.authorize(ctx, row.Kind, row.TargetID); e != nil {
			return nil, e
		}
	}
	if e == nil && row.ExpiresAt.Before(time.Now()) && row.State == "uploading" {
		row.State = "expired"
	}
	return &row, e
}

// Chunks exposes bounded digests so a reselected local file can verify every
// committed byte before sending its remaining suffix.
func (s *Service) Chunks(ctx context.Context, id string, offset int64) ([]Chunk, error) {
	row, err := s.Get(ctx, id)
	if err != nil {
		return nil, err
	}
	if offset < 0 || offset > row.Offset || row.ChunkSize <= 0 || (offset != row.Offset && offset%row.ChunkSize != 0) {
		return nil, ErrInvalid
	}
	chunks := []Chunk{}
	err = s.db.WithContext(ctx).Where("upload_id=? AND \"offset\">=? AND \"offset\"<?", row.ID, offset, row.Offset).Order("\"offset\"").Limit(128).Find(&chunks).Error
	return chunks, err
}

func (s *Service) List(ctx context.Context, kind, target string, summary bool) ([]Upload, error) {
	q, e := s.scoped(ctx)
	if e != nil {
		return nil, e
	}
	rows := []Upload{}
	q = q.Select(s.listSelect(summary)).Where("state NOT IN ('cancelled','expired')")
	if kind != "" {
		q = q.Where("kind=?", kind)
	}
	if target != "" {
		q = q.Where("target_id=?", target)
	}
	e = q.Order("created_at DESC").Limit(100).Find(&rows).Error
	if e != nil {
		return nil, e
	}
	visible := make([]Upload, 0, len(rows))
	type authKey struct{ kind, target string }
	authCache := map[authKey]error{}
	for _, row := range rows {
		if s.authorize == nil {
			continue
		}
		key := authKey{row.Kind, row.TargetID}
		if _, ok := authCache[key]; !ok {
			authCache[key] = s.authorize(ctx, row.Kind, row.TargetID)
		}
		if authCache[key] == nil {
			visible = append(visible, row)
		}
	}
	return visible, nil
}

// Put uses a bounded buffer and fsyncs before committing the offset. A retry is
// accepted only for the exact persisted chunk (including size and checksum).
func (s *Service) Put(ctx context.Context, id string, offset, size int64, checksum string, body io.Reader) (*Upload, error) {
	if offset < 0 || size < 1 || size > ChunkSize || !validHash(checksum) {
		return nil, ErrInvalid
	}
	t, u, e := identity(ctx)
	if e != nil {
		return nil, e
	}
	var row Upload
	e = s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		if e := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=? AND tenant_id=? AND owner_id=?", id, t, u).Take(&row).Error; e != nil {
			return e
		}
		if row.State != "uploading" || row.ExpiresAt.Before(time.Now()) {
			return ErrConflict
		}
		if s.authorize == nil {
			return ErrForbidden
		}
		if e := s.authorize(ctx, row.Kind, row.TargetID); e != nil {
			return e
		}
		if row.ChunkSize <= 0 || offset%row.ChunkSize != 0 || size != min(row.ChunkSize, row.Size-offset) {
			return ErrInvalid
		}
		if offset < row.Offset {
			var chunk Chunk
			if e := tx.Where("upload_id=? AND \"offset\"=?", id, offset).Take(&chunk).Error; e != nil {
				return ErrConflict
			}
			if chunk.Size != size || chunk.SHA256 != checksum {
				return ErrConflict
			}
			h := sha256.New()
			n, e := io.Copy(h, io.LimitReader(body, size+1))
			if e != nil {
				return e
			}
			if n != size || hex.EncodeToString(h.Sum(nil)) != checksum {
				return ErrInvalid
			}
			return nil
		}
		if offset != row.Offset || offset+size > row.Size {
			return ErrConflict
		}
		f, e := os.OpenFile(s.path(id), os.O_CREATE|os.O_RDWR, 0600)
		if e != nil {
			return e
		}
		defer f.Close()
		stat, e := f.Stat()
		if e != nil {
			return e
		}
		if stat.Size() < offset {
			return fmt.Errorf("上传暂存文件缺失，请重新上传")
		}
		if e = f.Truncate(offset); e != nil {
			return e
		}
		if _, e = f.Seek(offset, io.SeekStart); e != nil {
			return e
		}
		h := sha256.New()
		n, e := io.Copy(io.MultiWriter(f, h), io.LimitReader(body, size+1))
		if e != nil || n != size || hex.EncodeToString(h.Sum(nil)) != checksum {
			_ = f.Truncate(offset)
			if e != nil {
				return e
			}
			return ErrInvalid
		}
		if e = f.Sync(); e != nil {
			return e
		}
		if e = tx.Create(&Chunk{UploadID: id, Offset: offset, Size: size, SHA256: checksum}).Error; e != nil {
			return e
		}
		row.Offset += size
		row.ExpiresAt = time.Now().Add(7 * 24 * time.Hour)
		return tx.Model(&row).Updates(map[string]any{"offset": row.Offset, "expires_at": row.ExpiresAt, "updated_at": time.Now()}).Error
	})
	return &row, e
}

func (s *Service) Finalize(ctx context.Context, id string) (*Upload, error) {
	t, u, e := identity(ctx)
	if e != nil {
		return nil, e
	}
	var row Upload
	e = s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		if e := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=? AND tenant_id=? AND owner_id=?", id, t, u).Take(&row).Error; e != nil {
			return e
		}
		if e := s.authorize(ctx, row.Kind, row.TargetID); e != nil {
			return e
		}
		if row.State == "queued" || row.State == "verifying" || row.State == "processing" || row.State == "completed" {
			return nil
		}
		if (row.State != "uploading" && row.State != "failed") || row.Offset != row.Size || row.ExpiresAt.Before(time.Now()) {
			return ErrConflict
		}
		row.State = "queued"
		row.Error = ""
		row.ExpiresAt = time.Now().Add(7 * 24 * time.Hour)
		return tx.Model(&row).Updates(map[string]any{"state": row.State, "error": "", "expires_at": row.ExpiresAt, "updated_at": time.Now()}).Error
	})
	return &row, e
}
func (s *Service) Cancel(ctx context.Context, id string) error {
	t, u, e := identity(ctx)
	if e != nil {
		return e
	}
	return s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		var row Upload
		if e := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=? AND tenant_id=? AND owner_id=?", id, t, u).Take(&row).Error; e != nil {
			return e
		}
		if row.State == "verifying" || row.State == "processing" {
			return ErrConflict
		}
		if e := s.removeStorage(id); e != nil {
			return e
		}
		if e := tx.Where("upload_id=?", id).Delete(&Chunk{}).Error; e != nil {
			return e
		}
		return tx.Model(&row).Updates(map[string]any{"state": "cancelled", "result": types.JSON(`{}`), "updated_at": time.Now()}).Error
	})
}

func (s *Service) Cleanup(ctx context.Context) error {
	var rows []Upload
	if e := s.db.WithContext(ctx).Where("expires_at < ? AND (lease_until IS NULL OR lease_until < ?) AND state NOT IN ('expired','cancelled')", time.Now(), time.Now()).Limit(100).Find(&rows).Error; e != nil {
		return e
	}
	for _, row := range rows {
		if e := s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
			var current Upload
			if e := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=?", row.ID).Take(&current).Error; e != nil {
				return e
			}
			if current.ExpiresAt.After(time.Now()) || (current.LeaseUntil != nil && current.LeaseUntil.After(time.Now())) {
				return nil
			}
			if e := s.removeStorage(row.ID); e != nil {
				return e
			}
			if e := tx.Where("upload_id=?", row.ID).Delete(&Chunk{}).Error; e != nil {
				return e
			}
			return tx.Model(&current).Updates(map[string]any{"state": "expired", "result": types.JSON(`{}`), "updated_at": time.Now()}).Error
		}); e != nil {
			return e
		}
	}
	return nil
}