package storage
import (
"bytes"
"context"
"errors"
"io"
"strings"
"sync"
"testing"
"github.com/aws/aws-sdk-go-v2/service/s3"
"github.com/aws/aws-sdk-go-v2/service/s3/types"
)
type mockS3Client struct {
mu sync.Mutex
objects map[string][]byte
}
func (m *mockS3Client) PutObject(_ context.Context, params *s3.PutObjectInput, _ ...func(*s3.Options)) (*s3.PutObjectOutput, error) {
m.mu.Lock()
defer m.mu.Unlock()
data, err := io.ReadAll(params.Body)
if err != nil {
return nil, err
}
m.objects[*params.Bucket+"/"+*params.Key] = data
return &s3.PutObjectOutput{}, nil
}
func (m *mockS3Client) GetObject(_ context.Context, params *s3.GetObjectInput, _ ...func(*s3.Options)) (*s3.GetObjectOutput, error) {
m.mu.Lock()
defer m.mu.Unlock()
data, ok := m.objects[*params.Bucket+"/"+*params.Key]
if !ok {
return nil, &types.NoSuchKey{}
}
return &s3.GetObjectOutput{Body: io.NopCloser(bytes.NewReader(data))}, nil
}
func (m *mockS3Client) DeleteObject(_ context.Context, params *s3.DeleteObjectInput, _ ...func(*s3.Options)) (*s3.DeleteObjectOutput, error) {
m.mu.Lock()
defer m.mu.Unlock()
delete(m.objects, *params.Bucket+"/"+*params.Key)
return &s3.DeleteObjectOutput{}, nil
}
func TestValidateKey(t *testing.T) {
valid := []string{
"did:plc:xyz123/go-mod-v1",
"did:web:spindle.example.com/cache.tar",
"abc",
"a/b/c/d",
}
for _, k := range valid {
if err := ValidateKey(k); err != nil {
t.Errorf("ValidateKey(%q) = %v, want nil", k, err)
}
}
invalid := []string{
"",
"../escape",
"a/../../b",
"/leading",
"trailing/",
"double//slash",
"with space",
"with\\backslash",
"-leading-dash",
}
for _, k := range invalid {
if err := ValidateKey(k); err == nil {
t.Errorf("ValidateKey(%q) = nil, want error", k)
}
}
}
func TestDiskRoundTrip(t *testing.T) {
ctx := context.Background()
d, err := NewDisk(t.TempDir())
if err != nil {
t.Fatal(err)
}
key := "did:plc:xyz/go-mod-v1"
if err := d.Put(ctx, key, strings.NewReader("archive-bytes")); err != nil {
t.Fatal(err)
}
rc, err := d.Get(ctx, key)
if err != nil {
t.Fatal(err)
}
got, err := io.ReadAll(rc)
rc.Close()
if err != nil {
t.Fatal(err)
}
if string(got) != "archive-bytes" {
t.Fatalf("got %q, want %q", got, "archive-bytes")
}
// overwrite
if err := d.Put(ctx, key, strings.NewReader("new-bytes")); err != nil {
t.Fatal(err)
}
rc, _ = d.Get(ctx, key)
got, _ = io.ReadAll(rc)
rc.Close()
if string(got) != "new-bytes" {
t.Fatalf("overwrite: got %q, want %q", got, "new-bytes")
}
if err := d.Delete(ctx, key); err != nil {
t.Fatalf("delete: %v", err)
}
if _, err := d.Get(ctx, key); !errors.Is(err, ErrNotExist) {
t.Fatalf("get deleted: got %v, want ErrNotExist", err)
}
if err := d.Delete(ctx, key); err != nil {
t.Fatalf("delete missing: %v", err)
}
}
func TestDiskRejectsTraversal(t *testing.T) {
ctx := context.Background()
d, err := NewDisk(t.TempDir())
if err != nil {
t.Fatal(err)
}
if err := d.Put(ctx, "../evil", strings.NewReader("x")); err == nil {
t.Fatal("put with traversal key succeeded")
}
if _, err := d.Get(ctx, "../evil"); err == nil {
t.Fatal("get with traversal key succeeded")
}
}
func TestS3RoundTrip(t *testing.T) {
ctx := context.Background()
client := &mockS3Client{objects: make(map[string][]byte)}
store, err := NewS3(client, "bucket", "prefix")
if err != nil {
t.Fatal(err)
}
if err := store.Put(ctx, "logs/test.log", strings.NewReader("artifact")); err != nil {
t.Fatal(err)
}
rc, err := store.Get(ctx, "logs/test.log")
if err != nil {
t.Fatal(err)
}
got, err := io.ReadAll(rc)
rc.Close()
if err != nil {
t.Fatal(err)
}
if string(got) != "artifact" {
t.Fatalf("content = %q", got)
}
if err := store.Delete(ctx, "logs/test.log"); err != nil {
t.Fatal(err)
}
if _, err := store.Get(ctx, "logs/test.log"); !errors.Is(err, ErrNotExist) {
t.Fatalf("get deleted = %v, want ErrNotExist", err)
}
}