ent/entc/integration/blob_test.go

// Copyright 2019-present Facebook Inc. All rights reserved.
// This source code is licensed under the Apache 2.0 license found
// in the LICENSE file in the root directory of this source tree.

package integration

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"entgo.io/ent/dialect"
	"entgo.io/ent/entc/integration/blob"
	"entgo.io/ent/entc/integration/ent"
	"entgo.io/ent/entc/integration/ent/document"
	"entgo.io/ent/entc/integration/ent/enttest"
	"entgo.io/ent/entc/integration/ent/migrate"

	_ "github.com/mattn/go-sqlite3"
	"github.com/stretchr/testify/require"
	_ "gocloud.dev/blob/fileblob"
)

// readBlob reads all data from a blob field reader and closes it.
func readBlob(t *testing.T, rc io.ReadCloser, err error) []byte {
	t.Helper()
	require.NoError(t, err)
	if rc == nil {
		return nil
	}
	data, readErr := io.ReadAll(rc)
	require.NoError(t, readErr)
	require.NoError(t, rc.Close())
	return data
}

// blobContent is a shorthand for readBlob(t, entity.Field(ctx)).
func blobContent(t *testing.T, fn func(context.Context) (io.ReadCloser, error), ctx context.Context) []byte {
	t.Helper()
	rc, err := fn(ctx)
	return readBlob(t, rc, err)
}

// blobDir creates a temp directory with subdirectories for blob fields.
func blobDir(t *testing.T) string {
	t.Helper()
	dir := t.TempDir()
	require.NoError(t, os.MkdirAll(filepath.Join(dir, "documents"), 0o755))
	require.NoError(t, os.MkdirAll(filepath.Join(dir, "thumbnails"), 0o755))
	return dir
}

// newBlobOpeners returns BlobOpeners that route to the correct
// file:// bucket based on the field name, using absolute paths under dir.
func newBlobOpeners(dir string) ent.BlobOpeners {
	return ent.BlobOpeners{
		Document: func(ctx context.Context, field string) (ent.Blob, error) {
			switch field {
			case document.FieldContent:
				return blob.OpenBucket(ctx, "file://"+filepath.Join(dir, "documents"))
			case document.FieldThumbnail:
				return blob.OpenBucket(ctx, "file://"+filepath.Join(dir, "thumbnails"))
			default:
				return nil, fmt.Errorf("unknown blob field: %s", field)
			}
		},
	}
}

// setupBlob creates a temp directory with subdirectories for each blob field,
// opens an in-memory SQLite client with auto-migration, and registers cleanup.
func setupBlob(t *testing.T, opts ...ent.Option) (*ent.Client, context.Context, string) {
	t.Helper()
	dir := blobDir(t)
	allOpts := append([]ent.Option{ent.WithBlobOpeners(newBlobOpeners(dir))}, opts...)
	entOpts := []enttest.Option{
		enttest.WithMigrateOptions(migrate.WithDropIndex(true), migrate.WithDropColumn(true)),
		enttest.WithOptions(allOpts...),
	}
	client := enttest.Open(t, dialect.SQLite, "file:ent?mode=memory&cache=shared&_fk=1",
		entOpts...,
	)
	t.Cleanup(func() {
		client.Close()
	})
	return client, context.Background(), dir
}

func TestBlobCreateAndRead(t *testing.T) {
	client, ctx, _ := setupBlob(t)

	data := []byte("Hello from blob integration test!")
	doc := client.Document.Create().
		SetName("test-doc").
		SetContent(bytes.NewReader(data)).
		SetThumbnail(bytes.NewReader([]byte("thumb"))).
		SaveX(ctx)

	// Read the blob back through the entity method.
	got := blobContent(t, doc.ContentReader, ctx)
	require.Equal(t, data, got)
}

func TestBlobQueryAndRead(t *testing.T) {
	client, ctx, _ := setupBlob(t)

	data := []byte("queried blob content")
	created := client.Document.Create().
		SetName("query-doc").
		SetContent(bytes.NewReader(data)).
		SetThumbnail(bytes.NewReader([]byte("thumb"))).
		SaveX(ctx)

	// Query it back from the database (no content column, just ID/name).
	queried := client.Document.GetX(ctx, created.ID)

	// Read blob content from the queried entity.
	got := blobContent(t, queried.ContentReader, ctx)
	require.Equal(t, data, got)
}

func TestBlobUpdateData(t *testing.T) {
	client, ctx, _ := setupBlob(t)

	v1 := []byte("version 1")
	doc := client.Document.Create().
		SetName("update-doc").
		SetContent(bytes.NewReader(v1)).
		SetThumbnail(bytes.NewReader([]byte("thumb"))).
		SaveX(ctx)

	// Update blob data through mutation (overwrites same key).
	v2 := []byte("version 2 - updated via mutation")
	doc = doc.Update().
		SetContent(bytes.NewReader(v2)).
		SaveX(ctx)

	// Read the new blob content.
	got := blobContent(t, doc.ContentReader, ctx)
	require.Equal(t, v2, got)
}

func TestBlobRequiredValidation(t *testing.T) {
	client, ctx, _ := setupBlob(t)

	// Creating a document without required blob fields should fail.
	_, err := client.Document.Create().
		SetName("no-blob-doc").
		Save(ctx)
	require.Error(t, err)
	require.Contains(t, err.Error(), "missing required field")
}

func TestBlobMultipleDocuments(t *testing.T) {
	client, ctx, _ := setupBlob(t)

	contents := []string{
		"document one content",
		"document two content - larger payload with more data",
		"document three",
	}

	var ids []int
	for i, c := range contents {
		doc := client.Document.Create().
			SetName("multi-" + string(rune('a'+i))).
			SetContent(strings.NewReader(c)).
			SetThumbnail(bytes.NewReader([]byte("thumb"))).
			SaveX(ctx)
		ids = append(ids, doc.ID)
	}

	// Read each document and verify content is correct.
	for i, id := range ids {
		doc := client.Document.GetX(ctx, id)
		got := blobContent(t, doc.ContentReader, ctx)
		require.Equal(t, contents[i], string(got), "document %d content mismatch", i)
	}
}

func TestBlobBulkCreate(t *testing.T) {
	client, ctx, _ := setupBlob(t)

	bulk := make([]*ent.DocumentCreate, 5)
	for i := range bulk {
		bulk[i] = client.Document.Create().
			SetName(strings.Repeat("bulk-", i+1)).
			SetContent(bytes.NewReader([]byte(strings.Repeat("x", i+1)))).
			SetThumbnail(bytes.NewReader([]byte("thumb")))
	}
	docs, err := client.Document.CreateBulk(bulk...).Save(ctx)
	require.NoError(t, err)
	require.Len(t, docs, 5)

	// Verify each document's blob can be read back with correct data.
	for i, doc := range docs {
		got := blobContent(t, doc.ContentReader, ctx)
		require.Equal(t, []byte(strings.Repeat("x", i+1)), got)
	}
}

func TestBlobThumbnailCreateAndRead(t *testing.T) {
	client, ctx, _ := setupBlob(t)

	thumbData := []byte("fake-png-thumbnail-data")
	doc := client.Document.Create().
		SetName("doc-with-thumb").
		SetContent(bytes.NewReader([]byte("content"))).
		SetThumbnail(bytes.NewReader(thumbData)).
		SaveX(ctx)

	got := blobContent(t, doc.ThumbnailReader, ctx)
	require.Equal(t, thumbData, got)
}

func TestBlobBothFields(t *testing.T) {
	client, ctx, _ := setupBlob(t)

	contentData := []byte("document body content")
	thumbData := []byte("thumbnail image bytes")
	doc := client.Document.Create().
		SetName("doc-both").
		SetContent(bytes.NewReader(contentData)).
		SetThumbnail(bytes.NewReader(thumbData)).
		SaveX(ctx)

	// Read content.
	cGot := blobContent(t, doc.ContentReader, ctx)
	require.Equal(t, contentData, cGot)

	// Read thumbnail.
	tGot := blobContent(t, doc.ThumbnailReader, ctx)
	require.Equal(t, thumbData, tGot)

	// Update only thumbnail, content should remain unchanged.
	newThumb := []byte("updated thumbnail")
	doc = doc.Update().
		SetThumbnail(bytes.NewReader(newThumb)).
		SaveX(ctx)

	cGot2 := blobContent(t, doc.ContentReader, ctx)
	require.Equal(t, contentData, cGot2)

	tGot2 := blobContent(t, doc.ThumbnailReader, ctx)
	require.Equal(t, newThumb, tGot2)
}

func TestBlobBulkCreateBothFields(t *testing.T) {
	client, ctx, _ := setupBlob(t)

	bulk := make([]*ent.DocumentCreate, 3)
	for i := range bulk {
		bulk[i] = client.Document.Create().
			SetName(strings.Repeat("bulk-both-", i+1)).
			SetContent(bytes.NewReader([]byte(strings.Repeat("c", i+1)))).
			SetThumbnail(bytes.NewReader([]byte(strings.Repeat("t", i+1))))
	}
	docs, err := client.Document.CreateBulk(bulk...).Save(ctx)
	require.NoError(t, err)
	require.Len(t, docs, 3)

	for i, doc := range docs {
		cGot := blobContent(t, doc.ContentReader, ctx)
		require.Equal(t, []byte(strings.Repeat("c", i+1)), cGot)

		tGot := blobContent(t, doc.ThumbnailReader, ctx)
		require.Equal(t, []byte(strings.Repeat("t", i+1)), tGot)
	}
}

func TestBlobWriter(t *testing.T) {
	client, ctx, _ := setupBlob(t)

	// Create a document with content via mutation.
	doc := client.Document.Create().
		SetName("writer-doc").
		SetContent(strings.NewReader("initial")).
		SetThumbnail(bytes.NewReader([]byte("thumb"))).
		SaveX(ctx)

	// Overwrite via ContentWriter (bypasses mutation pipeline).
	w, err := doc.ContentWriter(ctx)
	require.NoError(t, err)
	_, err = io.Copy(w, strings.NewReader("overwritten via writer"))
	require.NoError(t, err)
	require.NoError(t, w.Close())

	// Read back should reflect the overwritten data.
	got := blobContent(t, doc.ContentReader, ctx)
	require.Equal(t, []byte("overwritten via writer"), got)
}

func TestBlobReader(t *testing.T) {
	client, ctx, _ := setupBlob(t)

	data := []byte("streaming read test data")
	doc := client.Document.Create().
		SetName("reader-doc").
		SetContent(bytes.NewReader(data)).
		SetThumbnail(bytes.NewReader([]byte("thumb"))).
		SaveX(ctx)

	// Read back via Content (io.ReadCloser).
	r, err := doc.ContentReader(ctx)
	require.NoError(t, err)
	defer r.Close()

	got, err := io.ReadAll(r)
	require.NoError(t, err)
	require.Equal(t, data, got)
}

func TestBlobWriterThenReader(t *testing.T) {
	client, ctx, _ := setupBlob(t)

	// Create a document with initial content, then overwrite via Writer.
	doc := client.Document.Create().
		SetName("write-then-read").
		SetContent(strings.NewReader("initial")).
		SetThumbnail(bytes.NewReader([]byte("thumb"))).
		SaveX(ctx)

	// Write via writer.
	w, err := doc.ContentWriter(ctx)
	require.NoError(t, err)
	_, err = w.Write([]byte("hello "))
	require.NoError(t, err)
	_, err = w.Write([]byte("world"))
	require.NoError(t, err)
	require.NoError(t, w.Close())

	// Read via reader.
	got := blobContent(t, doc.ContentReader, ctx)
	require.Equal(t, []byte("hello world"), got)
}

func TestBlobWriterOptions(t *testing.T) {
	// Verify that a custom opener with WriterOptions works for roundtrip.
	dir := blobDir(t)
	openerWithOpts := ent.BlobOpeners{
		Document: func(ctx context.Context, field string) (ent.Blob, error) {
			b, err := blob.OpenBucket(ctx, "file://"+filepath.Join(dir, "documents"))
			if err != nil {
				return nil, err
			}
			return b.WithWriterOptions(nil), nil
		},
	}
	client, ctx, _ := setupBlob(t, ent.WithBlobOpeners(openerWithOpts))

	data := []byte("content with writer options")
	doc := client.Document.Create().
		SetName("opts-doc").
		SetContent(bytes.NewReader(data)).
		SetThumbnail(bytes.NewReader([]byte("thumb"))).
		SaveX(ctx)

	got := blobContent(t, doc.ContentReader, ctx)
	require.Equal(t, data, got)
}

func TestBlobWriterOptionsApplied(t *testing.T) {
	// Verify that WriterOptions are applied on both create and update.
	dir := blobDir(t)
	openerWithOpts := ent.BlobOpeners{
		Document: func(ctx context.Context, field string) (ent.Blob, error) {
			b, err := blob.OpenBucket(ctx, "file://"+filepath.Join(dir, "documents"))
			if err != nil {
				return nil, err
			}
			return b.WithWriterOptions(nil), nil
		},
	}
	client, ctx, _ := setupBlob(t, ent.WithBlobOpeners(openerWithOpts))

	data := []byte("content with writer options")
	doc := client.Document.Create().
		SetName("opts-doc").
		SetContent(bytes.NewReader(data)).
		SetThumbnail(bytes.NewReader([]byte("thumb"))).
		SaveX(ctx)

	// Verify the data was written successfully and can be read back.
	got := blobContent(t, doc.ContentReader, ctx)
	require.Equal(t, data, got)

	// Update also uses WriterOpts.
	v2 := []byte("updated with writer options")
	doc = doc.Update().SetContent(bytes.NewReader(v2)).SaveX(ctx)
	got = blobContent(t, doc.ContentReader, ctx)
	require.Equal(t, v2, got)
}

func TestBlobEncryption(t *testing.T) {
	// Demonstrate per-tenant encryption using a master seed + tenant from context.
	// Each tenant derives a unique AES-256 key via SHA-256(tenant || seed).
	// Data written by one tenant cannot be decrypted by another.
	dir := blobDir(t)

	// Master seed shared across all tenants (kept secret server-side).
	masterSeed := []byte("super-secret-master-seed-for-test")

	encryptedOpeners := ent.BlobOpeners{
		Document: func(ctx context.Context, field string) (ent.Blob, error) {
			var subdir string
			switch field {
			case document.FieldContent:
				subdir = "documents"
			case document.FieldThumbnail:
				subdir = "thumbnails"
			default:
				return nil, fmt.Errorf("unknown blob field: %s", field)
			}
			b, err := blob.OpenBucket(ctx, "file://"+filepath.Join(dir, subdir))
			if err != nil {
				return nil, err
			}
			// Wrap the bucket with per-tenant encryption.
			return blob.NewEncrypted(b, masterSeed), nil
		},
	}

	// Tenant "acme" writes encrypted data.
	client, _, _ := setupBlob(t, ent.WithBlobOpeners(encryptedOpeners))
	acmeCtx := blob.WithTenant(context.Background(), "acme")

	plaintext := []byte("top secret document content for acme")
	doc := client.Document.Create().
		SetName("encrypted-doc").
		SetContent(bytes.NewReader(plaintext)).
		SetThumbnail(bytes.NewReader([]byte("secret-thumb"))).
		SaveX(acmeCtx)

	// Reading with the same tenant returns decrypted plaintext.
	got := blobContent(t, doc.ContentReader, acmeCtx)
	require.Equal(t, plaintext, got)

	gotThumb := blobContent(t, doc.ThumbnailReader, acmeCtx)
	require.Equal(t, []byte("secret-thumb"), gotThumb)

	// Verify the raw file on disk is NOT plaintext (it's encrypted).
	matches, err := filepath.Glob(filepath.Join(dir, "documents", "*", "content"))
	require.NoError(t, err)
	require.Len(t, matches, 1)
	raw, err := os.ReadFile(matches[0])
	require.NoError(t, err)
	require.NotEqual(t, plaintext, raw, "raw data on disk should be encrypted")
	// The raw data should be longer than plaintext (16-byte IV prepended).
	require.Equal(t, len(plaintext)+16, len(raw))

	// A different tenant ("evil") cannot decrypt acme's data — the derived key
	// differs, so AES-CTR produces garbage (not the original plaintext).
	evilCtx := blob.WithTenant(context.Background(), "evil")
	evilData := blobContent(t, doc.ContentReader, evilCtx)
	require.NotEqual(t, plaintext, evilData, "different tenant must not read acme's plaintext")

	// No tenant in context → error.
	noTenantCtx := context.Background()
	_, err = doc.ContentReader(noTenantCtx)
	require.Error(t, err)
	require.Contains(t, err.Error(), "requires a tenant")

	// Update works through encryption too.
	v2 := []byte("updated secret content")
	doc = doc.Update().SetContent(bytes.NewReader(v2)).SaveX(acmeCtx)
	got = blobContent(t, doc.ContentReader, acmeCtx)
	require.Equal(t, v2, got)

	// ContentWriter also encrypts with tenant key.
	w, err := doc.ContentWriter(acmeCtx)
	require.NoError(t, err)
	_, err = w.Write([]byte("written via writer"))
	require.NoError(t, err)
	require.NoError(t, w.Close())

	got = blobContent(t, doc.ContentReader, acmeCtx)
	require.Equal(t, []byte("written via writer"), got)
}

func TestBlobPrefix(t *testing.T) {
	dir := blobDir(t)
	openers := newBlobOpeners(dir)
	prefixedOpeners := ent.BlobOpeners{
		Document: func(ctx context.Context, field string) (ent.Blob, error) {
			if field == "content" {
				b, err := blob.OpenBucket(ctx, "file://"+filepath.Join(dir, "documents"))
				if err != nil {
					return nil, err
				}
				return b.Prefixed("tenant-1/"), nil
			}
			return openers.Document(ctx, field)
		},
	}
	client, ctx, _ := setupBlob(t, ent.WithBlobOpeners(prefixedOpeners))

	data := []byte("prefixed blob content")
	doc := client.Document.Create().
		SetName("prefixed-doc").
		SetContent(bytes.NewReader(data)).
		SetThumbnail(bytes.NewReader([]byte("thumb"))).
		SaveX(ctx)

	// Read through the entity — uses the same opener.
	got := blobContent(t, doc.ContentReader, ctx)
	require.Equal(t, data, got)

	// Create a second client with the default opener — reading should return nil (not found).
	defaultClient := enttest.Open(t, dialect.SQLite, "file:ent?mode=memory&cache=shared&_fk=1",
		enttest.WithOptions(ent.WithBlobOpeners(openers)),
	)
	t.Cleanup(func() { defaultClient.Close() })
	doc2 := defaultClient.Document.Query().OnlyX(ctx)
	got = blobContent(t, doc2.ContentReader, ctx)
	require.Nil(t, got, "expected nil when reading without prefix")
}

func TestBlobPrefixUpdate(t *testing.T) {
	dir := blobDir(t)
	openers := newBlobOpeners(dir)
	prefixedOpeners := ent.BlobOpeners{
		Document: func(ctx context.Context, field string) (ent.Blob, error) {
			if field == "content" {
				b, err := blob.OpenBucket(ctx, "file://"+filepath.Join(dir, "documents"))
				if err != nil {
					return nil, err
				}
				return b.Prefixed("tenant-2/"), nil
			}
			return openers.Document(ctx, field)
		},
	}
	client, ctx, _ := setupBlob(t, ent.WithBlobOpeners(prefixedOpeners))

	data := []byte("initial")
	doc := client.Document.Create().
		SetName("prefix-update-doc").
		SetContent(bytes.NewReader(data)).
		SetThumbnail(bytes.NewReader([]byte("thumb"))).
		SaveX(ctx)

	updated := []byte("updated under prefix")
	doc = doc.Update().SetContent(bytes.NewReader(updated)).SaveX(ctx)

	got := blobContent(t, doc.ContentReader, ctx)
	require.Equal(t, updated, got)
}

func TestBlobPrefixBulkCreate(t *testing.T) {
	dir := blobDir(t)
	openers := newBlobOpeners(dir)
	prefixedOpeners := ent.BlobOpeners{
		Document: func(ctx context.Context, field string) (ent.Blob, error) {
			if field == "content" {
				b, err := blob.OpenBucket(ctx, "file://"+filepath.Join(dir, "documents"))
				if err != nil {
					return nil, err
				}
				return b.Prefixed("bulk-tenant/"), nil
			}
			return openers.Document(ctx, field)
		},
	}
	client, ctx, _ := setupBlob(t, ent.WithBlobOpeners(prefixedOpeners))

	docs := client.Document.CreateBulk(
		client.Document.Create().SetName("bulk-p1").SetContent(strings.NewReader("b1")).SetThumbnail(bytes.NewReader([]byte("t1"))),
		client.Document.Create().SetName("bulk-p2").SetContent(strings.NewReader("b2")).SetThumbnail(bytes.NewReader([]byte("t2"))),
	).SaveX(ctx)

	for i, doc := range docs {
		got := blobContent(t, doc.ContentReader, ctx)
		require.Equal(t, []byte(fmt.Sprintf("b%d", i+1)), got)
	}
}

func TestBlobPrefixWriterReader(t *testing.T) {
	dir := blobDir(t)
	openers := newBlobOpeners(dir)
	prefixedOpeners := ent.BlobOpeners{
		Document: func(ctx context.Context, field string) (ent.Blob, error) {
			if field == "content" {
				b, err := blob.OpenBucket(ctx, "file://"+filepath.Join(dir, "documents"))
				if err != nil {
					return nil, err
				}
				return b.Prefixed("rw-tenant/"), nil
			}
			return openers.Document(ctx, field)
		},
	}
	client, ctx, _ := setupBlob(t, ent.WithBlobOpeners(prefixedOpeners))

	doc := client.Document.Create().
		SetName("prefix-rw-doc").
		SetContent(strings.NewReader("initial")).
		SetThumbnail(bytes.NewReader([]byte("thumb"))).
		SaveX(ctx)

	// Write via ContentWriter.
	w, err := doc.ContentWriter(ctx)
	require.NoError(t, err)
	_, err = io.Copy(w, strings.NewReader("writer-prefixed"))
	require.NoError(t, err)
	require.NoError(t, w.Close())

	// Read via Content.
	got := blobContent(t, doc.ContentReader, ctx)
	require.Equal(t, []byte("writer-prefixed"), got)
}