OpenList/drivers/s3/util.go

package s3

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"path"
	"strings"

	"github.com/OpenListTeam/OpenList/v4/internal/model"
	"github.com/OpenListTeam/OpenList/v4/internal/op"
	"github.com/OpenListTeam/OpenList/v4/pkg/utils"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/request"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	log "github.com/sirupsen/logrus"
)

const (
	maxCopyObjectSize   int64 = 5 * 1000 * 1000 * 1000
	defaultCopyPartSize int64 = 100 * 1024 * 1024
	maxCopyPartSize     int64 = 5 * 1024 * 1024 * 1024
	maxCopyParts        int64 = 10000
)

// do others that not defined in Driver interface

func (d *S3) initSession() error {
	var err error
	accessKeyID, secretAccessKey, sessionToken := d.AccessKeyID, d.SecretAccessKey, d.SessionToken
	if d.config.Name == "Doge" {
		credentialsTmp, err := getCredentials(d.AccessKeyID, d.SecretAccessKey)
		if err != nil {
			return err
		}
		accessKeyID, secretAccessKey, sessionToken = credentialsTmp.AccessKeyId, credentialsTmp.SecretAccessKey, credentialsTmp.SessionToken
	}
	cfg := &aws.Config{
		Credentials:      credentials.NewStaticCredentials(accessKeyID, secretAccessKey, sessionToken),
		Region:           &d.Region,
		Endpoint:         &d.Endpoint,
		S3ForcePathStyle: aws.Bool(d.ForcePathStyle),
	}
	d.Session, err = session.NewSession(cfg)
	return err
}

const (
	ClientTypeNormal = iota
	ClientTypeLink
	ClientTypeDirectUpload
)

func (d *S3) getClient(clientType int) *s3.S3 {
	client := s3.New(d.Session)
	if d.UserAgent != "" {
		client.Handlers.Build.PushBack(func(r *request.Request) {
			r.HTTPRequest.Header.Set("User-Agent", d.UserAgent)
		})
	}
	if clientType == ClientTypeLink && d.CustomHost != "" {
		client.Handlers.Build.PushBack(func(r *request.Request) {
			if r.HTTPRequest.Method != http.MethodGet {
				return
			}
			//判断CustomHost是否以http://或https://开头
			split := strings.SplitN(d.CustomHost, "://", 2)
			if utils.SliceContains([]string{"http", "https"}, split[0]) {
				r.HTTPRequest.URL.Scheme = split[0]
				r.HTTPRequest.URL.Host = split[1]
			} else {
				r.HTTPRequest.URL.Host = d.CustomHost
			}
		})
	}
	if clientType == ClientTypeDirectUpload && d.DirectUploadHost != "" {
		client.Handlers.Build.PushBack(func(r *request.Request) {
			if r.HTTPRequest.Method != http.MethodPut {
				return
			}
			split := strings.SplitN(d.DirectUploadHost, "://", 2)
			if utils.SliceContains([]string{"http", "https"}, split[0]) {
				r.HTTPRequest.URL.Scheme = split[0]
				r.HTTPRequest.URL.Host = split[1]
			} else {
				r.HTTPRequest.URL.Host = d.DirectUploadHost
			}
		})
	}
	return client
}

func getKey(path string, dir bool) string {
	path = strings.TrimPrefix(path, "/")
	if path != "" && dir {
		path += "/"
	}
	return path
}

var defaultPlaceholderName = ".openlist"

func getPlaceholderName(placeholder string) string {
	if placeholder == "" {
		return defaultPlaceholderName
	}
	return placeholder
}

func (d *S3) listV1(dirPath string, args model.ListArgs) ([]model.Obj, error) {
	prefix := getKey(dirPath, true)
	log.Debugf("list: %s", prefix)
	files := make([]model.Obj, 0)
	marker := ""
	for {
		input := &s3.ListObjectsInput{
			Bucket:    &d.Bucket,
			Marker:    &marker,
			Prefix:    &prefix,
			Delimiter: aws.String("/"),
		}
		listObjectsResult, err := d.client.ListObjects(input)
		if err != nil {
			return nil, err
		}
		for _, object := range listObjectsResult.CommonPrefixes {
			name := path.Base(strings.Trim(*object.Prefix, "/"))
			file := model.Object{
				Path:     path.Join(dirPath, name),
				Name:     name,
				Modified: d.Modified,
				IsFolder: true,
			}
			files = append(files, &file)
		}
		for _, object := range listObjectsResult.Contents {
			name := path.Base(*object.Key)
			if !args.S3ShowPlaceholder && (name == getPlaceholderName(d.Placeholder) || name == d.Placeholder) {
				continue
			}
			file := model.Object{
				Path:     path.Join(dirPath, name),
				Name:     name,
				Size:     *object.Size,
				Modified: *object.LastModified,
			}
			files = append(files, &file)
		}
		if listObjectsResult.IsTruncated == nil {
			return nil, errors.New("IsTruncated nil")
		}
		if *listObjectsResult.IsTruncated {
			marker = *listObjectsResult.NextMarker
		} else {
			break
		}
	}
	return files, nil
}

func (d *S3) listV2(dirPath string, args model.ListArgs) ([]model.Obj, error) {
	prefix := getKey(dirPath, true)
	files := make([]model.Obj, 0)
	var continuationToken, startAfter *string
	for {
		input := &s3.ListObjectsV2Input{
			Bucket:            &d.Bucket,
			ContinuationToken: continuationToken,
			Prefix:            &prefix,
			Delimiter:         aws.String("/"),
			StartAfter:        startAfter,
		}
		listObjectsResult, err := d.client.ListObjectsV2(input)
		if err != nil {
			return nil, err
		}
		log.Debugf("resp: %+v", listObjectsResult)
		for _, object := range listObjectsResult.CommonPrefixes {
			name := path.Base(strings.Trim(*object.Prefix, "/"))
			file := model.Object{
				Path:     path.Join(dirPath, name),
				Name:     name,
				Modified: d.Modified,
				IsFolder: true,
			}
			files = append(files, &file)
		}
		for _, object := range listObjectsResult.Contents {
			if strings.HasSuffix(*object.Key, "/") {
				continue
			}
			name := path.Base(*object.Key)
			if !args.S3ShowPlaceholder && (name == getPlaceholderName(d.Placeholder) || name == d.Placeholder) {
				continue
			}
			file := model.Object{
				Path:     path.Join(dirPath, name),
				Name:     name,
				Size:     *object.Size,
				Modified: *object.LastModified,
			}
			files = append(files, &file)
		}
		if !aws.BoolValue(listObjectsResult.IsTruncated) {
			break
		}
		if listObjectsResult.NextContinuationToken != nil {
			continuationToken = listObjectsResult.NextContinuationToken
			continue
		}
		if len(listObjectsResult.Contents) == 0 {
			break
		}
		startAfter = listObjectsResult.Contents[len(listObjectsResult.Contents)-1].Key
	}
	return files, nil
}

func (d *S3) copy(ctx context.Context, src string, dst string, size int64, isDir bool) error {
	if isDir {
		return d.copyDir(ctx, src, dst)
	}
	return d.copyFile(ctx, src, dst, size)
}

func (d *S3) copyFile(ctx context.Context, src string, dst string, size int64) error {
	srcKey := getKey(src, false)
	dstKey := getKey(dst, false)
	encodedKey := strings.ReplaceAll(url.PathEscape(d.Bucket+"/"+srcKey), "+", "%2B")
	if size > maxCopyObjectSize {
		return d.copyFileMultipart(ctx, srcKey, dstKey, encodedKey, size)
	}
	input := &s3.CopyObjectInput{
		Bucket:     &d.Bucket,
		CopySource: aws.String(encodedKey),
		Key:        &dstKey,
	}
	_, err := d.client.CopyObjectWithContext(ctx, input)
	return err
}

func (d *S3) copyFileMultipart(ctx context.Context, srcKey, dstKey, encodedKey string, size int64) (err error) {
	head, err := d.client.HeadObjectWithContext(ctx, &s3.HeadObjectInput{
		Bucket: &d.Bucket,
		Key:    &srcKey,
	})
	if err != nil {
		return err
	}
	if head.ContentLength != nil {
		size = *head.ContentLength
	}
	partSize, err := getCopyPartSize(size)
	if err != nil {
		return err
	}
	createInput := &s3.CreateMultipartUploadInput{
		Bucket:                  &d.Bucket,
		Key:                     &dstKey,
		CacheControl:            head.CacheControl,
		ContentDisposition:      head.ContentDisposition,
		ContentEncoding:         head.ContentEncoding,
		ContentLanguage:         head.ContentLanguage,
		ContentType:             head.ContentType,
		Metadata:                head.Metadata,
		WebsiteRedirectLocation: head.WebsiteRedirectLocation,
	}
	if head.Expires != nil {
		if expires, parseErr := http.ParseTime(*head.Expires); parseErr == nil {
			createInput.Expires = &expires
		}
	}
	created, err := d.client.CreateMultipartUploadWithContext(ctx, createInput)
	if err != nil {
		return err
	}
	uploadID := aws.StringValue(created.UploadId)
	if uploadID == "" {
		return errors.New("create multipart upload returned an empty upload ID")
	}
	completed := false
	defer func() {
		if completed {
			return
		}
		_, abortErr := d.client.AbortMultipartUploadWithContext(context.WithoutCancel(ctx), &s3.AbortMultipartUploadInput{
			Bucket:   &d.Bucket,
			Key:      &dstKey,
			UploadId: &uploadID,
		})
		if abortErr != nil {
			err = errors.Join(err, fmt.Errorf("failed to abort multipart copy: %w", abortErr))
		}
	}()

	parts := make([]*s3.CompletedPart, 0, (size+partSize-1)/partSize)
	for start, partNumber := int64(0), int64(1); start < size; start, partNumber = start+partSize, partNumber+1 {
		end := min(start+partSize, size) - 1
		copied, copyErr := d.client.UploadPartCopyWithContext(ctx, &s3.UploadPartCopyInput{
			Bucket:          &d.Bucket,
			CopySource:      &encodedKey,
			CopySourceRange: aws.String(fmt.Sprintf("bytes=%d-%d", start, end)),
			Key:             &dstKey,
			PartNumber:      &partNumber,
			UploadId:        &uploadID,
		})
		if copyErr != nil {
			return copyErr
		}
		if copied.CopyPartResult == nil || aws.StringValue(copied.CopyPartResult.ETag) == "" {
			return fmt.Errorf("multipart copy part %d returned an empty ETag", partNumber)
		}
		parts = append(parts, &s3.CompletedPart{
			ETag:       copied.CopyPartResult.ETag,
			PartNumber: &partNumber,
		})
	}

	_, err = d.client.CompleteMultipartUploadWithContext(ctx, &s3.CompleteMultipartUploadInput{
		Bucket:   &d.Bucket,
		Key:      &dstKey,
		UploadId: &uploadID,
		MultipartUpload: &s3.CompletedMultipartUpload{
			Parts: parts,
		},
	})
	if err != nil {
		return err
	}
	completed = true
	return nil
}

func getCopyPartSize(size int64) (int64, error) {
	partSize := max(defaultCopyPartSize, (size-1)/maxCopyParts+1)
	if partSize > maxCopyPartSize {
		return 0, fmt.Errorf("object size %d exceeds multipart copy limit", size)
	}
	return partSize, nil
}

func (d *S3) copyDir(ctx context.Context, src string, dst string) error {
	objs, err := op.List(ctx, d, src, model.ListArgs{S3ShowPlaceholder: true})
	if err != nil {
		return err
	}
	for _, obj := range objs {
		cSrc := path.Join(src, obj.GetName())
		cDst := path.Join(dst, obj.GetName())
		if obj.IsDir() {
			err = d.copyDir(ctx, cSrc, cDst)
		} else {
			err = d.copyFile(ctx, cSrc, cDst, obj.GetSize())
		}
		if err != nil {
			return err
		}
	}
	return nil
}

func (d *S3) removeDir(ctx context.Context, src string) error {
	objs, err := op.List(ctx, d, src, model.ListArgs{})
	if err != nil {
		return err
	}
	for _, obj := range objs {
		cSrc := path.Join(src, obj.GetName())
		if obj.IsDir() {
			err = d.removeDir(ctx, cSrc)
		} else {
			err = d.removeFile(ctx, cSrc)
		}
		if err != nil {
			return err
		}
	}
	_ = d.removeFile(ctx, path.Join(src, getPlaceholderName(d.Placeholder)))
	_ = d.removeFile(ctx, path.Join(src, d.Placeholder))
	return nil
}

func (d *S3) removeFile(ctx context.Context, src string) error {
	key := getKey(src, false)
	input := &s3.DeleteObjectInput{
		Bucket: &d.Bucket,
		Key:    &key,
	}
	_, err := d.client.DeleteObjectWithContext(ctx, input)
	return err
}