JCS-pub/common/pkgs/storage/s3/multipart_upload.go

package s3

import (
	"context"
	"crypto/sha256"
	"io"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	s3types "github.com/aws/aws-sdk-go-v2/service/s3/types"
	"gitlink.org.cn/cloudream/common/utils/io2"
	"gitlink.org.cn/cloudream/common/utils/os2"
	"gitlink.org.cn/cloudream/common/utils/sort2"
	stgtypes "gitlink.org.cn/cloudream/jcs-pub/common/pkgs/storage/types"
	jcstypes "gitlink.org.cn/cloudream/jcs-pub/common/types"
)

type Multiparter struct {
	detail *jcstypes.UserSpaceDetail
	feat   *jcstypes.MultipartUploadFeature
	bucket string
	cli    *s3.Client
}

func NewMultiparter(detail *jcstypes.UserSpaceDetail, feat *jcstypes.MultipartUploadFeature, bkt string, cli *s3.Client) *Multiparter {
	return &Multiparter{
		detail: detail,
		feat:   feat,
		bucket: bkt,
		cli:    cli,
	}
}

func (*Multiparter) MinPartSize() int64 {
	return 5 * 1024 * 1024 // 5MB
}

func (*Multiparter) MaxPartSize() int64 {
	return 5 * 1024 * 1024 * 1024 // 5GB
}

func (m *Multiparter) Initiate(ctx context.Context) (stgtypes.MultipartTask, error) {
	tempFileName := os2.GenerateRandomFileName(10)
	tempDir := m.detail.UserSpace.WorkingDir.Clone()
	tempDir.Push(stgtypes.TempWorkingDir)
	tempFilePath := tempDir.Clone()
	tempFilePath.Push(tempFileName)

	resp, err := m.cli.CreateMultipartUpload(ctx, &s3.CreateMultipartUploadInput{
		Bucket:            aws.String(m.bucket),
		Key:               aws.String(tempFilePath.String()),
		ChecksumAlgorithm: s3types.ChecksumAlgorithmSha256,
	})
	if err != nil {
		return nil, err
	}

	return &MultipartTask{
		cli:          m.cli,
		bucket:       m.bucket,
		tempDir:      tempDir,
		tempFileName: tempFileName,
		tempFilePath: tempFilePath,
		uploadID:     *resp.UploadId,
	}, nil
}

func (m *Multiparter) UploadPart(ctx context.Context, init stgtypes.MultipartInitState, partSize int64, partNumber int, stream io.Reader) (stgtypes.UploadedPartInfo, error) {
	hashStr := io2.NewReadHasher(sha256.New(), stream)
	resp, err := m.cli.UploadPart(ctx, &s3.UploadPartInput{
		Bucket:     aws.String(init.Bucket),
		Key:        aws.String(init.Key),
		UploadId:   aws.String(init.UploadID),
		PartNumber: aws.Int32(int32(partNumber)),
		Body:       hashStr,
	})
	if err != nil {
		return stgtypes.UploadedPartInfo{}, err
	}

	return stgtypes.UploadedPartInfo{
		ETag:       *resp.ETag,
		PartNumber: partNumber,
		PartHash:   hashStr.Sum(),
	}, nil
}

type MultipartTask struct {
	cli          *s3.Client
	bucket       string
	tempDir      jcstypes.JPath
	tempFileName string
	tempFilePath jcstypes.JPath
	uploadID     string
}

func (i *MultipartTask) InitState() stgtypes.MultipartInitState {
	return stgtypes.MultipartInitState{
		UploadID: i.uploadID,
		Bucket:   i.bucket,
		Key:      i.tempFilePath.String(),
	}
}

func (i *MultipartTask) JoinParts(ctx context.Context, parts []stgtypes.UploadedPartInfo) (stgtypes.FileInfo, error) {
	parts = sort2.Sort(parts, func(l, r stgtypes.UploadedPartInfo) int {
		return l.PartNumber - r.PartNumber
	})

	s3Parts := make([]s3types.CompletedPart, len(parts))
	for i, part := range parts {
		s3Parts[i] = s3types.CompletedPart{
			ETag:       aws.String(part.ETag),
			PartNumber: aws.Int32(int32(part.PartNumber)),
		}
	}
	partHashes := make([][]byte, len(parts))
	for i, part := range parts {
		partHashes[i] = part.PartHash
	}

	_, err := i.cli.CompleteMultipartUpload(ctx, &s3.CompleteMultipartUploadInput{
		Bucket:   aws.String(i.bucket),
		Key:      aws.String(i.tempFilePath.String()),
		UploadId: aws.String(i.uploadID),
		MultipartUpload: &s3types.CompletedMultipartUpload{
			Parts: s3Parts,
		},
	})
	if err != nil {
		return stgtypes.FileInfo{}, err
	}

	headResp, err := i.cli.HeadObject(ctx, &s3.HeadObjectInput{
		Bucket: aws.String(i.bucket),
		Key:    aws.String(i.tempFilePath.String()),
	})
	if err != nil {
		return stgtypes.FileInfo{}, err
	}

	hash := jcstypes.CalculateCompositeHash(partHashes)

	return stgtypes.FileInfo{
		Path: i.tempFilePath,
		Size: *headResp.ContentLength,
		Hash: hash,
	}, nil

}

func (i *MultipartTask) Close() {
	i.cli.AbortMultipartUpload(context.Background(), &s3.AbortMultipartUploadInput{
		Bucket:   aws.String(i.bucket),
		Key:      aws.String(i.tempFilePath.String()),
		UploadId: aws.String(i.uploadID),
	})
}