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\" 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(¤t).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(¤t).Updates(map[string]any{"state": "expired", "result": types.JSON(`{}`), "updated_at": time.Now()}).Error }); e != nil { return e } } return nil }