probo/pkg/filemanager/s3.go

// Copyright (c) 2025-2026 Probo Inc <hello@probo.com>.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

package filemanager

import (
	"context"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go-v2/feature/s3/transfermanager"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	smithyhttp "github.com/aws/smithy-go/transport/http"
	"go.probo.inc/probo/pkg/coredata"
)

type FileObject struct {
	Body                io.ReadCloser
	ContentType         string
	ContentLength       int64
	ContentRange        string
	ETag                string
	LastModified        time.Time
	NotModified         bool
	PartialContent      bool
	RangeNotSatisfiable bool
}

type FileConditions struct {
	IfNoneMatch     string
	IfModifiedSince time.Time
	IfRange         string
	Range           string
}

func (s *Service) GetFileBase64(
	ctx context.Context,
	file *coredata.File,
) (base64Data string, mimeType string, err error) {
	result, err := s.s3Client.GetObject(
		ctx,
		&s3.GetObjectInput{
			Bucket: new(file.BucketName),
			Key:    new(file.FileKey),
		},
	)
	if err != nil {
		return "", "", fmt.Errorf("cannot get file from S3: %w", err)
	}

	defer func() { _ = result.Body.Close() }()

	fileData, err := io.ReadAll(result.Body)
	if err != nil {
		return "", "", fmt.Errorf("cannot read file data: %w", err)
	}

	return base64.StdEncoding.EncodeToString(fileData), file.MimeType, nil
}

func (s *Service) GetFileBytes(
	ctx context.Context,
	file *coredata.File,
) ([]byte, error) {
	result, err := s.s3Client.GetObject(
		ctx,
		&s3.GetObjectInput{
			Bucket: new(file.BucketName),
			Key:    new(file.FileKey),
		},
	)
	if err != nil {
		return nil, fmt.Errorf("cannot get file from S3: %w", err)
	}

	defer func() { _ = result.Body.Close() }()

	data, err := io.ReadAll(result.Body)
	if err != nil {
		return nil, fmt.Errorf("cannot read file data: %w", err)
	}

	return data, nil
}

func (s *Service) OpenFile(
	ctx context.Context,
	file *coredata.File,
	conds FileConditions,
) (*FileObject, error) {
	input := &s3.GetObjectInput{
		Bucket: new(file.BucketName),
		Key:    new(file.FileKey),
	}
	if conds.IfNoneMatch != "" {
		input.IfNoneMatch = &conds.IfNoneMatch
	}

	if !conds.IfModifiedSince.IsZero() {
		input.IfModifiedSince = &conds.IfModifiedSince
	}

	if conds.Range != "" {
		honorRange := true

		if conds.IfRange != "" {
			head, err := s.s3Client.HeadObject(
				ctx,
				&s3.HeadObjectInput{
					Bucket: new(file.BucketName),
					Key:    new(file.FileKey),
				},
			)
			if err != nil {
				return nil, fmt.Errorf("cannot head file in S3: %w", err)
			}

			etag := ""
			if head.ETag != nil {
				etag = *head.ETag
			}

			lastModified := file.UpdatedAt
			if head.LastModified != nil {
				lastModified = *head.LastModified
			}

			honorRange = ifRangeMatches(conds.IfRange, etag, lastModified)
		}

		if honorRange {
			input.Range = &conds.Range
		}
	}

	result, err := s.s3Client.GetObject(ctx, input)
	if err != nil {
		if respErr, ok := errors.AsType[*smithyhttp.ResponseError](err); ok {
			switch respErr.HTTPStatusCode() {
			case http.StatusNotModified:
				return &FileObject{NotModified: true}, nil
			case http.StatusRequestedRangeNotSatisfiable:
				return &FileObject{
					RangeNotSatisfiable: true,
					ContentLength:       file.FileSize,
				}, nil
			}
		}

		return nil, fmt.Errorf("cannot get file from S3: %w", err)
	}

	obj := &FileObject{
		Body:          result.Body,
		ContentType:   file.MimeType,
		ContentLength: file.FileSize,
		LastModified:  file.UpdatedAt,
	}
	if result.ETag != nil {
		obj.ETag = *result.ETag
	}

	if result.LastModified != nil {
		obj.LastModified = *result.LastModified
	}

	// A Range that S3 honors comes back as 206 Partial Content with a
	// Content-Range header and a ContentLength scoped to the returned slice. An
	// unranged request (or one whose Range we dropped above after an If-Range
	// mismatch) yields a normal 200 with the full object, so we only override
	// the length/status when Content-Range is present.
	if result.ContentRange != nil {
		obj.ContentRange = *result.ContentRange
		obj.PartialContent = true
	}

	if result.ContentLength != nil {
		obj.ContentLength = *result.ContentLength
	}

	return obj, nil
}

type putFileConfig struct {
	attachmentDisposition bool
}

// PutFileOption configures optional PutFile behavior.
type PutFileOption func(*putFileConfig)

// WithAttachmentContentDisposition stores Content-Disposition: attachment on the
// uploaded object so browsers download it when served from object storage.
func WithAttachmentContentDisposition() PutFileOption {
	return func(cfg *putFileConfig) {
		cfg.attachmentDisposition = true
	}
}

func (s *Service) PutFile(
	ctx context.Context,
	file *coredata.File,
	content io.Reader,
	metadata map[string]string,
	opts ...PutFileOption,
) (int64, error) {
	cfg := putFileConfig{}
	for _, opt := range opts {
		opt(&cfg)
	}

	// Transfer manager accepts unseekable readers (e.g. io.Pipe) by buffering
	// parts in memory, which works against plain-HTTP S3-compatible endpoints
	// where PutObject checksums require a seekable body.
	uploader := transfermanager.New(s.s3Client)

	input := &transfermanager.UploadObjectInput{
		Bucket:       new(file.BucketName),
		Key:          new(file.FileKey),
		Body:         content,
		ContentType:  new(file.MimeType),
		CacheControl: new("private, max-age=3600"),
		Metadata:     metadata,
	}
	if cfg.attachmentDisposition && file.FileName != "" {
		input.ContentDisposition = new(attachmentContentDisposition(file.FileName))
	}

	_, err := uploader.UploadObject(ctx, input)
	if err != nil {
		return 0, fmt.Errorf("cannot upload file to S3: %w", err)
	}

	headOutput, err := s.s3Client.HeadObject(
		ctx,
		&s3.HeadObjectInput{
			Bucket: new(file.BucketName),
			Key:    new(file.FileKey),
		},
	)
	if err != nil {
		return 0, fmt.Errorf("cannot get object metadata: %w", err)
	}

	return *headOutput.ContentLength, nil
}

func (s *Service) GeneratePresignedURL(
	ctx context.Context,
	file *coredata.File,
	expiresIn time.Duration,
) (string, error) {
	presignClient := s3.NewPresignClient(s.s3Client)

	contentDisposition := attachmentContentDisposition(file.FileName)

	presignedReq, err := presignClient.PresignGetObject(
		ctx,
		&s3.GetObjectInput{
			Bucket:                     new(file.BucketName),
			Key:                        new(file.FileKey),
			ResponseCacheControl:       new("max-age=3600, public"),
			ResponseContentType:        new(file.MimeType),
			ResponseContentDisposition: &contentDisposition,
		},
		func(opts *s3.PresignOptions) {
			opts.Expires = expiresIn
		},
	)
	if err != nil {
		return "", fmt.Errorf("cannot presign GetObject request: %w", err)
	}

	return presignedReq.URL, nil
}

// ifRangeMatches reports whether an If-Range validator still matches the
// current representation, following the strong-comparison rules of RFC 9110
// section 13.1.3. S3 has no native If-Range support, so callers evaluate it
// before deciding whether to forward a Range header. An entity-tag validator
// must strong-match the current ETag (weak tags on either side never satisfy
// If-Range); otherwise the validator is an HTTP-date compared for equality
// against Last-Modified.
func ifRangeMatches(ifRange, etag string, lastModified time.Time) bool {
	ifRange = strings.TrimSpace(ifRange)
	if ifRange == "" {
		return true
	}

	if strings.HasPrefix(ifRange, `"`) {
		return etag != "" && !strings.HasPrefix(etag, `W/`) && ifRange == etag
	}

	if strings.HasPrefix(ifRange, `W/`) {
		return false
	}

	t, err := http.ParseTime(ifRange)
	if err != nil {
		return false
	}

	return !lastModified.IsZero() && lastModified.Truncate(time.Second).Equal(t)
}

func attachmentContentDisposition(filename string) string {
	return fmt.Sprintf(
		"attachment; filename=%q; filename*=UTF-8''%s",
		asciiFilename(filename),
		url.PathEscape(filename),
	)
}

func asciiFilename(filename string) string {
	var b strings.Builder
	b.Grow(len(filename))

	for _, r := range filename {
		if r < 0x20 || r > 0x7e {
			b.WriteByte('_')
			continue
		}

		b.WriteRune(r)
	}

	return b.String()
}

// GetFileSize determines the byte size of a seekable io.Reader by seeking to
// the end and back. Returns an error if content is not seekable.
func GetFileSize(content io.Reader) (int64, error) {
	seeker, ok := content.(io.Seeker)
	if !ok {
		return 0, fmt.Errorf("cannot determine file size: content is not seekable")
	}

	size, err := seeker.Seek(0, io.SeekEnd)
	if err != nil {
		return 0, fmt.Errorf("cannot determine file size: %w", err)
	}

	_, err = seeker.Seek(0, io.SeekStart)
	if err != nil {
		return 0, fmt.Errorf("cannot reset file position: %w", err)
	}

	return size, nil
}