Warp/server/internal/storage/pocketbase/blobs.go

package pocketbase

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"strconv"

	"github.com/dariuszjaniak/warp/server/internal/storage"
)

// mediaCollection holds media content: one record per distinct content,
// keyed by its SHA-256, with the bytes in a protected file field. Where
// PocketBase keeps the file — its pb_data/storage or an S3 bucket — is its
// own storage setting.
const mediaCollection = "warp_media"

// blobStore is storage.BlobStore over warp_media.
type blobStore struct{ *base }

// validSum reports whether s is a lower-case hex SHA-256: the only names
// content is stored under, and so the only ones worth asking about.
func validSum(s string) bool {
	if len(s) != sha256.Size*2 {
		return false
	}
	for _, c := range s {
		if (c < '0' || c > '9') && (c < 'a' || c > 'f') {
			return false
		}
	}
	return true
}

// Put implements storage.BlobStore.
//
// The content is spooled to a temporary file first: its name is its hash,
// which is not known until the last byte has been read, and the upload has
// to say both. The spool is also what keeps an oversize upload from reaching
// PocketBase at all.
func (b *blobStore) Put(ctx context.Context, r io.Reader, limit int64) (storage.BlobInfo, error) {
	tmp, err := os.CreateTemp("", "warp-media-*")
	if err != nil {
		return storage.BlobInfo{}, fmt.Errorf("create a temporary file: %w", err)
	}
	defer func() {
		_ = tmp.Close()
		_ = os.Remove(tmp.Name())
	}()

	hash := sha256.New()
	n, err := io.Copy(io.MultiWriter(tmp, hash), &ctxReader{ctx: ctx, r: io.LimitReader(r, limit+1)})
	if err != nil {
		return storage.BlobInfo{}, fmt.Errorf("store content: %w", err)
	}
	if n > limit {
		return storage.BlobInfo{}, storage.ErrBlobTooLarge
	}
	info := storage.BlobInfo{SHA256: hex.EncodeToString(hash.Sum(nil)), Size: n}
	id := rowID(mediaCollection, info.SHA256)

	if _, err := b.c.getRecord(ctx, mediaCollection, id); err == nil {
		// The same content is already stored.
		return info, nil
	} else if !isNotFound(err) {
		return storage.BlobInfo{}, fmt.Errorf("look up stored content: %w", err)
	}

	if _, err := tmp.Seek(0, io.SeekStart); err != nil {
		return storage.BlobInfo{}, fmt.Errorf("store content: %w", err)
	}
	err = b.c.createRecordWithFile(ctx, mediaCollection, map[string]string{
		"id":     id,
		"sha256": info.SHA256,
		"size":   strconv.FormatInt(n, 10),
	}, "content", info.SHA256, tmp)
	if err != nil {
		// Two uploads of the same content can race to create the record;
		// the loser finds the winner's.
		if _, getErr := b.c.getRecord(ctx, mediaCollection, id); getErr == nil {
			return info, nil
		}
		return storage.BlobInfo{}, fmt.Errorf("store content: %w", err)
	}
	return info, nil
}

// Open implements storage.BlobStore.
func (b *blobStore) Open(ctx context.Context, sum string) (io.ReadSeekCloser, error) {
	if !validSum(sum) {
		return nil, storage.ErrNoSuchBlob
	}
	id := rowID(mediaCollection, sum)
	record, err := b.c.getRecord(ctx, mediaCollection, id)
	switch {
	case isNotFound(err):
		return nil, storage.ErrNoSuchBlob
	case err != nil:
		return nil, fmt.Errorf("open media content: %w", err)
	}
	name, _ := record["content"].(string)
	size, _ := record["size"].(float64)
	if name == "" {
		return nil, storage.ErrNoSuchBlob
	}
	token, err := b.c.fileToken(ctx)
	if err != nil {
		return nil, fmt.Errorf("open media content: %w", err)
	}
	return &blobReader{ctx: ctx, c: b.c, id: id, name: name, token: token, size: int64(size)}, nil
}

// Delete implements storage.BlobStore.
func (b *blobStore) Delete(ctx context.Context, sum string) error {
	if !validSum(sum) {
		return nil
	}
	if err := b.c.deleteRecord(ctx, mediaCollection, rowID(mediaCollection, sum)); err != nil && !isNotFound(err) {
		return fmt.Errorf("delete media content: %w", err)
	}
	return nil
}

// blobReader reads a stored file over HTTP. A seek only moves the offset;
// the next read starts a ranged download from there, so reading a file
// through costs one request however it is read.
type blobReader struct {
	ctx   context.Context
	c     *client
	id    string
	name  string
	token string
	size  int64

	offset int64
	body   io.ReadCloser
}

func (r *blobReader) Read(p []byte) (int, error) {
	if r.offset >= r.size {
		return 0, io.EOF
	}
	if r.body == nil {
		resp, err := r.c.openFile(r.ctx, mediaCollection, r.id, r.name, r.token, r.offset)
		if err != nil {
			return 0, err
		}
		if r.offset > 0 && resp.StatusCode != http.StatusPartialContent {
			_ = resp.Body.Close()
			return 0, fmt.Errorf("pocketbase: the file server ignored a range request (HTTP %d)", resp.StatusCode)
		}
		r.body = resp.Body
	}
	n, err := r.body.Read(p)
	r.offset += int64(n)
	if errors.Is(err, io.EOF) && r.offset < r.size {
		err = io.ErrUnexpectedEOF
	}
	return n, err
}

func (r *blobReader) Seek(offset int64, whence int) (int64, error) {
	var next int64
	switch whence {
	case io.SeekStart:
		next = offset
	case io.SeekCurrent:
		next = r.offset + offset
	case io.SeekEnd:
		next = r.size + offset
	default:
		return 0, errors.New("pocketbase: invalid whence")
	}
	if next < 0 {
		return 0, errors.New("pocketbase: negative position")
	}
	if next != r.offset && r.body != nil {
		_ = r.body.Close()
		r.body = nil
	}
	r.offset = next
	return next, nil
}

func (r *blobReader) Close() error {
	if r.body != nil {
		return r.body.Close()
	}
	return nil
}

// ctxReader stops a copy when its context ends.
type ctxReader struct {
	ctx context.Context
	r   io.Reader
}

func (c *ctxReader) Read(p []byte) (int, error) {
	if err := c.ctx.Err(); err != nil {
		return 0, err
	}
	return c.r.Read(p)
}

var _ storage.BlobStore = (*blobStore)(nil)