mirror of
https://github.com/multipleof4/ntfy.git
synced 2026-10-17 17:07:53 +00:00
S3 WIP
This commit is contained in:
binwiederhier
+42
-144
@@ -4,24 +4,20 @@ import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/url"
|
||||
"strings"
|
||||
"os"
|
||||
"sync"
|
||||
|
||||
"github.com/aws/aws-sdk-go-v2/aws"
|
||||
"github.com/aws/aws-sdk-go-v2/credentials"
|
||||
"github.com/aws/aws-sdk-go-v2/service/s3"
|
||||
s3types "github.com/aws/aws-sdk-go-v2/service/s3/types"
|
||||
"heckel.io/ntfy/v2/log"
|
||||
"heckel.io/ntfy/v2/s3"
|
||||
"heckel.io/ntfy/v2/util"
|
||||
)
|
||||
|
||||
const tagS3Store = "s3_store"
|
||||
const (
|
||||
tagS3Store = "s3_store"
|
||||
)
|
||||
|
||||
type s3Store struct {
|
||||
client *s3.Client
|
||||
bucket string
|
||||
prefix string
|
||||
totalSizeCurrent int64
|
||||
totalSizeLimit int64
|
||||
mu sync.Mutex
|
||||
@@ -31,14 +27,12 @@ type s3Store struct {
|
||||
//
|
||||
// s3://ACCESS_KEY:SECRET_KEY@BUCKET[/PREFIX]?region=REGION[&endpoint=ENDPOINT]
|
||||
func NewS3Store(s3URL string, totalSizeLimit int64) (Store, error) {
|
||||
bucket, prefix, client, err := parseS3URL(s3URL)
|
||||
cfg, err := s3.ParseURL(s3URL)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
store := &s3Store{
|
||||
client: client,
|
||||
bucket: bucket,
|
||||
prefix: prefix,
|
||||
client: s3.New(cfg),
|
||||
totalSizeLimit: totalSizeLimit,
|
||||
}
|
||||
if totalSizeLimit > 0 {
|
||||
@@ -51,98 +45,40 @@ func NewS3Store(s3URL string, totalSizeLimit int64) (Store, error) {
|
||||
return store, nil
|
||||
}
|
||||
|
||||
func parseS3URL(s3URL string) (bucket string, prefix string, client *s3.Client, err error) {
|
||||
u, err := url.Parse(s3URL)
|
||||
if err != nil {
|
||||
return "", "", nil, fmt.Errorf("s3 store: invalid URL: %w", err)
|
||||
}
|
||||
if u.Scheme != "s3" {
|
||||
return "", "", nil, fmt.Errorf("s3 store: URL scheme must be 's3', got '%s'", u.Scheme)
|
||||
}
|
||||
if u.Host == "" {
|
||||
return "", "", nil, fmt.Errorf("s3 store: bucket name must be specified as host")
|
||||
}
|
||||
bucket = u.Host
|
||||
prefix = strings.TrimPrefix(u.Path, "/")
|
||||
|
||||
accessKey := u.User.Username()
|
||||
secretKey, _ := u.User.Password()
|
||||
if accessKey == "" || secretKey == "" {
|
||||
return "", "", nil, fmt.Errorf("s3 store: access key and secret key must be specified in URL")
|
||||
}
|
||||
|
||||
region := u.Query().Get("region")
|
||||
if region == "" {
|
||||
return "", "", nil, fmt.Errorf("s3 store: region query parameter is required")
|
||||
}
|
||||
endpoint := u.Query().Get("endpoint")
|
||||
|
||||
cfg := aws.Config{
|
||||
Region: region,
|
||||
Credentials: credentials.NewStaticCredentialsProvider(accessKey, secretKey, ""),
|
||||
}
|
||||
var opts []func(*s3.Options)
|
||||
if endpoint != "" {
|
||||
opts = append(opts, func(o *s3.Options) {
|
||||
o.BaseEndpoint = aws.String(endpoint)
|
||||
o.UsePathStyle = true
|
||||
})
|
||||
}
|
||||
client = s3.NewFromConfig(cfg, opts...)
|
||||
return bucket, prefix, client, nil
|
||||
}
|
||||
|
||||
func (c *s3Store) objectKey(id string) string {
|
||||
if c.prefix != "" {
|
||||
return c.prefix + "/" + id
|
||||
}
|
||||
return id
|
||||
}
|
||||
|
||||
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")
|
||||
|
||||
// Use io.Pipe so we can apply limiters while streaming to S3
|
||||
pr, pw := io.Pipe()
|
||||
var writeErr error
|
||||
var size int64
|
||||
|
||||
// Write through limiters into a temp file. This avoids buffering the full attachment in
|
||||
// memory while still giving us the Content-Length that PutObject requires.
|
||||
limiters = append(limiters, util.NewFixedLimiter(c.Remaining()))
|
||||
go func() {
|
||||
limitWriter := util.NewLimitWriter(pw, limiters...)
|
||||
size, writeErr = io.Copy(limitWriter, in)
|
||||
if writeErr != nil {
|
||||
pw.CloseWithError(writeErr)
|
||||
} else {
|
||||
pw.Close()
|
||||
}
|
||||
}()
|
||||
|
||||
key := c.objectKey(id)
|
||||
_, err := c.client.PutObject(context.Background(), &s3.PutObjectInput{
|
||||
Bucket: aws.String(c.bucket),
|
||||
Key: aws.String(key),
|
||||
Body: pr,
|
||||
})
|
||||
tmpFile, err := os.CreateTemp("", "ntfy-s3-upload-*")
|
||||
if err != nil {
|
||||
// If the limiter caused the error, return the original write error
|
||||
if writeErr != nil {
|
||||
return 0, writeErr
|
||||
}
|
||||
return 0, fmt.Errorf("s3 store: PutObject failed: %w", err)
|
||||
return 0, fmt.Errorf("s3 store: failed to create temp file: %w", err)
|
||||
}
|
||||
if writeErr != nil {
|
||||
// The write goroutine failed but PutObject somehow succeeded; clean up
|
||||
_, _ = c.client.DeleteObject(context.Background(), &s3.DeleteObjectInput{
|
||||
Bucket: aws.String(c.bucket),
|
||||
Key: aws.String(key),
|
||||
})
|
||||
return 0, writeErr
|
||||
tmpPath := tmpFile.Name()
|
||||
defer os.Remove(tmpPath)
|
||||
limitWriter := util.NewLimitWriter(tmpFile, limiters...)
|
||||
size, err := io.Copy(limitWriter, in)
|
||||
if err != nil {
|
||||
tmpFile.Close()
|
||||
return 0, err
|
||||
}
|
||||
if err := tmpFile.Close(); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
// Re-open the temp file for reading and stream it to S3
|
||||
f, err := os.Open(tmpPath)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer f.Close()
|
||||
if err := c.client.PutObject(context.Background(), id, f, size); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
c.mu.Lock()
|
||||
c.totalSizeCurrent += size
|
||||
c.mu.Unlock()
|
||||
@@ -153,19 +89,7 @@ func (c *s3Store) Read(id string) (io.ReadCloser, int64, error) {
|
||||
if !fileIDRegex.MatchString(id) {
|
||||
return nil, 0, errInvalidFileID
|
||||
}
|
||||
key := c.objectKey(id)
|
||||
resp, err := c.client.GetObject(context.Background(), &s3.GetObjectInput{
|
||||
Bucket: aws.String(c.bucket),
|
||||
Key: aws.String(key),
|
||||
})
|
||||
if err != nil {
|
||||
return nil, 0, fmt.Errorf("s3 store: GetObject failed: %w", err)
|
||||
}
|
||||
var size int64
|
||||
if resp.ContentLength != nil {
|
||||
size = *resp.ContentLength
|
||||
}
|
||||
return resp.Body, size, nil
|
||||
return c.client.GetObject(context.Background(), id)
|
||||
}
|
||||
|
||||
func (c *s3Store) Remove(ids ...string) error {
|
||||
@@ -181,23 +105,11 @@ func (c *s3Store) Remove(ids ...string) error {
|
||||
end = len(ids)
|
||||
}
|
||||
batch := ids[i:end]
|
||||
objects := make([]s3types.ObjectIdentifier, len(batch))
|
||||
for j, id := range batch {
|
||||
for _, id := range batch {
|
||||
log.Tag(tagS3Store).Field("message_id", id).Debug("Deleting attachment from S3")
|
||||
key := c.objectKey(id)
|
||||
objects[j] = s3types.ObjectIdentifier{
|
||||
Key: aws.String(key),
|
||||
}
|
||||
}
|
||||
_, err := c.client.DeleteObjects(context.Background(), &s3.DeleteObjectsInput{
|
||||
Bucket: aws.String(c.bucket),
|
||||
Delete: &s3types.Delete{
|
||||
Objects: objects,
|
||||
Quiet: aws.Bool(true),
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
return fmt.Errorf("s3 store: DeleteObjects failed: %w", err)
|
||||
if err := c.client.DeleteObjects(context.Background(), batch); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
// Recalculate totalSizeCurrent via ListObjectsV2 (matches fileStore's dirSize rescan pattern)
|
||||
@@ -227,29 +139,15 @@ func (c *s3Store) Remaining() int64 {
|
||||
return remaining
|
||||
}
|
||||
|
||||
// computeSize uses ListAllObjects to sum up the total size of all objects with our prefix.
|
||||
func (c *s3Store) computeSize() (int64, error) {
|
||||
var size int64
|
||||
paginator := s3.NewListObjectsV2Paginator(c.client, &s3.ListObjectsV2Input{
|
||||
Bucket: aws.String(c.bucket),
|
||||
Prefix: aws.String(c.prefixForList()),
|
||||
})
|
||||
for paginator.HasMorePages() {
|
||||
page, err := paginator.NextPage(context.Background())
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
for _, obj := range page.Contents {
|
||||
if obj.Size != nil {
|
||||
size += *obj.Size
|
||||
}
|
||||
}
|
||||
objects, err := c.client.ListAllObjects(context.Background())
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return size, nil
|
||||
}
|
||||
|
||||
func (c *s3Store) prefixForList() string {
|
||||
if c.prefix != "" {
|
||||
return c.prefix + "/"
|
||||
var totalSize int64
|
||||
for _, obj := range objects {
|
||||
totalSize += obj.Size
|
||||
}
|
||||
return ""
|
||||
return totalSize, nil
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user