fix: 初验针对修改

internal/storage/local/upload.go

@@ -0,0 +1,190 @@
+package local
+
+import (
+	"context"
+	"crypto/sha256"
+	"encoding/hex"
+	"errors"
+	"fmt"
+	"io"
+	"os"
+	"path/filepath"
+	"strings"
+	"time"
+
+	"git.apinb.com/ops/files/internal/signing"
+	"git.apinb.com/ops/files/internal/storage"
+)
+
+func (c *Client) PrepareUpload(_ context.Context, request storage.UploadRequest) (storage.UploadInstruction, error) {
+	if err := c.ensureCapacity(request.ExpectedSize); err != nil {
+		return storage.UploadInstruction{}, err
+	}
+	signature := signing.SignUpload(c.signingSecret, request)
+	return storage.UploadInstruction{
+		Method: "PUT",
+		URL: fmt.Sprintf("%s/v1/local/uploads/%s?expires=%d&signature=%s",
+			c.externalBaseURL, request.FileID, request.ExpiresAt.Unix(), signature),
+		Headers: map[string]string{
+			"Content-Type":   request.ContentType,
+			"Content-Length": fmt.Sprintf("%d", request.ExpectedSize),
+		},
+		ExpiresAt: request.ExpiresAt,
+	}, nil
+}
+
+func (c *Client) AcquireUpload(ctx context.Context, fileID string) (func() error, error) {
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+	if !validFileID(fileID) {
+		return nil, fmt.Errorf("文件标识无效")
+	}
+	lockPath := filepath.Join(c.locksPath, fileID+".lock")
+	lockFile, err := os.OpenFile(lockPath, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o640)
+	if err != nil {
+		if os.IsExist(err) {
+			return nil, storage.ErrUploadLocked
+		}
+		return nil, fmt.Errorf("获取上传锁: %w", err)
+	}
+	return func() error {
+		closeErr := lockFile.Close()
+		removeErr := os.Remove(lockPath)
+		if os.IsNotExist(removeErr) {
+			removeErr = nil
+		}
+		if closeErr != nil {
+			return closeErr
+		}
+		return removeErr
+	}, nil
+}
+
+func (c *Client) ReceiveUpload(ctx context.Context, request storage.UploadRequest, source io.Reader) (storage.ObjectInfo, error) {
+	if err := c.ensureCapacity(request.ExpectedSize); err != nil {
+		return storage.ObjectInfo{}, err
+	}
+	targetPath, err := c.resolveObjectPath(request.ObjectKey)
+	if err != nil {
+		return storage.ObjectInfo{}, err
+	}
+	if err := c.ensureDirectory(filepath.Dir(targetPath)); err != nil {
+		return storage.ObjectInfo{}, err
+	}
+	if existing, statErr := c.Stat(ctx, request.ObjectKey); statErr == nil {
+		if existing.Size != request.ExpectedSize {
+			return storage.ObjectInfo{}, fmt.Errorf("已存在文件的大小与上传请求不一致")
+		}
+		return existing, nil
+	} else if !errors.Is(statErr, storage.ErrObjectNotFound) {
+		return storage.ObjectInfo{}, statErr
+	}
+	nonce, err := signing.UploadNonce(c.signingSecret, request)
+	if err != nil {
+		return storage.ObjectInfo{}, err
+	}
+	stagingPath := filepath.Join(c.stagingPath, request.FileID+"."+nonce+".part")
+	staging, err := os.OpenFile(stagingPath, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o640)
+	if err != nil {
+		return storage.ObjectInfo{}, fmt.Errorf("创建上传暂存文件: %w", err)
+	}
+	committed := false
+	defer func() {
+		_ = staging.Close()
+		if !committed {
+			_ = os.Remove(stagingPath)
+		}
+	}()
+
+	hasher := sha256.New()
+	written, err := io.Copy(io.MultiWriter(staging, hasher), io.LimitReader(&contextReader{ctx: ctx, reader: source}, request.ExpectedSize+1))
+	if err != nil {
+		return storage.ObjectInfo{}, fmt.Errorf("写入上传暂存文件: %w", err)
+	}
+	if written != request.ExpectedSize {
+		return storage.ObjectInfo{}, fmt.Errorf("上传文件大小不匹配: 期望 %d，实际 %d", request.ExpectedSize, written)
+	}
+	if err := staging.Sync(); err != nil {
+		return storage.ObjectInfo{}, fmt.Errorf("同步上传暂存文件: %w", err)
+	}
+	if err := staging.Close(); err != nil {
+		return storage.ObjectInfo{}, fmt.Errorf("关闭上传暂存文件: %w", err)
+	}
+	if err := os.Link(stagingPath, targetPath); err != nil {
+		return storage.ObjectInfo{}, fmt.Errorf("提交上传文件: %w", err)
+	}
+	if err := os.Chmod(targetPath, 0o640); err != nil {
+		_ = os.Remove(targetPath)
+		return storage.ObjectInfo{}, fmt.Errorf("设置上传文件权限: %w", err)
+	}
+	if err := os.Remove(stagingPath); err != nil {
+		_ = os.Remove(targetPath)
+		return storage.ObjectInfo{}, fmt.Errorf("清理上传暂存文件: %w", err)
+	}
+	if err := syncDirectory(filepath.Dir(targetPath)); err != nil {
+		_ = os.Remove(targetPath)
+		return storage.ObjectInfo{}, err
+	}
+	committed = true
+	return storage.ObjectInfo{
+		Size:         written,
+		StorageETag:  "sha256:" + hex.EncodeToString(hasher.Sum(nil)),
+		LastModified: time.Now().UTC(),
+	}, nil
+}
+
+func (c *Client) DiscardUpload(_ context.Context, fileID string) error {
+	if !validFileID(fileID) {
+		return fmt.Errorf("文件标识无效")
+	}
+	entries, err := os.ReadDir(c.stagingPath)
+	if err != nil {
+		return err
+	}
+	prefix := fileID + "."
+	for _, entry := range entries {
+		if !entry.IsDir() && strings.HasPrefix(entry.Name(), prefix) && strings.HasSuffix(entry.Name(), ".part") {
+			if err := os.Remove(filepath.Join(c.stagingPath, entry.Name())); err != nil && !os.IsNotExist(err) {
+				return err
+			}
+		}
+	}
+	return nil
+}
+
+type contextReader struct {
+	ctx    context.Context
+	reader io.Reader
+}
+
+func (r *contextReader) Read(buffer []byte) (int, error) {
+	if err := r.ctx.Err(); err != nil {
+		return 0, err
+	}
+	return r.reader.Read(buffer)
+}
+
+func validFileID(fileID string) bool {
+	if fileID == "" {
+		return false
+	}
+	for _, character := range fileID {
+		if (character < '0' || character > '9') && (character < 'A' || character > 'Z') && (character < 'a' || character > 'z') && character != '-' {
+			return false
+		}
+	}
+	return true
+}
+
+func syncDirectory(path string) error {
+	directory, err := os.Open(path)
+	if err != nil {
+		return fmt.Errorf("打开目录进行同步: %w", err)
+	}
+	defer directory.Close()
+	if err := directory.Sync(); err != nil {
+		return fmt.Errorf("同步文件目录: %w", err)
+	}
+	return nil
+}