diff --git a/bridgev2/matrix/connector.go b/bridgev2/matrix/connector.go index c2b3f7cd..8c35fb4f 100644 --- a/bridgev2/matrix/connector.go +++ b/bridgev2/matrix/connector.go @@ -28,6 +28,7 @@ import ( _ "go.mau.fi/util/dbutil/litestream" "go.mau.fi/util/exsync" "go.mau.fi/util/random" + "golang.org/x/sync/semaphore" "maunium.net/go/mautrix" "maunium.net/go/mautrix/appservice" @@ -70,6 +71,7 @@ type Connector struct { DoublePuppet *doublePuppetUtil MediaProxy *mediaproxy.MediaProxy + uploadSema *semaphore.Weighted dmaSigKey [32]byte pubMediaSigKey []byte @@ -108,6 +110,7 @@ func NewConnector(cfg *bridgeconfig.Config) *Connector { c.Config = cfg c.userIDRegex = cfg.MakeUserIDRegex("(.+)") c.MediaConfig.UploadSize = 50 * 1024 * 1024 + c.uploadSema = semaphore.NewWeighted(c.MediaConfig.UploadSize * 2) c.Capabilities = &bridgev2.MatrixCapabilities{} c.doublePuppetIntents = exsync.NewMap[id.UserID, *appservice.IntentAPI]() return c @@ -366,6 +369,7 @@ func (br *Connector) fetchMediaConfig(ctx context.Context) { if ok { mfsn.SetMaxFileSize(br.MediaConfig.UploadSize) } + br.uploadSema = semaphore.NewWeighted(br.MediaConfig.UploadSize * 2) } } diff --git a/bridgev2/matrix/intent.go b/bridgev2/matrix/intent.go index 0c28ad50..5eb31436 100644 --- a/bridgev2/matrix/intent.go +++ b/bridgev2/matrix/intent.go @@ -211,6 +211,9 @@ func (as *ASIntent) DownloadMedia(ctx context.Context, uri id.ContentURIString, } func (as *ASIntent) UploadMedia(ctx context.Context, roomID id.RoomID, data []byte, fileName, mimeType string) (url id.ContentURIString, file *event.EncryptedFileInfo, err error) { + if int64(len(data)) > as.Connector.MediaConfig.UploadSize { + return "", nil, fmt.Errorf("file too large (%.2f MB > %.2f MB)", float64(len(data))/1000/1000, float64(as.Connector.MediaConfig.UploadSize)/1000/1000) + } if roomID != "" { var encrypted bool if encrypted, err = as.Matrix.StateStore.IsEncrypted(ctx, roomID); err != nil { @@ -231,6 +234,11 @@ func (as *ASIntent) UploadMedia(ctx context.Context, roomID id.RoomID, data []by FileName: fileName, } if as.Connector.Config.Homeserver.AsyncMedia { + // Prevent too many background uploads at once + err = as.Connector.uploadSema.Acquire(ctx, int64(len(data))) + if err != nil { + return + } var resp *mautrix.RespCreateMXC resp, err = as.Matrix.UploadAsync(ctx, req) if resp != nil { diff --git a/go.mod b/go.mod index 4ad4143c..6ff2233d 100644 --- a/go.mod +++ b/go.mod @@ -20,6 +20,7 @@ require ( golang.org/x/crypto v0.26.0 golang.org/x/exp v0.0.0-20240808152545-0cdaa3abc0fa golang.org/x/net v0.28.0 + golang.org/x/sync v0.8.0 gopkg.in/yaml.v3 v3.0.1 maunium.net/go/mauflag v1.0.0 ) diff --git a/go.sum b/go.sum index 0adc7117..53105fc7 100644 --- a/go.sum +++ b/go.sum @@ -58,6 +58,8 @@ golang.org/x/exp v0.0.0-20240808152545-0cdaa3abc0fa h1:ELnwvuAXPNtPk1TJRuGkI9fDT golang.org/x/exp v0.0.0-20240808152545-0cdaa3abc0fa/go.mod h1:akd2r19cwCdwSwWeIdzYQGa/EZZyqcOdwWiwj5L5eKQ= golang.org/x/net v0.28.0 h1:a9JDOJc5GMUJ0+UDqmLT86WiEy7iWyIhz8gz8E4e5hE= golang.org/x/net v0.28.0/go.mod h1:yqtgsTWOOnlGLG9GFRrK3++bGOUEkNBoHZc8MEDWPNg= +golang.org/x/sync v0.8.0 h1:3NFvSEYkUoMifnESzZl15y791HH1qU2xm6eCJU5ZPXQ= +golang.org/x/sync v0.8.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk= golang.org/x/sys v0.0.0-20220310020820-b874c991c1a5/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=