multipleof4/ntfy
1
0
Fork 0
mirror of https://github.com/multipleof4/ntfy.git synced 2026-10-16 16:37:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

Allow streaming to S3

Browse Source
This commit is contained in:
binwiederhier
2026-03-21 21:14:49 -04:00
parent b3a8f18019
commit b81218953a
12 changed files with 181 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files
attachment/backend.go
+1 -1
View File
@@ -15,7 +15,7 @@ type object struct {
// backend is a minimal I/O interface for storing and retrieving attachment files.
// It has no knowledge of size tracking, limiting, or ID validation.
type backend interface {
Put(id string, in io.Reader) error
Put(id string, reader io.Reader, untrustedLength int64) error
Get(id string) (io.ReadCloser, int64, error)
List() ([]object, error)
Delete(ids ...string) error
attachment/backend_file.go
+10 -2
View File
@@ -1,6 +1,7 @@
package attachment
import (
"fmt"
"io"
"os"
"path/filepath"
@@ -24,16 +25,23 @@ func newFileBackend(dir string) (*fileBackend, error) {
return &fileBackend{dir: dir}, nil
}
func (b *fileBackend) Put(id string, in io.Reader) error {
func (b *fileBackend) Put(id string, reader io.Reader, untrustedLength int64) error {
if untrustedLength > 0 {
reader = io.LimitReader(reader, untrustedLength)
}
file := filepath.Join(b.dir, id)
f, err := os.OpenFile(file, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0600)
if err != nil {
return err
}
defer f.Close()
if _, err := io.Copy(f, in); err != nil {
n, err := io.Copy(f, reader)
if err != nil {
os.Remove(file)
return err
} else if untrustedLength > 0 && n != untrustedLength {
os.Remove(file)
return fmt.Errorf("content length mismatch: claimed %d, got %d", untrustedLength, n)
}
if err := f.Close(); err != nil {
os.Remove(file)
attachment/backend_s3.go
+2 -2
View File
@@ -24,8 +24,8 @@ func newS3Backend(client *s3.Client) *s3Backend {
return &s3Backend{client: client}
}
func (b *s3Backend) Put(id string, in io.Reader) error {
return b.client.PutObject(context.Background(), id, in)
func (b *s3Backend) Put(id string, reader io.Reader, untrustedLength int64) error {
return b.client.PutObject(context.Background(), id, reader, untrustedLength)
}
func (b *s3Backend) Get(id string) (io.ReadCloser, int64, error) {
attachment/store.go
+8 -6
View File
@@ -72,20 +72,22 @@ func newStore(backend backend, totalSizeLimit int64, localIDs func() ([]string,
}
// Write stores an attachment file. The id is validated, and the write is subject to
// the total size limit and any additional limiters.
func (c *Store) Write(id string, in io.Reader, limiters ...util.Limiter) (int64, error) {
// the total size limit and any additional limiters. The untrustedLength is a hint
// from the client's Content-Length header; backends may use it to optimize uploads (e.g.
// streaming directly to S3 without buffering).
func (c *Store) Write(id string, reader io.Reader, untrustedLength int64, limiters ...util.Limiter) (int64, error) {
if !fileIDRegex.MatchString(id) {
return 0, errInvalidFileID
}
log.Tag(tagStore).Field("message_id", id).Debug("Writing attachment")
limiters = append(limiters, util.NewFixedLimiter(c.Remaining()))
cr := util.NewCountingReader(in)
lr := util.NewLimitReader(cr, limiters...)
if err := c.backend.Put(id, lr); err != nil {
countingReader := util.NewCountingReader(reader)
limitReader := util.NewLimitReader(countingReader, limiters...)
if err := c.backend.Put(id, limitReader, untrustedLength); err != nil {
c.backend.Delete(id) //nolint:errcheck
return 0, err
}
size := cr.Total()
size := countingReader.Total()
c.mu.Lock()
c.size += size
c.sizes[id] = size
attachment/store_file_test.go
+36 -10
View File
@@ -19,7 +19,7 @@ var (
func TestFileStore_Write_Success(t *testing.T) {
dir, c := newTestFileStore(t)
size, err := c.Write("abcdefghijkl", strings.NewReader("normal file"), util.NewFixedLimiter(999))
size, err := c.Write("abcdefghijkl", strings.NewReader("normal file"), 0, util.NewFixedLimiter(999))
require.Nil(t, err)
require.Equal(t, int64(11), size)
require.Equal(t, "normal file", readFile(t, dir+"/abcdefghijkl"))
@@ -29,7 +29,7 @@ func TestFileStore_Write_Success(t *testing.T) {
func TestFileStore_Write_Read_Success(t *testing.T) {
_, c := newTestFileStore(t)
size, err := c.Write("abcdefghijkl", strings.NewReader("hello world"))
size, err := c.Write("abcdefghijkl", strings.NewReader("hello world"), 0)
require.Nil(t, err)
require.Equal(t, int64(11), size)
@@ -45,7 +45,7 @@ func TestFileStore_Write_Read_Success(t *testing.T) {
func TestFileStore_Write_Remove_Success(t *testing.T) {
dir, c := newTestFileStore(t) // max = 10k (10240), each = 1k (1024)
for i := 0; i < 10; i++ { // 10x999 = 9990
size, err := c.Write(fmt.Sprintf("abcdefghijk%d", i), bytes.NewReader(make([]byte, 999)))
size, err := c.Write(fmt.Sprintf("abcdefghijk%d", i), bytes.NewReader(make([]byte, 999)), 0)
require.Nil(t, err)
require.Equal(t, int64(999), size)
}
@@ -64,22 +64,48 @@ func TestFileStore_Write_Remove_Success(t *testing.T) {
func TestFileStore_Write_FailedTotalSizeLimit(t *testing.T) {
dir, c := newTestFileStore(t)
for i := 0; i < 10; i++ {
size, err := c.Write(fmt.Sprintf("abcdefghijk%d", i), bytes.NewReader(oneKilobyteArray))
size, err := c.Write(fmt.Sprintf("abcdefghijk%d", i), bytes.NewReader(oneKilobyteArray), 0)
require.Nil(t, err)
require.Equal(t, int64(1024), size)
}
_, err := c.Write("abcdefghijkX", bytes.NewReader(oneKilobyteArray))
_, err := c.Write("abcdefghijkX", bytes.NewReader(oneKilobyteArray), 0)
require.Equal(t, util.ErrLimitReached, err)
require.NoFileExists(t, dir+"/abcdefghijkX")
}
func TestFileStore_Write_FailedAdditionalLimiter(t *testing.T) {
dir, c := newTestFileStore(t)
_, err := c.Write("abcdefghijkl", bytes.NewReader(make([]byte, 1001)), util.NewFixedLimiter(1000))
_, err := c.Write("abcdefghijkl", bytes.NewReader(make([]byte, 1001)), 0, util.NewFixedLimiter(1000))
require.Equal(t, util.ErrLimitReached, err)
require.NoFileExists(t, dir+"/abcdefghijkl")
}
func TestFileStore_Write_UntrustedContentLengthExact(t *testing.T) {
dir, c := newTestFileStore(t)
size, err := c.Write("abcdefghijkl", strings.NewReader("hello world"), 11)
require.Nil(t, err)
require.Equal(t, int64(11), size)
require.Equal(t, "hello world", readFile(t, dir+"/abcdefghijkl"))
}
func TestFileStore_Write_UntrustedContentLengthBodyLonger(t *testing.T) {
dir, c := newTestFileStore(t)
// Body has 11 bytes, but we claim 5 — only first 5 bytes should be stored
size, err := c.Write("abcdefghijkl", strings.NewReader("hello world"), 5)
require.Nil(t, err)
require.Equal(t, int64(5), size)
require.Equal(t, "hello", readFile(t, dir+"/abcdefghijkl"))
}
func TestFileStore_Write_UntrustedContentLengthBodyShorter(t *testing.T) {
dir, c := newTestFileStore(t)
// Body has 5 bytes, but we claim 100 — should fail with content length mismatch
_, err := c.Write("abcdefghijkl", strings.NewReader("hello"), 100)
require.Error(t, err)
require.Contains(t, err.Error(), "content length mismatch")
require.NoFileExists(t, dir+"/abcdefghijkl")
}
func TestFileStore_Read_NotFound(t *testing.T) {
_, c := newTestFileStore(t)
_, _, err := c.Read("abcdefghijkl")
@@ -90,11 +116,11 @@ func TestFileStore_Sync(t *testing.T) {
dir, c := newTestFileStore(t)
// Write some files
_, err := c.Write("abcdefghijk0", strings.NewReader("file0"))
_, err := c.Write("abcdefghijk0", strings.NewReader("file0"), 0)
require.Nil(t, err)
_, err = c.Write("abcdefghijk1", strings.NewReader("file1"))
_, err = c.Write("abcdefghijk1", strings.NewReader("file1"), 0)
require.Nil(t, err)
_, err = c.Write("abcdefghijk2", strings.NewReader("file2"))
_, err = c.Write("abcdefghijk2", strings.NewReader("file2"), 0)
require.Nil(t, err)
require.Equal(t, int64(15), c.Size())
@@ -124,7 +150,7 @@ func TestFileStore_Sync_SkipsRecentFiles(t *testing.T) {
dir, c := newTestFileStore(t)
// Write a file
_, err := c.Write("abcdefghijk0", strings.NewReader("file0"))
_, err := c.Write("abcdefghijk0", strings.NewReader("file0"), 0)
require.Nil(t, err)
// Set the ID provider to return empty (no valid IDs)
attachment/store_s3_test.go
+11 -11
View File
@@ -26,7 +26,7 @@ func TestS3Store_WriteReadRemove(t *testing.T) {
cache := newTestS3Store(t, server, "my-bucket", "pfx", 10*1024)
// Write
size, err := cache.Write("abcdefghijkl", strings.NewReader("hello world"))
size, err := cache.Write("abcdefghijkl", strings.NewReader("hello world"), 0)
require.Nil(t, err)
require.Equal(t, int64(11), size)
require.Equal(t, int64(11), cache.Size())
@@ -55,7 +55,7 @@ func TestS3Store_WriteNoPrefix(t *testing.T) {
cache := newTestS3Store(t, server, "my-bucket", "", 10*1024)
size, err := cache.Write("abcdefghijkl", strings.NewReader("test"))
size, err := cache.Write("abcdefghijkl", strings.NewReader("test"), 0)
require.Nil(t, err)
require.Equal(t, int64(4), size)
@@ -74,13 +74,13 @@ func TestS3Store_WriteTotalSizeLimit(t *testing.T) {
cache := newTestS3Store(t, server, "my-bucket", "pfx", 100)
// First write fits
_, err := cache.Write("abcdefghijk0", bytes.NewReader(make([]byte, 80)))
_, err := cache.Write("abcdefghijk0", bytes.NewReader(make([]byte, 80)), 0)
require.Nil(t, err)
require.Equal(t, int64(80), cache.Size())
require.Equal(t, int64(20), cache.Remaining())
// Second write exceeds total limit
_, err = cache.Write("abcdefghijk1", bytes.NewReader(make([]byte, 50)))
_, err = cache.Write("abcdefghijk1", bytes.NewReader(make([]byte, 50)), 0)
require.ErrorIs(t, err, util.ErrLimitReached)
}
@@ -90,7 +90,7 @@ func TestS3Store_WriteFileSizeLimit(t *testing.T) {
cache := newTestS3Store(t, server, "my-bucket", "pfx", 10*1024)
_, err := cache.Write("abcdefghijkl", bytes.NewReader(make([]byte, 200)), util.NewFixedLimiter(100))
_, err := cache.Write("abcdefghijkl", bytes.NewReader(make([]byte, 200)), 0, util.NewFixedLimiter(100))
require.ErrorIs(t, err, util.ErrLimitReached)
}
@@ -101,7 +101,7 @@ func TestS3Store_WriteRemoveMultiple(t *testing.T) {
cache := newTestS3Store(t, server, "my-bucket", "pfx", 10*1024)
for i := 0; i < 5; i++ {
_, err := cache.Write(fmt.Sprintf("abcdefghijk%d", i), bytes.NewReader(make([]byte, 100)))
_, err := cache.Write(fmt.Sprintf("abcdefghijk%d", i), bytes.NewReader(make([]byte, 100)), 0)
require.Nil(t, err)
}
require.Equal(t, int64(500), cache.Size())
@@ -126,7 +126,7 @@ func TestS3Store_InvalidID(t *testing.T) {
cache := newTestS3Store(t, server, "my-bucket", "pfx", 10*1024)
_, err := cache.Write("bad", strings.NewReader("x"))
_, err := cache.Write("bad", strings.NewReader("x"), 0)
require.Equal(t, errInvalidFileID, err)
_, _, err = cache.Read("bad")
@@ -143,11 +143,11 @@ func TestS3Store_Sync(t *testing.T) {
cache := newTestS3Store(t, server, "my-bucket", "pfx", 10*1024)
// Write some files
_, err := cache.Write("abcdefghijk0", strings.NewReader("file0"))
_, err := cache.Write("abcdefghijk0", strings.NewReader("file0"), 0)
require.Nil(t, err)
_, err = cache.Write("abcdefghijk1", strings.NewReader("file1"))
_, err = cache.Write("abcdefghijk1", strings.NewReader("file1"), 0)
require.Nil(t, err)
_, err = cache.Write("abcdefghijk2", strings.NewReader("file2"))
_, err = cache.Write("abcdefghijk2", strings.NewReader("file2"), 0)
require.Nil(t, err)
require.Equal(t, int64(15), cache.Size())
@@ -175,7 +175,7 @@ func TestS3Store_Sync_SkipsRecentFiles(t *testing.T) {
cache := newTestS3Store(t, mockServer, "my-bucket", "pfx", 10*1024)
_, err := cache.Write("abcdefghijk0", strings.NewReader("file0"))
_, err := cache.Write("abcdefghijk0", strings.NewReader("file0"), 0)
require.Nil(t, err)
// Set the ID provider to return empty (no valid IDs)