bridgev2/matrix: prevent too many async uploads at once

This commit is contained in:
Tulir Asokan 2024-08-19 17:52:54 +03:00
commit 7f392b17b6
4 changed files with 15 additions and 0 deletions

View file

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

View file

@ -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 {

1
go.mod
View file

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

2
go.sum
View file

@ -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=