files/internal/storage/local/upload.go

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
}