From e9249d6ff9d75975605331110ff8592605f43159 Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Sat, 28 Jan 2023 15:43:16 +0200
Subject: [PATCH] Upload files before sending message

---
 attachments.go | 21 +++++++++++++++++++++
 go.mod         |  2 +-
 go.sum         |  4 ++--
 portal.go      | 34 +++++++++++++++++++++++++++-------
 user.go        | 11 +++++++++++
 5 files changed, 62 insertions(+), 10 deletions(-)

diff --git a/attachments.go b/attachments.go
index 4cafcc6..a47cba8 100644
--- a/attachments.go
+++ b/attachments.go
@@ -41,6 +41,27 @@ func downloadDiscordAttachment(url string) ([]byte, error) {
 	return io.ReadAll(resp.Body)
 }
 
+func uploadDiscordAttachment(url string, data []byte) error {
+	req, err := http.NewRequest(http.MethodPut, url, bytes.NewReader(data))
+	if err != nil {
+		return err
+	}
+	for key, value := range discordgo.DroidFetchHeaders {
+		req.Header.Set(key, value)
+	}
+
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode > 300 {
+		respData, _ := io.ReadAll(resp.Body)
+		return fmt.Errorf("unexpected status %d: %s", resp.StatusCode, respData)
+	}
+	return nil
+}
+
 func (portal *Portal) downloadMatrixAttachment(content *event.MessageEventContent) ([]byte, error) {
 	var file *event.EncryptedFileInfo
 	rawMXC := content.URL
diff --git a/go.mod b/go.mod
index bfd9fdd..697f542 100644
--- a/go.mod
+++ b/go.mod
@@ -33,4 +33,4 @@ require (
 	maunium.net/go/mauflag v1.0.0 // indirect
 )
 
-replace github.com/bwmarrin/discordgo => github.com/beeper/discordgo v0.0.0-20230128130408-ec0d4a43f122
+replace github.com/bwmarrin/discordgo => github.com/beeper/discordgo v0.0.0-20230128134018-766d08cb045e
diff --git a/go.sum b/go.sum
index f28b7a6..d6230a3 100644
--- a/go.sum
+++ b/go.sum
@@ -1,6 +1,6 @@
 github.com/DATA-DOG/go-sqlmock v1.5.0 h1:Shsta01QNfFxHCfpW6YH2STWB0MudeXXEWMr20OEh60=
-github.com/beeper/discordgo v0.0.0-20230128130408-ec0d4a43f122 h1:VbdWGvPQbbxURlc6bmEbluJo2jv/FwwBGVb1zX9E+7o=
-github.com/beeper/discordgo v0.0.0-20230128130408-ec0d4a43f122/go.mod h1:NJZpH+1AfhIcyQsPeuBKsUtYrRnjkyu0kIVMCHkZtRY=
+github.com/beeper/discordgo v0.0.0-20230128134018-766d08cb045e h1:R0Db6p3gANvV2Hk8lbSSlPDNG3zzeOM8nyZHmLl3tkI=
+github.com/beeper/discordgo v0.0.0-20230128134018-766d08cb045e/go.mod h1:NJZpH+1AfhIcyQsPeuBKsUtYrRnjkyu0kIVMCHkZtRY=
 github.com/coreos/go-systemd/v22 v22.3.3-0.20220203105225-a9a7ef127534/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
diff --git a/portal.go b/portal.go
index 9449f88..df82762 100644
--- a/portal.go
+++ b/portal.go
@@ -1,7 +1,6 @@
 package main
 
 import (
-	"bytes"
 	"errors"
 	"fmt"
 	"html"
@@ -1513,9 +1512,11 @@ func (portal *Portal) handleMatrixMessage(sender *User, evt *event.Event) {
 
 	var sendReq discordgo.MessageSend
 
+	var description string
 	if evt.Type == event.EventSticker {
 		content.MsgType = event.MsgImage
 		if mimeData := mimetype.Lookup(content.Info.MimeType); mimeData != nil {
+			description = content.Body
 			content.Body = "sticker" + mimeData.Extension()
 		}
 	}
@@ -1539,15 +1540,34 @@ func (portal *Portal) handleMatrixMessage(sender *User, evt *event.Event) {
 			return
 		}
 
-		sendReq.Files = []*discordgo.File{{
-			Name:        content.Body,
-			ContentType: content.Info.MimeType,
-			Reader:      bytes.NewReader(data),
-		}}
+		att := &discordgo.MessageAttachment{
+			ID:          "0",
+			Filename:    content.Body,
+			Description: description,
+		}
+		sendReq.Attachments = []*discordgo.MessageAttachment{att}
 		if content.FileName != "" && content.FileName != content.Body {
-			sendReq.Files[0].Name = content.FileName
+			att.Filename = content.FileName
 			sendReq.Content = portal.parseMatrixHTML(sender, content)
 		}
+		prep, err := sender.Session.ChannelAttachmentCreate(channelID, &discordgo.ReqPrepareAttachments{
+			Files: []*discordgo.FilePrepare{{
+				Size: len(data),
+				Name: att.Filename,
+				ID:   sender.NextDiscordUploadID(),
+			}},
+		})
+		if err != nil {
+			go portal.sendMessageMetrics(evt, err, "Error preparing to reupload media in")
+			return
+		}
+		prepared := prep.Attachments[0]
+		att.UploadedFilename = prepared.UploadFilename
+		err = uploadDiscordAttachment(prepared.UploadURL, data)
+		if err != nil {
+			go portal.sendMessageMetrics(evt, err, "Error reuploading media in")
+			return
+		}
 	default:
 		go portal.sendMessageMetrics(evt, fmt.Errorf("%w %q", errUnknownMsgType, content.MsgType), "Ignoring")
 		return
diff --git a/user.go b/user.go
index 12c7301..b4c1ff4 100644
--- a/user.go
+++ b/user.go
@@ -3,11 +3,14 @@ package main
 import (
 	"errors"
 	"fmt"
+	"math/rand"
 	"net/http"
 	"os"
 	"runtime"
+	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/gorilla/websocket"
@@ -55,6 +58,8 @@ type User struct {
 
 	markedOpened     map[string]time.Time
 	markedOpenedLock sync.Mutex
+
+	nextDiscordUploadID atomic.Int32
 }
 
 func (user *User) GetRemoteID() string {
@@ -193,6 +198,7 @@ func (br *DiscordBridge) NewUser(dbUser *database.User) *User {
 		markedOpened:    make(map[string]time.Time),
 		PermissionLevel: br.Config.Bridge.Permissions.Get(dbUser.MXID),
 	}
+	user.nextDiscordUploadID.Store(rand.Int31n(100))
 	user.BridgeState = br.NewBridgeStateQueue(user, user.log)
 	return user
 }
@@ -433,6 +439,11 @@ func (user *User) syncChatDoublePuppetDetails(portal *Portal, justCreated bool)
 	}
 }
 
+func (user *User) NextDiscordUploadID() string {
+	val := user.nextDiscordUploadID.Add(2)
+	return strconv.Itoa(int(val))
+}
+
 func (user *User) Login(token string) error {
 	user.bridgeStateLock.Lock()
 	user.wasLoggedOut = false