s3store.go

package dstore

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"os"
	"path"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	awsconfig "github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/credentials"
	"github.com/aws/aws-sdk-go-v2/feature/s3/manager"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/aws/aws-sdk-go-v2/service/s3/types"
	"go.uber.org/zap"
	"golang.org/x/net/http2"
)

type s3ReadCloser struct {
	outer    io.ReadCloser
	httpBody io.ReadCloser
}

func (s *s3ReadCloser) Read(p []byte) (int, error) {
	return s.outer.Read(p)
}

func (s *s3ReadCloser) Close() error {
	// Drain the raw HTTP body BEFORE closing the outer reader chain. If outer
	// closes httpBody first (which it does when there is no compression layer),
	// the subsequent drain becomes a no-op and the connection is never returned
	// to the pool.
	io.Copy(io.Discard, s.httpBody)
	return s.outer.Close()
}

var retryS3PushLocalFilesDelay time.Duration
var s3ReadAttempts = 1
var bufferedS3Read bool
var s3MaxIdleConns = 100
var s3MaxIdleConnsPerHost = 10
var s3IdleConnTimeout = 90 * time.Second

var sharedTransport http.RoundTripper
var sharedTransportOnce sync.Once

func getSharedTransport() http.RoundTripper {
	sharedTransportOnce.Do(func() {
		t := &http.Transport{
			ForceAttemptHTTP2:     true,
			MaxIdleConns:          s3MaxIdleConns,
			MaxIdleConnsPerHost:   s3MaxIdleConnsPerHost,
			IdleConnTimeout:       s3IdleConnTimeout,
			ResponseHeaderTimeout: 30 * time.Second,
			TLSHandshakeTimeout:   10 * time.Second,
			ExpectContinueTimeout: 1 * time.Second,
			DialContext: (&net.Dialer{
				Timeout:   30 * time.Second,
				KeepAlive: 30 * time.Second,
			}).DialContext,
		}

		if t2, err := http2.ConfigureTransports(t); err == nil {
			t2.ReadIdleTimeout = 31 * time.Second
			t2.PingTimeout = 15 * time.Second
		}

		sharedTransport = t
	})
	return sharedTransport
}

func init() {
	retry := os.Getenv("DSTORE_S3_RETRY_PUSH_DELAY")
	if retry != "" {
		retryS3PushLocalFilesDelay, _ = time.ParseDuration(retry)
	}
	if os.Getenv("DSTORE_S3_BUFFERED_READ") == "true" {
		bufferedS3Read = true
	}

	readAttempts := os.Getenv("DSTORE_S3_READ_ATTEMPTS")
	if readAttempts != "" {
		attempts, _ := strconv.ParseUint(readAttempts, 10, 64)
		if attempts > 0 {
			s3ReadAttempts = int(attempts)
		}
	}

	if v := os.Getenv("DSTORE_S3_MAX_IDLE_CONNS"); v != "" {
		if n, err := strconv.Atoi(v); err == nil && n > 0 {
			s3MaxIdleConns = n
		}
	}
	if v := os.Getenv("DSTORE_S3_MAX_IDLE_CONNS_PER_HOST"); v != "" {
		if n, err := strconv.Atoi(v); err == nil && n > 0 {
			s3MaxIdleConnsPerHost = n
		}
	}
	if v := os.Getenv("DSTORE_S3_IDLE_CONN_TIMEOUT"); v != "" {
		if d, err := time.ParseDuration(v); err == nil {
			s3IdleConnTimeout = d
		}
	}

	zlog.Info("S3 storage configured",
		zap.Bool("buffered_read", bufferedS3Read),
		zap.Int("read_attempts", s3ReadAttempts),
		zap.Duration("retry_push_local_files_delay", retryS3PushLocalFilesDelay),
		zap.Int("max_idle_conns", s3MaxIdleConns),
		zap.Int("max_idle_conns_per_host", s3MaxIdleConnsPerHost),
		zap.Duration("idle_conn_timeout", s3IdleConnTimeout),
	)
}

type S3Store struct {
	baseURL *url.URL

	bucket       string
	path         string
	storageClass string
	client       *s3.Client
	uploader     *manager.Uploader
	downloader   *manager.Downloader
	context      context.Context

	*commonStore
}

func NewS3Store(baseURL *url.URL, extension, compressionType string, overwrite bool, opts ...Option) (*S3Store, error) {
	ctx := context.Background()
	return newS3StoreContext(ctx, baseURL, extension, compressionType, overwrite, opts...)
}

func newS3StoreContext(ctx context.Context, baseURL *url.URL, extension, compressionType string, overwrite bool, opts ...Option) (*S3Store, error) {
	conf := config{}
	for _, opt := range opts {
		opt.apply(&conf)
	}

	common := &commonStore{
		compressionType:           compressionType,
		extension:                 extension,
		overwrite:                 overwrite,
		uncompressedReadCallback:  conf.uncompressedReadCallback,
		compressedReadCallback:    conf.compressedReadCallback,
		uncompressedWriteCallback: conf.uncompressedWriteCallback,
		compressedWriteCallback:   conf.compressedWriteCallback,
	}

	s := &S3Store{
		baseURL:     baseURL,
		commonStore: common,
		context:     ctx,
	}

	awsConfig, bucket, path, storageClass, err := ParseS3URL(baseURL)
	if err != nil {
		return nil, fmt.Errorf("invalid s3 url: %w", err)
	}

	awsConfig = append(awsConfig, awsconfig.WithHTTPClient(&http.Client{
		Transport: getSharedTransport(),
	}))

	cfg, err := awsconfig.LoadDefaultConfig(ctx, awsConfig...)
	if err != nil {
		return nil, fmt.Errorf("error loading AWS config: %w", err)
	}

	s.client = s3.NewFromConfig(cfg)
	s.uploader = manager.NewUploader(s.client)
	s.downloader = manager.NewDownloader(s.client)
	s.bucket = bucket
	s.path = path
	s.storageClass = storageClass

	return s, nil
}

func (s *S3Store) Clone(ctx context.Context, opts ...Option) (Store, error) {
	return newS3StoreContext(ctx, s.baseURL, s.extension, s.compressionType, s.overwrite, opts...)
}

func (s *S3Store) SubStore(subFolder string) (Store, error) {
	url, err := url.Parse(s.baseURL.String())
	if err != nil {
		return nil, fmt.Errorf("s3 store parsing base url: %w", err)
	}
	url.Path = path.Join(url.Path, subFolder)
	newPath := path.Join(s.path, subFolder)
	return &S3Store{
		baseURL:      url,
		commonStore:  s.commonStore,
		client:       s.client,
		uploader:     s.uploader,
		downloader:   s.downloader,
		bucket:       s.bucket,
		storageClass: s.storageClass,
		path:         newPath,
	}, nil
}

func ParseS3URL(s3URL *url.URL) (configOptions []func(*awsconfig.LoadOptions) error, bucket, path, storageClass string, err error) {
	region := s3URL.Query().Get("region")
	if region == "" {
		return nil, "", "", "", fmt.Errorf("specify s3 bucket like: s3://bucket/path?region=us-east-1")
	}

	configOptions = append(configOptions, awsconfig.WithRegion(region))

	hasEndpoint := hasCustomEndpoint(s3URL)
	if hasEndpoint {
		endpoint := s3URL.Host
		if s3URL.Query().Get("insecure") != "" {
			endpoint = "http://" + endpoint
		} else {
			endpoint = "https://" + endpoint
		}

		configOptions = append(configOptions, awsconfig.WithEndpointResolverWithOptions(
			aws.EndpointResolverWithOptionsFunc(func(service, region string, options ...interface{}) (aws.Endpoint, error) {
				return aws.Endpoint{
					URL:               endpoint,
					HostnameImmutable: true,
				}, nil
			}),
		))

		pathParts := strings.Split(strings.TrimLeft(s3URL.Path, "/"), "/")
		bucket = pathParts[0]
		path = strings.Replace(s3URL.Path, bucket, "", 1)
	} else {
		bucket = s3URL.Hostname()
		path = s3URL.Path
	}

	accessKeyID := s3URL.Query().Get("access_key_id")
	secretAccessKey := s3URL.Query().Get("secret_access_key")
	if accessKeyID != "" && secretAccessKey != "" {
		configOptions = append(configOptions, awsconfig.WithCredentialsProvider(
			credentials.NewStaticCredentialsProvider(accessKeyID, secretAccessKey, ""),
		))
	}

	return configOptions, bucket, strings.Trim(path, "/"), getStorageClass(s3URL.Query()), nil
}

func hasCustomEndpoint(s3URL *url.URL) bool {
	// As soon as there is a port in the url, we are sure that's it's the
	// hostname that should be configured, so move along
	if s3URL.Port() != "" {
		return true
	}

	// If there is no `.` in the hostname, we assume it's a bucket. It could still be
	// problematic for `localhost`, we are expecting people to use `:<port>` to go
	// in the condition above.
	host := s3URL.Hostname()
	if !strings.Contains(host, ".") {
		return false
	}

	// Otherwise, by default we assume it's an hostname followed by the bucket. If
	// operator really intent to use the bucket directly an it contains dot, the
	// query parameter `infer_aws_endpoint=true` can be used to tell the store
	// implementation that the hostname is the actual bucket
	return s3URL.Query().Get("infer_aws_endpoint") == ""
}

func getStorageClass(q url.Values) string {
	if v := q.Get("storage_class"); v != "" {
		return v
	}
	if v := q.Get("storageClass"); v != "" {
		zlog.Warn("query parameter 'storageClass' is deprecated, use 'storage_class' instead")
		return v
	}
	return ""
}

func (s *S3Store) BaseURL() *url.URL {
	return s.baseURL
}

func (s *S3Store) ObjectPath(name string) string {
	return path.Join(s.path, s.pathWithExt(name))
}

func (s *S3Store) ObjectURL(name string) string {
	return fmt.Sprintf("%s/%s", strings.TrimRight(s.baseURL.String(), "/"), strings.TrimLeft(s.pathWithExt(name), "/"))
}

func (s *S3Store) WriteObject(ctx context.Context, base string, f io.Reader, metadataKeyValues ...string) (err error) {
	ctx = withFileName(ctx, base)
	ctx = withStoreType(ctx, "s3store")
	ctx = withLogger(ctx, zlog, tracer)

	// Parse metadataKeyValues array
	if len(metadataKeyValues)%2 != 0 {
		return fmt.Errorf("metadataKeyValues must have an even number of strings (key-value pairs), got %d", len(metadataKeyValues))
	}

	objPath := s.ObjectPath(base)

	exists, err := s.FileExists(ctx, base)
	if err != nil {
		return err
	}

	if !s.overwrite && exists {
		// We silently ignore when we ask not to overwrite
		return nil
	}

	pr, pw := io.Pipe()
	writeDone := make(chan error, 1)
	ctx, cancel := context.WithCancel(ctx)

	wg := sync.WaitGroup{}
	wg.Add(1)

	go func(ctx context.Context) {
		defer wg.Done()

		err := s.compressedCopy(ctx, pw, f)
		writeDone <- err
		pw.Close() // required to allow the uploader to complete

		if err != nil {
			cancel()
		}
	}(ctx)

	uploadInput := &s3.PutObjectInput{
		Bucket: aws.String(s.bucket),
		Key:    aws.String(objPath),
		Body:   pr,
	}

	if s.storageClass != "" {
		uploadInput.StorageClass = types.StorageClass(s.storageClass)
	}

	// Add metadata if provided
	if len(metadataKeyValues) > 0 {
		metadata := make(map[string]string)
		for i := 0; i < len(metadataKeyValues); i += 2 {
			key := metadataKeyValues[i]
			value := metadataKeyValues[i+1]
			metadata[key] = value
		}
		uploadInput.Metadata = metadata
	}

	_, err = s.uploader.Upload(ctx, uploadInput)
	if err != nil {
		select {
		case err2 := <-writeDone:
			if err2 != nil {
				return fmt.Errorf("writing through pipe: %w", err2)
			}
		default:
			// error was generated in the Upload (s3 or context timeout), compressedCopy is not finished,
			// we make it fail. double closing is safe here
			pw.Close()
		}
		return fmt.Errorf("uploading to S3 through manager: %w", err)
	}

	wg.Wait()

	return nil
}

func (s *S3Store) CopyObject(ctx context.Context, src, dest string) error {
	// TODO optimize this
	reader, err := s.OpenObject(ctx, src)
	if err != nil {
		return err
	}
	defer reader.Close()

	return s.WriteObject(ctx, dest, reader)
}

func (s *S3Store) FileExists(ctx context.Context, base string) (bool, error) {
	path := s.ObjectPath(base)

	_, err := s.client.HeadObject(ctx, &s3.HeadObjectInput{
		Bucket: aws.String(s.bucket),
		Key:    aws.String(path),
	})
	if err != nil {
		var notFound *types.NotFound
		var noSuchKey *types.NoSuchKey
		if errors.As(err, &notFound) || errors.As(err, &noSuchKey) {
			return false, nil
		}
		return false, err
	}

	return true, nil
}

func (s *S3Store) ObjectAttributes(ctx context.Context, base string) (*ObjectAttributes, error) {
	path := s.ObjectPath(base)

	output, err := s.client.HeadObject(ctx, &s3.HeadObjectInput{
		Bucket: aws.String(s.bucket),
		Key:    aws.String(path),
	})
	if err != nil {
		return nil, err
	}

	return &ObjectAttributes{
		LastModified: *output.LastModified,
		Size:         *output.ContentLength,
		Metadata:     output.Metadata,
	}, nil
}

func (s *S3Store) SetMetadata(ctx context.Context, base string, metadata map[string]string) error {
	path := s.ObjectPath(base)

	// Use CopyObject to update metadata (copy to itself with new metadata)
	copySource := fmt.Sprintf("%s/%s", s.bucket, path)
	_, err := s.client.CopyObject(ctx, &s3.CopyObjectInput{
		Bucket:            aws.String(s.bucket),
		Key:               aws.String(path),
		CopySource:        aws.String(copySource),
		Metadata:          metadata,
		MetadataDirective: types.MetadataDirectiveReplace,
	})

	return err
}

func (s *S3Store) OpenObject(ctx context.Context, name string) (out io.ReadCloser, err error) {
	ctx = withStoreType(ctx, "s3store")
	ctx = withLogger(ctx, zlog, tracer)

	path := s.ObjectPath(name)
	ctx = withFileName(ctx, path)

	if tracer.Enabled() {
		zlog.Debug("opening dstore file", zap.String("path", path))
	}

	for i := 0; i < s3ReadAttempts; i++ {
		if i > 0 { // small wait on retry
			zlog.Debug("got an error on s3 OpenObject, retrying",
				zap.Error(err),
				zap.Int("attempt", i),
				zap.Int("max_attempts", s3ReadAttempts),
				zap.String("name", name),
				zap.String("path", path),
			)
			time.Sleep(500 * time.Millisecond)
		}
		var reader *s3.GetObjectOutput
		reader, err = s.client.GetObject(ctx, &s3.GetObjectInput{
			Bucket: aws.String(s.bucket),
			Key:    aws.String(path),
		})
		if err != nil {
			var noSuchBucket *types.NoSuchBucket
			var noSuchKey *types.NoSuchKey
			if errors.As(err, &noSuchBucket) {
				err = fmt.Errorf("s3 bucket %s does not exist", s.bucket)
			} else if errors.As(err, &noSuchKey) {
				err = ErrNotFound
			}
			continue
		}
		if bufferedS3Read {
			var data []byte
			data, err = io.ReadAll(reader.Body)
			if err != nil {
				continue
			}
			if err = reader.Body.Close(); err != nil {
				continue
			}
			out, err = s.uncompressedReader(ctx, io.NopCloser(bytes.NewReader(data)))
		} else {
			httpBody := reader.Body
			out, err = s.uncompressedReader(ctx, httpBody)
			if err == nil {
				out = &s3ReadCloser{outer: out, httpBody: httpBody}
			}
		}
		if tracer.Enabled() {
			out = wrapReadCloser(out, func() {
				zlog.Debug("closing dstore file", zap.String("path", path))
			})
		}
		return out, err
	}
	return nil, fmt.Errorf("s3 open object (%d attempts, buffered_read: %v): %w", s3ReadAttempts, bufferedS3Read, err)
}

func (s *S3Store) WalkFrom(ctx context.Context, prefix, startingPoint string, f func(filename string) (err error)) error {
	return s.WalkFromTo(ctx, prefix, startingPoint, "", f)
}

func (s *S3Store) WalkFromTo(ctx context.Context, prefix, startingPoint, exclusiveEndPoint string, f func(filename string) (err error)) error {
	targetPrefix := s.path
	if targetPrefix != "" {
		targetPrefix += "/"
	}
	if prefix != "" {
		targetPrefix = filepath.Join(targetPrefix, prefix)
		if prefix[len(prefix)-1:] == "/" {
			targetPrefix += "/"
		}
	}

	input := &s3.ListObjectsV2Input{
		Bucket: aws.String(s.bucket),
		Prefix: aws.String(targetPrefix),
	}

	if startingPoint != "" {
		if !strings.HasPrefix(startingPoint, prefix) {
			return fmt.Errorf("starting point %q must start with prefix %q", startingPoint, prefix)
		}

		// "startingPoint" is known to start with "prefix" (checked above), but our the prefix received do
		// not contain the "baseURL" which is required because it contains the "path" of the store. So we remove the
		// "original prefix" from the "startingPoint" and append it to the real "final" prefix instead.
		relativeStartingPoint := strings.TrimPrefix(startingPoint, prefix)

		// to match 'helloworld.html' by using startAfter, we use 'helloworld.htm' (and we filter again in the walk function  to filter out 'helloworld.htm0')
		if len(relativeStartingPoint) > 1 {
			rightBeforeStartingPoint := relativeStartingPoint[0 : len(relativeStartingPoint)-1]
			startAfter := targetPrefix + rightBeforeStartingPoint

			// StartAfter is also known as 'marker' within S3 compatible layer
			input.StartAfter = aws.String(startAfter)
		}
	}

	var relativeEndPoint string
	if exclusiveEndPoint != "" {
		if !strings.HasPrefix(exclusiveEndPoint, prefix) {
			return fmt.Errorf("exclusive end point %q must start with prefix %q", exclusiveEndPoint, prefix)
		}
		relativeEndPoint = strings.TrimPrefix(exclusiveEndPoint, prefix)
	}

	if tracer.Enabled() {
		zlog.Info("walking files from", zap.String("original_prefix", targetPrefix), zap.String("prefix", targetPrefix), zap.Stringp("start_after", input.StartAfter))
	}

	paginator := s3.NewListObjectsV2Paginator(s.client, input)

	for paginator.HasMorePages() {
		page, err := paginator.NextPage(ctx)
		if err != nil {
			return fmt.Errorf("listing objects: %w", err)
		}

		for _, obj := range page.Contents {
			filename := s.toBaseName(*obj.Key)
			if filename == "" {
				zlog.Debug("got an empty filename from s3 store, ignoring it", zap.String("key", *obj.Key))
				continue
			}

			if startingPoint != "" && filename < startingPoint {
				continue
			}
			if relativeEndPoint != "" && filename >= relativeEndPoint {
				return nil
			}

			if err := f(filename); err != nil {
				if errors.Is(err, StopIteration) {
					return nil
				}
				return fmt.Errorf("processing object list: %w", err)
			}
		}
	}

	return nil
}

func (s *S3Store) Walk(ctx context.Context, prefix string, f func(filename string) (err error)) error {
	return s.WalkFrom(ctx, prefix, "", f)
}

func (s *S3Store) toBaseName(filename string) string {
	return strings.TrimPrefix(strings.TrimSuffix(filename, s.pathWithExt("")), s.path+"/")
}

func (s *S3Store) DeleteObject(ctx context.Context, base string) error {
	path := s.ObjectPath(base)
	_, err := s.client.DeleteObject(ctx, &s3.DeleteObjectInput{
		Bucket: aws.String(s.bucket),
		Key:    aws.String(path),
	})
	if err != nil {
		var noSuchKey *types.NoSuchKey
		if errors.As(err, &noSuchKey) {
			return ErrNotFound
		}
	}
	return err
}

func (s *S3Store) PushLocalFile(ctx context.Context, localFile, toBaseName string) error {
	remove, err := pushLocalFile(ctx, s, localFile, toBaseName)
	if retryS3PushLocalFilesDelay != 0 {
		time.Sleep(retryS3PushLocalFilesDelay)
		exists, err := s.FileExists(ctx, toBaseName)
		if err != nil {
			zlog.Debug("just pushed file to dstore, but cannot check if it is still there after 500 milliseconds and retryS3PushLocalFiles is set", zap.Error(err))
			return err
		}
		if !exists {
			zlog.Debug("just pushed file to dstore, but it disappeared. Pushing again because retryS3PushLocalFiles is set", zap.String("dest basename", toBaseName))
			rem, err := pushLocalFile(ctx, s, localFile, toBaseName)
			if err != nil {
				return err
			}
			return rem()
		}
	}

	if err != nil {
		return err
	}
	return remove()
}

func (s *S3Store) ListFiles(ctx context.Context, prefix string, max int) ([]string, error) {
	return listFiles(ctx, s, prefix, max)
}