BountyBoard/internal/files/store.go

package files

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"net/http"
	"time"

	"go.mongodb.org/mongo-driver/v2/bson"
	"go.mongodb.org/mongo-driver/v2/mongo"
	"go.mongodb.org/mongo-driver/v2/mongo/options"

	"bountyboard/internal/ulid"
)

// ErrTooLarge is returned when an upload exceeds MAX_UPLOAD_MB.
var ErrTooLarge = errors.New("file exceeds maximum upload size")

// ErrNotFound is returned when no stored file has the given id.
var ErrNotFound = errors.New("file not found")

// Valid metadata scopes (§4.7); access control at the HTTP layer keys off
// these.
const (
	ScopeAvatar = "avatar"
	ScopeChat   = "chat"
	ScopeTask   = "task"
)

type Meta struct {
	ID       string
	Name     string
	OwnerID  string
	Scope    string
	MimeType string
	Size     int64
	SHA256   string
}

// Store wraps the GridFS bucket with MIME sniffing, a size cap, and sha256
// recording.
type Store struct {
	bucket   *mongo.GridFSBucket
	maxBytes int64
}

func NewStore(db *mongo.Database, maxUploadMB int) *Store {
	return &Store{
		bucket:   db.GridFSBucket(),
		maxBytes: int64(maxUploadMB) * 1024 * 1024,
	}
}

// Save streams r into GridFS. The content type is sniffed from the first
// 512 bytes; declaredMime is used only when sniffing is inconclusive.
func (s *Store) Save(ctx context.Context, scope, ownerID, name, declaredMime string, r io.Reader) (*Meta, error) {
	head := make([]byte, 512)
	n, err := io.ReadFull(r, head)
	if err != nil && err != io.ErrUnexpectedEOF && err != io.EOF {
		return nil, fmt.Errorf("read upload head: %w", err)
	}
	head = head[:n]

	mimeType := http.DetectContentType(head)
	if mimeType == "application/octet-stream" && declaredMime != "" {
		mimeType = declaredMime
	}

	id := ulid.New()
	hasher := sha256.New()
	// Allow one byte over the cap so an oversize upload is detectable.
	body := &countingReader{r: io.TeeReader(
		io.LimitReader(io.MultiReader(bytes.NewReader(head), r), s.maxBytes+1),
		hasher)}
	err = s.bucket.UploadFromStreamWithID(ctx, id, name, body,
		options.GridFSUpload().SetMetadata(bson.M{
			"ownerId":  ownerID,
			"scope":    scope,
			"mimeType": mimeType,
		}))
	if err != nil {
		return nil, fmt.Errorf("gridfs upload: %w", err)
	}
	size := body.n
	if size > s.maxBytes {
		if delErr := s.bucket.Delete(ctx, id); delErr != nil {
			return nil, fmt.Errorf("delete oversize upload: %w", delErr)
		}
		return nil, fmt.Errorf("%w (limit %d bytes)", ErrTooLarge, s.maxBytes)
	}

	sum := hex.EncodeToString(hasher.Sum(nil))
	_, err = s.bucket.GetFilesCollection().UpdateOne(ctx,
		bson.M{"_id": id},
		bson.M{"$set": bson.M{"metadata.sha256": sum}})
	if err != nil {
		return nil, fmt.Errorf("record sha256: %w", err)
	}

	return &Meta{
		ID: id, Name: name, OwnerID: ownerID, Scope: scope,
		MimeType: mimeType, Size: size, SHA256: sum,
	}, nil
}

// Open returns a reader over the stored file plus its metadata. The caller
// must Close the stream.
func (s *Store) Open(ctx context.Context, id string) (io.ReadCloser, *Meta, error) {
	ds, err := s.bucket.OpenDownloadStream(ctx, id)
	if err != nil {
		if errors.Is(err, mongo.ErrFileNotFound) {
			return nil, nil, fmt.Errorf("%s: %w", id, ErrNotFound)
		}
		return nil, nil, fmt.Errorf("gridfs open %s: %w", id, err)
	}
	f := ds.GetFile()
	var md struct {
		OwnerID  string `bson:"ownerId"`
		Scope    string `bson:"scope"`
		MimeType string `bson:"mimeType"`
		SHA256   string `bson:"sha256"`
	}
	if f.Metadata != nil {
		if err := bson.Unmarshal(f.Metadata, &md); err != nil {
			ds.Close()
			return nil, nil, fmt.Errorf("decode file metadata %s: %w", id, err)
		}
	}
	return ds, &Meta{
		ID: id, Name: f.Name, OwnerID: md.OwnerID, Scope: md.Scope,
		MimeType: md.MimeType, Size: f.Length, SHA256: md.SHA256,
	}, nil
}

// Delete removes the file and its chunks.
func (s *Store) Delete(ctx context.Context, id string) error {
	if err := s.bucket.Delete(ctx, id); err != nil {
		if errors.Is(err, mongo.ErrFileNotFound) {
			return fmt.Errorf("%s: %w", id, ErrNotFound)
		}
		return fmt.Errorf("gridfs delete %s: %w", id, err)
	}
	return nil
}

// countingReader tracks how many bytes the upload actually consumed, since
// the driver's upload helper does not report a size.
type countingReader struct {
	r io.Reader
	n int64
}

func (c *countingReader) Read(p []byte) (int, error) {
	n, err := c.r.Read(p)
	c.n += int64(n)
	return n, err
}

// SignedURLFor is a convenience for handing attachments to external services.
func (s *Store) SignedURLFor(appBaseURL string, secret []byte, fileID string) string {
	return SignedURL(appBaseURL, secret, fileID, time.Now().Add(DefaultTokenTTL))
}