distributor/internal/adapters/s3/backend_test.go

package s3

import (
	"context"
	"io"
	"sort"
	"strings"
	"testing"

	"gitea.maximumdirect.net/eric/distributor/internal/storage"
	"github.com/aws/aws-sdk-go-v2/aws"
	awss3 "github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/aws/aws-sdk-go-v2/service/s3/types"
)

func TestKeyAndPrefixNormalization(t *testing.T) {
	backend := newTestBackend(t, "root/prefix", nil)
	key, err := backend.objectKey("bundle/report.md", false)
	if err != nil {
		t.Fatalf("objectKey() error = %v", err)
	}
	if got, want := key, "root/prefix/bundle/report.md"; got != want {
		t.Fatalf("objectKey() = %q, want %q", got, want)
	}
	listPrefix, err := backend.listPrefix("bundle")
	if err != nil {
		t.Fatalf("listPrefix() error = %v", err)
	}
	if got, want := listPrefix, "root/prefix/bundle/"; got != want {
		t.Fatalf("listPrefix() = %q, want %q", got, want)
	}
}

func TestPathTraversalRejected(t *testing.T) {
	backend := newTestBackend(t, "", nil)
	for _, logicalPath := range []string{"/absolute", "../escape", "a/../b", `a\b`} {
		t.Run(logicalPath, func(t *testing.T) {
			if _, err := backend.objectKey(logicalPath, false); err == nil || !storage.IsInvalidPath(err) {
				t.Fatalf("objectKey() error = %v, want invalid path", err)
			}
		})
	}
}

func TestContentType(t *testing.T) {
	tests := map[string]string{
		"report.md":   "text/markdown; charset=utf-8",
		"report.html": "text/html; charset=utf-8",
		"state.json":  "application/json",
		"summary.txt": "text/plain; charset=utf-8",
		"data.bin":    "application/octet-stream",
	}
	for path, want := range tests {
		if got := ContentType(path); got != want {
			t.Fatalf("ContentType(%q) = %q, want %q", path, got, want)
		}
	}
}

func TestStatRequiresExactObject(t *testing.T) {
	client := newFakeClient(map[string]string{"root/dir/file.txt": "data"})
	backend := newTestBackend(t, "root", client)
	_, err := backend.Stat(context.Background(), "dir")
	if err == nil || !storage.IsNotFound(err) {
		t.Fatalf("Stat() error = %v, want not found", err)
	}
}

func TestWriteFromChecksOverwriteBeforePut(t *testing.T) {
	client := newFakeClient(map[string]string{"root/report.md": "old"})
	backend := newTestBackend(t, "root", client)
	_, err := backend.WriteFile(context.Background(), "report.md", []byte("new"), storage.WriteOptions{})
	if err == nil || !storage.IsAlreadyExists(err) {
		t.Fatalf("WriteFile() error = %v, want already exists", err)
	}
	if len(client.putKeys) != 0 {
		t.Fatalf("put keys = %v, want none", client.putKeys)
	}
}

func TestWriteFromPutsNewObjectWithContentType(t *testing.T) {
	client := newFakeClient(nil)
	backend := newTestBackend(t, "root", client)
	entry, err := backend.WriteFile(context.Background(), "report.html", []byte("<p>ok</p>"), storage.WriteOptions{})
	if err != nil {
		t.Fatalf("WriteFile() error = %v", err)
	}
	if entry.Path != "report.html" || entry.Size != 9 || entry.Type != storage.EntryTypeFile {
		t.Fatalf("entry = %#v", entry)
	}
	if got, want := client.objects["root/report.html"], "<p>ok</p>"; got != want {
		t.Fatalf("object = %q, want %q", got, want)
	}
	if got, want := client.contentTypes["root/report.html"], "text/html; charset=utf-8"; got != want {
		t.Fatalf("content type = %q, want %q", got, want)
	}
}

func TestWalkUsesPagination(t *testing.T) {
	client := newFakeClient(nil)
	client.listPages = []awss3.ListObjectsV2Output{
		{
			Contents:              []types.Object{{Key: aws.String("root/a.txt"), Size: aws.Int64(1)}},
			IsTruncated:           aws.Bool(true),
			NextContinuationToken: aws.String("next"),
		},
		{
			Contents:    []types.Object{{Key: aws.String("root/b.txt"), Size: aws.Int64(2)}},
			IsTruncated: aws.Bool(false),
		},
	}
	backend := newTestBackend(t, "root", client)
	entries, err := storage.List(context.Background(), backend, "", storage.WalkOptions{Recursive: true})
	if err != nil {
		t.Fatalf("List() error = %v", err)
	}
	paths := entryPaths(entries)
	if got, want := paths, []string{"a.txt", "b.txt"}; !equalStrings(got, want) {
		t.Fatalf("paths = %v, want %v", got, want)
	}
	if got, want := client.tokens, []string{"", "next"}; !equalStrings(got, want) {
		t.Fatalf("tokens = %v, want %v", got, want)
	}
}

func TestHasAnyStopsAfterFirstPage(t *testing.T) {
	client := newFakeClient(nil)
	client.listPages = []awss3.ListObjectsV2Output{
		{
			Contents:              []types.Object{{Key: aws.String("root/a.txt"), Size: aws.Int64(1)}},
			IsTruncated:           aws.Bool(true),
			NextContinuationToken: aws.String("next"),
		},
		{
			Contents:    []types.Object{{Key: aws.String("root/b.txt"), Size: aws.Int64(2)}},
			IsTruncated: aws.Bool(false),
		},
	}
	backend := newTestBackend(t, "root", client)
	found, err := backend.HasAny(context.Background(), "")
	if err != nil {
		t.Fatalf("HasAny() error = %v", err)
	}
	if !found {
		t.Fatal("HasAny() = false, want true")
	}
	if got, want := len(client.tokens), 1; got != want {
		t.Fatalf("list calls = %d, want %d", got, want)
	}
}

func TestWalkNonRecursiveUsesPrefixBoundary(t *testing.T) {
	client := newFakeClient(map[string]string{
		"base/dir/file.txt": "nested",
		"base/file.txt":     "file",
		"baseball/file.txt": "wrong",
	})
	backend := newTestBackend(t, "base", client)
	entries, err := storage.List(context.Background(), backend, "", storage.WalkOptions{})
	if err != nil {
		t.Fatalf("List() error = %v", err)
	}
	paths := entryPaths(entries)
	if got, want := paths, []string{"dir", "file.txt"}; !equalStrings(got, want) {
		t.Fatalf("paths = %v, want %v", got, want)
	}
}

func TestDeleteManagedBundleDeletesOnlyManagedTargets(t *testing.T) {
	client := newFakeClient(map[string]string{
		"root/report.md":         "report",
		"root/.distributor.json": "state",
		"root/keep.txt":          "keep",
	})
	backend := newTestBackend(t, "root", client)
	err := backend.DeleteManagedBundle(context.Background(), "", []string{"report.md"}, storage.DeleteOptions{IgnoreMissing: true})
	if err != nil {
		t.Fatalf("DeleteManagedBundle() error = %v", err)
	}
	if _, ok := client.objects["root/report.md"]; ok {
		t.Fatal("managed output still exists")
	}
	if _, ok := client.objects["root/.distributor.json"]; ok {
		t.Fatal("state file still exists")
	}
	if _, ok := client.objects["root/keep.txt"]; !ok {
		t.Fatal("unmanaged object was deleted")
	}
	if got, want := sortedStrings(client.deleteKeys), []string{"root/.distributor.json", "root/report.md"}; !equalStrings(got, want) {
		t.Fatalf("deleted keys = %v, want %v", got, want)
	}
}

func newTestBackend(t *testing.T, prefix string, client *fakeClient) *Backend {
	t.Helper()
	if client == nil {
		client = newFakeClient(nil)
	}
	backend, err := NewWithClient(client, Options{
		Endpoint:       "https://s3.example.com",
		Bucket:         "reports",
		Prefix:         prefix,
		Region:         DefaultRegion,
		ForcePathStyle: true,
	})
	if err != nil {
		t.Fatalf("NewWithClient() error = %v", err)
	}
	return backend
}

type fakeClient struct {
	objects      map[string]string
	contentTypes map[string]string
	listPages    []awss3.ListObjectsV2Output
	tokens       []string
	putKeys      []string
	deleteKeys   []string
}

func newFakeClient(objects map[string]string) *fakeClient {
	copied := make(map[string]string)
	for key, value := range objects {
		copied[key] = value
	}
	return &fakeClient{
		objects:      copied,
		contentTypes: make(map[string]string),
	}
}

func (c *fakeClient) HeadObject(ctx context.Context, input *awss3.HeadObjectInput, optFns ...func(*awss3.Options)) (*awss3.HeadObjectOutput, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	value, ok := c.objects[aws.ToString(input.Key)]
	if !ok {
		return nil, &types.NotFound{}
	}
	return &awss3.HeadObjectOutput{ContentLength: aws.Int64(int64(len(value)))}, nil
}

func (c *fakeClient) GetObject(ctx context.Context, input *awss3.GetObjectInput, optFns ...func(*awss3.Options)) (*awss3.GetObjectOutput, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	value, ok := c.objects[aws.ToString(input.Key)]
	if !ok {
		return nil, &types.NotFound{}
	}
	return &awss3.GetObjectOutput{
		Body:          io.NopCloser(stringsReader(value)),
		ContentLength: aws.Int64(int64(len(value))),
	}, nil
}

func (c *fakeClient) PutObject(ctx context.Context, input *awss3.PutObjectInput, optFns ...func(*awss3.Options)) (*awss3.PutObjectOutput, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	data, err := io.ReadAll(input.Body)
	if err != nil {
		return nil, err
	}
	key := aws.ToString(input.Key)
	c.objects[key] = string(data)
	c.contentTypes[key] = aws.ToString(input.ContentType)
	c.putKeys = append(c.putKeys, key)
	return &awss3.PutObjectOutput{}, nil
}

func (c *fakeClient) ListObjectsV2(ctx context.Context, input *awss3.ListObjectsV2Input, optFns ...func(*awss3.Options)) (*awss3.ListObjectsV2Output, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	c.tokens = append(c.tokens, aws.ToString(input.ContinuationToken))
	if len(c.listPages) > 0 {
		index := len(c.tokens) - 1
		if index >= len(c.listPages) {
			return &awss3.ListObjectsV2Output{IsTruncated: aws.Bool(false)}, nil
		}
		page := c.listPages[index]
		return &page, nil
	}
	return c.dynamicList(input), nil
}

func (c *fakeClient) DeleteObject(ctx context.Context, input *awss3.DeleteObjectInput, optFns ...func(*awss3.Options)) (*awss3.DeleteObjectOutput, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	key := aws.ToString(input.Key)
	delete(c.objects, key)
	c.deleteKeys = append(c.deleteKeys, key)
	return &awss3.DeleteObjectOutput{}, nil
}

func (c *fakeClient) dynamicList(input *awss3.ListObjectsV2Input) *awss3.ListObjectsV2Output {
	prefix := aws.ToString(input.Prefix)
	delimiter := aws.ToString(input.Delimiter)
	var contents []types.Object
	commonPrefixes := make(map[string]struct{})
	for key, value := range c.objects {
		if !strings.HasPrefix(key, prefix) {
			continue
		}
		remainder := strings.TrimPrefix(key, prefix)
		if delimiter != "" {
			if index := strings.Index(remainder, delimiter); index >= 0 {
				commonPrefixes[prefix+remainder[:index+1]] = struct{}{}
				continue
			}
		}
		contents = append(contents, types.Object{Key: aws.String(key), Size: aws.Int64(int64(len(value)))})
	}
	sort.Slice(contents, func(i, j int) bool { return aws.ToString(contents[i].Key) < aws.ToString(contents[j].Key) })
	prefixes := make([]types.CommonPrefix, 0, len(commonPrefixes))
	for prefix := range commonPrefixes {
		prefixes = append(prefixes, types.CommonPrefix{Prefix: aws.String(prefix)})
	}
	sort.Slice(prefixes, func(i, j int) bool { return aws.ToString(prefixes[i].Prefix) < aws.ToString(prefixes[j].Prefix) })
	return &awss3.ListObjectsV2Output{
		Contents:       contents,
		CommonPrefixes: prefixes,
		IsTruncated:    aws.Bool(false),
	}
}

func stringsReader(value string) io.Reader {
	return strings.NewReader(value)
}

func entryPaths(entries []storage.Entry) []string {
	paths := make([]string, 0, len(entries))
	for _, entry := range entries {
		paths = append(paths, entry.Path)
	}
	return paths
}

func sortedStrings(values []string) []string {
	copied := append([]string(nil), values...)
	sort.Strings(copied)
	return copied
}

func equalStrings(a, b []string) bool {
	if len(a) != len(b) {
		return false
	}
	for index := range a {
		if a[index] != b[index] {
			return false
		}
	}
	return true
}