WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes

internal/resumable/service.go

@@ -0,0 +1,335 @@
+package resumable
+
+import (
+	"context"
+	"crypto/sha256"
+	"encoding/hex"
+	"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) 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) ([]Upload, error) {
+	q, e := s.scoped(ctx)
+	if e != nil {
+		return nil, e
+	}
+	rows := []Upload{}
+	q = q.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))
+	for _, row := range rows {
+		if s.authorize != nil && s.authorize(ctx, row.Kind, row.TargetID) == 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 := os.Remove(s.path(id)); e != nil && !errors.Is(e, os.ErrNotExist) {
+			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 := os.Remove(s.path(row.ID)); e != nil && !errors.Is(e, os.ErrNotExist) {
+				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
+}