package attachment import ( "context" "fmt" "io" "sync" "heckel.io/ntfy/v2/log" "heckel.io/ntfy/v2/s3" "heckel.io/ntfy/v2/util" ) const ( tagS3Store = "s3_store" ) type s3Store struct { client *s3.Client totalSizeCurrent int64 totalSizeLimit int64 mu sync.Mutex } // NewS3Store creates a new S3-backed attachment store. The s3URL must be in the format: // // s3://ACCESS_KEY:SECRET_KEY@BUCKET[/PREFIX]?region=REGION[&endpoint=ENDPOINT] func NewS3Store(s3URL string, totalSizeLimit int64) (Store, error) { cfg, err := s3.ParseURL(s3URL) if err != nil { return nil, err } store := &s3Store{ client: s3.New(cfg), totalSizeLimit: totalSizeLimit, } if totalSizeLimit > 0 { size, err := store.computeSize() if err != nil { return nil, fmt.Errorf("s3 store: failed to compute initial size: %w", err) } store.totalSizeCurrent = size } return store, nil } func (c *s3Store) Write(id string, in io.Reader, limiters ...util.Limiter) (int64, error) { if !fileIDRegex.MatchString(id) { return 0, errInvalidFileID } log.Tag(tagS3Store).Field("message_id", id).Debug("Writing attachment to S3") // Stream through limiters via an io.Pipe directly to S3. PutObject supports chunked // uploads, so no temp file or Content-Length is needed. limiters = append(limiters, util.NewFixedLimiter(c.Remaining())) pr, pw := io.Pipe() lw := util.NewLimitWriter(pw, limiters...) var size int64 var copyErr error done := make(chan struct{}) go func() { defer close(done) size, copyErr = io.Copy(lw, in) if copyErr != nil { pw.CloseWithError(copyErr) } else { pw.Close() } }() putErr := c.client.PutObject(context.Background(), id, pr) pr.Close() <-done if copyErr != nil { return 0, copyErr } if putErr != nil { return 0, putErr } c.mu.Lock() c.totalSizeCurrent += size c.mu.Unlock() return size, nil } func (c *s3Store) Read(id string) (io.ReadCloser, int64, error) { if !fileIDRegex.MatchString(id) { return nil, 0, errInvalidFileID } return c.client.GetObject(context.Background(), id) } func (c *s3Store) Remove(ids ...string) error { for _, id := range ids { if !fileIDRegex.MatchString(id) { return errInvalidFileID } } // S3 DeleteObjects supports up to 1000 keys per call for i := 0; i < len(ids); i += 1000 { end := i + 1000 if end > len(ids) { end = len(ids) } batch := ids[i:end] for _, id := range batch { log.Tag(tagS3Store).Field("message_id", id).Debug("Deleting attachment from S3") } if err := c.client.DeleteObjects(context.Background(), batch); err != nil { return err } } // Recalculate totalSizeCurrent via ListObjectsV2 (matches fileStore's dirSize rescan pattern) size, err := c.computeSize() if err != nil { return fmt.Errorf("s3 store: failed to compute size after remove: %w", err) } c.mu.Lock() c.totalSizeCurrent = size c.mu.Unlock() return nil } func (c *s3Store) Size() int64 { c.mu.Lock() defer c.mu.Unlock() return c.totalSizeCurrent } func (c *s3Store) Remaining() int64 { c.mu.Lock() defer c.mu.Unlock() remaining := c.totalSizeLimit - c.totalSizeCurrent if remaining < 0 { return 0 } return remaining } // computeSize uses ListAllObjects to sum up the total size of all objects with our prefix. func (c *s3Store) computeSize() (int64, error) { objects, err := c.client.ListAllObjects(context.Background()) if err != nil { return 0, err } var totalSize int64 for _, obj := range objects { totalSize += obj.Size } return totalSize, nil }