go.charczuk.com@v0.0.0-20240327042549-bc490516bd1a/projects/blogctl/pkg/aws/s3/manager.go

/*

Copyright (c) 2023 - Present. Will Charczuk. All rights reserved.
Use of this source code is governed by a MIT license that can be found in the LICENSE file at the root of the repository.

*/

package s3

import (
	"bytes"
	"context"
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"sync"

	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"

	"go.charczuk.com/sdk/async"
	"go.charczuk.com/sdk/logutil"

	"go.charczuk.com/projects/blogctl/pkg/aws"
)

// New returns a new manager.
func New(cfg aws.Config) *Manager {
	return &Manager{
		Config: cfg,
		Ignores: []string{
			".DS_Store",
			".git",
		},
		Session:     aws.NewSession(cfg),
		Parallelism: runtime.NumCPU(),
	}
}

// Manager is a helper for uploading files to s3.
type Manager struct {
	Log               *log.Logger
	Ignores           []string
	Config            aws.Config
	Session           *session.Session
	PutObjectDefaults File
	DryRun            bool
	Parallelism       int
}

// ParallelismOrDefault returns the parallelism or a default.
func (m Manager) ParallelismOrDefault() int {
	if m.Parallelism > 0 {
		return m.Parallelism
	}
	return runtime.NumCPU()
}

// GetKey returns the relative path for a given file.
func (m Manager) GetKey(rootPath, workingPath string) string {
	if !strings.HasPrefix(workingPath, "./") {
		workingPath = "./" + workingPath
	}
	if !strings.HasPrefix(rootPath, "./") {
		rootPath = "./" + rootPath
	}
	return strings.TrimPrefix(workingPath, rootPath)
}

// SyncDirectory sync's a directory.
// It returns a list of invalidated keys (i.e. keys to update or remove), and an error.
func (m Manager) SyncDirectory(ctx context.Context, directoryPath, bucket string) (invalidations []string, err error) {
	if m.DryRun {
		logutil.Debug(m.Log, "sync directory (dry run): not realizing changes")
	}
	localFiles, err := m.DiscoverFiles(ctx, directoryPath)
	if err != nil {
		return nil, err
	}
	invalidations, err = m.ProcessFiles(ctx, localFiles, directoryPath, bucket)
	return
}

// DiscoverFiles discovers local files.
func (m Manager) DiscoverFiles(ctx context.Context, directoryPath string) (localFiles []string, err error) {
	err = filepath.Walk(directoryPath, func(currentPath string, fileInfo os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if currentPath == directoryPath {
			return nil
		}
		for _, ignore := range m.Ignores {
			if strings.HasSuffix(currentPath, ignore) {
				return nil
			}
		}
		if fileInfo.IsDir() {
			return nil
		}
		localFiles = append(localFiles, currentPath)
		return nil
	})
	return
}

// ProcessFiles processes the files list.
func (m Manager) ProcessFiles(ctx context.Context, localFiles []string, directoryPath, bucket string) (invalidated []string, err error) {
	remoteETags := make(map[string]string)
	localKeys := new(Set)

	remoteFiles, err := m.List(ctx, bucket)
	if err != nil {
		return nil, err
	}

	var remoteFileBatch = make([]File, 0, len(remoteFiles))
	for _, remoteFile := range remoteFiles {
		key := remoteFile.Key
		if !strings.HasPrefix(key, "/") {
			key = "/" + key
		}
		logutil.Debugf(m.Log, "%s: remote file has etag %s", key, remoteFile.ETag)
		remoteETags[key] = aws.StripQuotes(remoteFile.ETag)
		remoteFileBatch = append(remoteFileBatch, remoteFile)
	}

	b := new(async.Batch)
	b.SetLimit(m.ParallelismOrDefault())

	var processLocalFile = func(file string) func() error {
		return func() error {
			key := m.GetKey(directoryPath, file)
			localKeys.Set(key)

			var localETag string
			remoteETag, hasRemoteFile := remoteETags[key]
			if hasRemoteFile { // if we need to compare against a remote etag
				logutil.Debugf(m.Log, "%s: generating local file etag for key", key)
				localETag, err = m.GenerateETag(file)
				if err != nil {
					return err
				}
			} else {
				logutil.Debugf(m.Log, "%s: missing remote file etag for key", key)
			}

			if !hasRemoteFile || remoteETag != localETag {
				if !hasRemoteFile {
					logutil.Debugf(m.Log, "%s: local file is not present on remote", key)
				} else if hasRemoteFile && remoteETag != localETag {
					logutil.Debugf(m.Log, "%s: local file has different etag than remote; %s vs. %v", key, localETag, remoteETag)
				}

				contentType, err := DetectContentType(file)
				if err != nil {
					return err
				}

				if !m.DryRun {
					if err := m.Put(ctx, File{
						FilePath:    file,
						Key:         key,
						Bucket:      bucket,
						ContentType: contentType,
					}); err != nil {
						return err
					}
					logutil.Infof(m.Log, "%s: put file to remote", key)
				} else {
					logutil.Infof(m.Log, "%s: (dry run) put file to remote", key)
				}
				if hasRemoteFile {
					logutil.Infof(m.Log, "%s: marking to be invalidated", key)
					invalidated = append(invalidated, key)
				}
			} else {
				logutil.Debugf(m.Log, "%s: skipping (unchanged)", key)
			}
			return nil
		}
	}
	for _, file := range localFiles {
		b.Go(processLocalFile(file))
	}
	if err := b.Wait(); err != nil {
		return nil, err
	}

	rb := new(async.Batch)
	rb.SetLimit(m.ParallelismOrDefault())

	var invalidatedSync sync.Mutex
	var invalidateRemoteFile = func(remoteFile File) func() error {
		return func() error {
			key := remoteFile.Key
			if !strings.HasPrefix(key, "/") {
				key = "/" + key
			}
			if !localKeys.Has(key) {
				if !m.DryRun {
					logutil.Infof(m.Log, "%s: removing remote file", remoteFile.Key)
					if err := m.Delete(ctx, bucket, remoteFile.Key); err != nil {
						return err
					}
				} else {
					logutil.Infof(m.Log, "%s: (dry run) removing remote file", remoteFile.Key)
				}

				invalidatedSync.Lock()
				invalidated = append(invalidated, key)
				invalidatedSync.Unlock()
			} else {
				logutil.Debugf(m.Log, "%s: keeping remote file", remoteFile.Key)
			}
			return nil
		}
	}
	for _, remoteFile := range remoteFileBatch {
		rb.Go(invalidateRemoteFile(remoteFile))
	}
	if err := rb.Wait(); err != nil {
		return nil, err
	}
	return invalidated, nil
}

// List lists all files in a bucket.
func (m Manager) List(ctx context.Context, bucket string) ([]File, error) {
	remoteFiles, err := s3.New(m.Session).ListObjectsWithContext(ctx, &s3.ListObjectsInput{
		Bucket: &bucket,
	})
	if IsNotFound(err) {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}

	var files []File
	for _, file := range remoteFiles.Contents {
		files = append(files, File{
			Bucket: bucket,
			Key:    aws.DerefStr(file.Key),
			ETag:   aws.DerefStr(file.ETag),
		})
	}
	return files, nil
}

// Get fetches a file at a given key
func (m Manager) Get(ctx context.Context, bucket, key string) (file File, contents io.ReadCloser, err error) {
	remoteFile, getErr := s3.New(m.Session).GetObjectWithContext(ctx, &s3.GetObjectInput{
		Bucket: &bucket,
		Key:    &key,
	})
	if IsNotFound(getErr) {
		return
	}
	if getErr != nil {
		err = getErr
		return
	}

	file = File{
		Bucket:               bucket,
		Key:                  key,
		ContentType:          aws.DerefStr(remoteFile.ContentType),
		ContentDisposition:   aws.DerefStr(remoteFile.ContentDisposition),
		ServerSideEncryption: aws.DerefStr(remoteFile.ServerSideEncryption),
		ETag:                 aws.DerefStr(remoteFile.ETag),
	}
	contents = remoteFile.Body
	return
}

// GetMeta fetches file metadata at a given key
func (m Manager) GetMeta(ctx context.Context, bucket, key string) (meta File, err error) {
	var contents io.ReadCloser
	meta, contents, err = m.Get(ctx, bucket, key)
	if err != nil {
		return
	}
	if contents != nil {
		defer contents.Close()
	}
	return
}

// Put uploads a file to s3.
func (m Manager) Put(ctx context.Context, fileInfo File) error {
	var size int64
	var contentType, contentDisposition, acl, serverSideEncryption string
	var contents io.ReadSeeker

	if fileInfo.FilePath != "" {
		file, err := os.Open(fileInfo.FilePath)
		if err != nil {
			return err
		}
		defer file.Close()
		stats, err := file.Stat()
		if err != nil {
			return err
		}
		contents = file
		size = int64(stats.Size())
	} else if len(fileInfo.Contents) > 0 {
		size = int64(len(fileInfo.Contents))
		contents = bytes.NewReader(fileInfo.Contents)
	} else {
		return fmt.Errorf("invalid put object; must set either the path or the contents")
	}

	if fileInfo.ContentType != "" {
		contentType = fileInfo.ContentType
	} else if m.PutObjectDefaults.ContentDisposition != "" {
		contentType = m.PutObjectDefaults.ContentType
	}

	if fileInfo.ContentDisposition != "" {
		contentDisposition = fileInfo.ContentDisposition
	} else if m.PutObjectDefaults.ContentDisposition != "" {
		contentDisposition = m.PutObjectDefaults.ContentDisposition
	}

	if fileInfo.ACL != "" {
		acl = fileInfo.ACL
	} else if m.PutObjectDefaults.ACL != "" {
		acl = m.PutObjectDefaults.ACL
	}

	if fileInfo.ServerSideEncryption != "" {
		serverSideEncryption = fileInfo.ServerSideEncryption
	} else if m.PutObjectDefaults.ServerSideEncryption != "" {
		serverSideEncryption = m.PutObjectDefaults.ServerSideEncryption
	}

	_, err := s3.New(m.Session).PutObject(&s3.PutObjectInput{
		Bucket:               aws.RefStr(fileInfo.Bucket),
		Key:                  aws.RefStr(fileInfo.Key),
		Body:                 contents,
		ContentLength:        &size,
		ContentType:          aws.RefStr(contentType),
		ContentDisposition:   aws.RefStr(contentDisposition),
		ACL:                  aws.RefStr(acl),
		ServerSideEncryption: aws.RefStr(serverSideEncryption),
	})
	return err
}

// Delete removes an object with a given key.
func (m Manager) Delete(ctx context.Context, bucket, key string) error {
	_, err := s3.New(m.Session).DeleteObject(&s3.DeleteObjectInput{
		Bucket: aws.RefStr(bucket),
		Key:    aws.RefStr(key),
	})
	return err
}

// GenerateETag generate an etag for a give file by path.
func (m Manager) GenerateETag(path string) (string, error) {
	f, err := os.Open(path)
	if err != nil {
		return "", err
	}

	hash := md5.New()
	_, err = io.Copy(hash, f)
	if err != nil {
		return "", err
	}

	return hex.EncodeToString(hash.Sum(nil)), nil
}