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)) }