summercms/modules/lagoon/attach/store_test.go

package attach_test

import (
	"bytes"
	"context"
	"errors"
	"image"
	"image/png"
	"io"
	"strings"
	"testing"

	"git.golem15.com/golem15/summercms/modules/lagoon"
	"git.golem15.com/golem15/summercms/modules/lagoon/attach"
	"gocloud.dev/blob"
	"gocloud.dev/blob/memblob"
	"gorm.io/gorm"
)

func smokePNG(t *testing.T) []byte {
	t.Helper()
	var buf bytes.Buffer
	if err := png.Encode(&buf, image.NewRGBA(image.Rect(0, 0, 4, 3))); err != nil {
		t.Fatal(err)
	}
	return buf.Bytes()
}

func bucketKeys(t *testing.T, bucket *blob.Bucket) []string {
	t.Helper()
	var keys []string
	iter := bucket.List(nil)
	for {
		obj, err := iter.Next(context.Background())
		if err == io.EOF {
			break
		}
		if err != nil {
			t.Fatal(err)
		}
		keys = append(keys, obj.Key)
	}
	return keys
}

// TestStoreSmoke stores a guarded PNG and a body of exactly MaxBytes.
func TestStoreSmoke(t *testing.T) {
	if testing.Short() {
		t.Skip("requires testcontainers postgres")
	}
	ctx := t.Context()
	gdb := attachGorm(t)
	if err := lagoon.Migrate(gdb, nil); err != nil {
		t.Fatal(err)
	}
	bucket := memblob.OpenBucket(nil)
	t.Cleanup(func() { _ = bucket.Close() })

	data := smokePNG(t)
	f, err := attach.Store(ctx, gdb, bucket, attach.Upload{FileName: `C:\photos\Cover.PNG`, Body: bytes.NewReader(data), Public: false}, attach.Limits{Image: true})
	if err != nil {
		t.Fatal(err)
	}
	if f.ID == 0 || f.SortOrder != int(f.ID) {
		t.Fatalf("sort_order %d, id %d", f.SortOrder, f.ID)
	}
	if !strings.HasSuffix(f.DiskName, ".png") || len(f.DiskName) != 26 {
		t.Fatalf("disk name %q", f.DiskName)
	}
	if f.FileName != "Cover.PNG" || f.ContentType != "image/png" || f.FileSize != int64(len(data)) || f.Public() {
		t.Fatalf("row %+v", f)
	}
	var stored attach.File
	if err := gdb.First(&stored, f.ID).Error; err != nil {
		t.Fatal(err)
	}
	if stored.SortOrder != int(f.ID) || stored.Public() || stored.AttachmentID != "" {
		t.Fatalf("stored row %+v", stored)
	}
	got, err := bucket.ReadAll(ctx, attach.BlobKey(f.DiskName))
	if err != nil || !bytes.Equal(got, data) {
		t.Fatalf("blob %v", err)
	}

	exact := bytes.Repeat([]byte("a"), 64)
	g, err := attach.Store(ctx, gdb, bucket, attach.Upload{FileName: "notes.txt", Body: bytes.NewReader(exact)}, attach.Limits{MaxBytes: 64})
	if err != nil {
		t.Fatalf("exactly MaxBytes: %v", err)
	}
	if g.FileSize != 64 || g.Public() || g.ContentType != "text/plain" {
		t.Fatalf("row %+v", g)
	}
}

// TestStoreSmokeRefusals covers the refusals that happen before any row is
// written: an SVG in image mode and bodies of MaxBytes+1 bytes.
func TestStoreSmokeRefusals(t *testing.T) {
	ctx := t.Context()
	bucket := memblob.OpenBucket(nil)
	t.Cleanup(func() { _ = bucket.Close() })
	db := &gorm.DB{}

	svg := []byte(`<svg xmlns="http://www.w3.org/2000/svg"><script>alert(1)</script></svg>`)
	_, err := attach.Store(ctx, db, bucket, attach.Upload{FileName: "x.png", Body: bytes.NewReader(svg)}, attach.Limits{Image: true})
	if !errors.Is(err, attach.ErrNotImage) {
		t.Fatalf("svg bytes: %v", err)
	}
	_, err = attach.Store(ctx, db, bucket, attach.Upload{FileName: "x.svg", Body: bytes.NewReader(svg)}, attach.Limits{Image: true})
	if !errors.Is(err, attach.ErrFileType) {
		t.Fatalf("svg extension: %v", err)
	}

	_, err = attach.Store(ctx, db, bucket, attach.Upload{FileName: "notes.txt", Body: bytes.NewReader(bytes.Repeat([]byte("a"), 65))}, attach.Limits{MaxBytes: 64})
	if !errors.Is(err, attach.ErrTooLarge) {
		t.Fatalf("small body: %v", err)
	}
	// Past the 1 MiB read-ahead the limit is enforced while streaming.
	const limit = 2 << 20
	_, err = attach.Store(ctx, db, bucket, attach.Upload{FileName: "notes.txt", Body: bytes.NewReader(bytes.Repeat([]byte("a"), limit+1))}, attach.Limits{MaxBytes: limit})
	if !errors.Is(err, attach.ErrTooLarge) {
		t.Fatalf("streamed body: %v", err)
	}
	if keys := bucketKeys(t, bucket); len(keys) != 0 {
		t.Fatalf("blobs left behind: %v", keys)
	}
}