From 8b61dc5352c508ce48350fc70a01496429c2c5d1 Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Fri, 29 Nov 2024 19:26:47 +0200
Subject: [PATCH] config: add support for using a proxy

---
 attachments.go      | 10 +++++-----
 config/bridge.go    |  2 ++
 config/upgrade.go   |  1 +
 example-config.yaml |  2 ++
 portal.go           |  2 +-
 user.go             | 14 ++++++++++++++
 6 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/attachments.go b/attachments.go
index 50d3c0b..befef9e 100644
--- a/attachments.go
+++ b/attachments.go
@@ -29,7 +29,7 @@ import (
 	"go.mau.fi/mautrix-discord/database"
 )
 
-func downloadDiscordAttachment(url string, maxSize int64) ([]byte, error) {
+func downloadDiscordAttachment(cli *http.Client, url string, maxSize int64) ([]byte, error) {
 	req, err := http.NewRequest(http.MethodGet, url, nil)
 	if err != nil {
 		return nil, err
@@ -38,7 +38,7 @@ func downloadDiscordAttachment(url string, maxSize int64) ([]byte, error) {
 		req.Header.Set(key, value)
 	}
 
-	resp, err := http.DefaultClient.Do(req)
+	resp, err := cli.Do(req)
 	if err != nil {
 		return nil, err
 	}
@@ -65,7 +65,7 @@ func downloadDiscordAttachment(url string, maxSize int64) ([]byte, error) {
 	}
 }
 
-func uploadDiscordAttachment(url string, data []byte) error {
+func uploadDiscordAttachment(cli *http.Client, url string, data []byte) error {
 	req, err := http.NewRequest(http.MethodPut, url, bytes.NewReader(data))
 	if err != nil {
 		return err
@@ -80,7 +80,7 @@ func uploadDiscordAttachment(url string, data []byte) error {
 	req.Header.Del("X-Discord-Timezone")
 	req.Header.Del("X-Super-Properties")
 
-	resp, err := http.DefaultClient.Do(req)
+	resp, err := cli.Do(req)
 	if err != nil {
 		return err
 	}
@@ -301,7 +301,7 @@ func (br *DiscordBridge) copyAttachmentToMatrix(intent *appservice.IntentAPI, ur
 			}()
 
 			var data []byte
-			data, onceErr = downloadDiscordAttachment(url, br.MediaConfig.UploadSize)
+			data, onceErr = downloadDiscordAttachment(http.DefaultClient, url, br.MediaConfig.UploadSize)
 			if onceErr != nil {
 				return
 			}
diff --git a/config/bridge.go b/config/bridge.go
index dd6ad4e..2f78ed7 100644
--- a/config/bridge.go
+++ b/config/bridge.go
@@ -57,6 +57,8 @@ type BridgeConfig struct {
 	EnableWebhookAvatars        bool `yaml:"enable_webhook_avatars"`
 	UseDiscordCDNUpload         bool `yaml:"use_discord_cdn_upload"`
 
+	Proxy string `yaml:"proxy"`
+
 	CacheMedia  string      `yaml:"cache_media"`
 	DirectMedia DirectMedia `yaml:"direct_media"`
 
diff --git a/config/upgrade.go b/config/upgrade.go
index f84d9ef..041e8e6 100644
--- a/config/upgrade.go
+++ b/config/upgrade.go
@@ -63,6 +63,7 @@ func DoUpgrade(helper *up.Helper) {
 	helper.Copy(up.Bool, "bridge", "prefix_webhook_messages")
 	helper.Copy(up.Bool, "bridge", "enable_webhook_avatars")
 	helper.Copy(up.Bool, "bridge", "use_discord_cdn_upload")
+	helper.Copy(up.Str|up.Null, "bridge", "proxy")
 	helper.Copy(up.Str, "bridge", "cache_media")
 	helper.Copy(up.Bool, "bridge", "direct_media", "enabled")
 	helper.Copy(up.Str, "bridge", "direct_media", "server_name")
diff --git a/example-config.yaml b/example-config.yaml
index 7fca846..54792f9 100644
--- a/example-config.yaml
+++ b/example-config.yaml
@@ -168,6 +168,8 @@ bridge:
     # like the official client does? The other option is sending the media in the message send request as a form part
     # (which is always used by bots and webhooks).
     use_discord_cdn_upload: true
+    # Proxy for Discord connections
+    proxy:
     # Should mxc uris copied from Discord be cached?
     # This can be `never` to never cache, `unencrypted` to only cache unencrypted mxc uris, or `always` to cache everything.
     # If you have a media repo that generates non-unique mxc uris, you should set this to never.
diff --git a/portal.go b/portal.go
index 2c91156..9a094d3 100644
--- a/portal.go
+++ b/portal.go
@@ -1647,7 +1647,7 @@ func (portal *Portal) handleMatrixMessage(sender *User, evt *event.Event) {
 			}
 			prepared := prep.Attachments[0]
 			att.UploadedFilename = prepared.UploadFilename
-			err = uploadDiscordAttachment(prepared.UploadURL, data)
+			err = uploadDiscordAttachment(sender.Session.Client, prepared.UploadURL, data)
 			if err != nil {
 				go portal.sendMessageMetrics(evt, err, "Error reuploading media in")
 				return
diff --git a/user.go b/user.go
index d3c85e5..63ede61 100644
--- a/user.go
+++ b/user.go
@@ -2,10 +2,12 @@ package main
 
 import (
 	"context"
+	"crypto/tls"
 	"errors"
 	"fmt"
 	"math/rand"
 	"net/http"
+	"net/url"
 	"os"
 	"runtime/debug"
 	"sort"
@@ -547,6 +549,18 @@ func (user *User) Connect() error {
 	if err != nil {
 		return err
 	}
+	if user.bridge.Config.Bridge.Proxy != "" {
+		u, _ := url.Parse(user.bridge.Config.Bridge.Proxy)
+		tlsConf := &tls.Config{
+			InsecureSkipVerify: os.Getenv("DISCORD_SKIP_TLS_VERIFICATION") == "true",
+		}
+		session.Client.Transport = &http.Transport{
+			Proxy:           http.ProxyURL(u),
+			TLSClientConfig: tlsConf,
+		}
+		session.Dialer.Proxy = http.ProxyURL(u)
+		session.Dialer.TLSClientConfig = tlsConf
+	}
 	// TODO move to config
 	if os.Getenv("DISCORD_DEBUG") == "1" {
 		session.LogLevel = discordgo.LogDebug