From 741b4b9befcbc00e4bb17e5cbb6a8197cc55c421 Mon Sep 17 00:00:00 2001
From: Travis Ralston <travpc@gmail.com>
Date: Sun, 14 Jan 2024 16:01:55 -0700
Subject: [PATCH] Support being redirected over federation

Fixes https://github.com/turt2live/matrix-media-repo/issues/515
---
 CHANGELOG.md                              |  1 +
 matrix/requests.go                        | 15 +++++++++++----
 pipelines/_steps/download/try_download.go |  2 +-
 3 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index bc685a08..6d244238 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -12,6 +12,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 * Dendrite homeservers can now have their media imported safely, and `adminApiKind` may be set to `dendrite`.
 * Exporting MMR's data to Synapse is now possible with `import_to_synapse`. To use it, first run `gdpr_export` or similar.
 * Errors encountered during a background task, such as an API-induced export, are exposed as `error_message` in the admin API.
+* MMR will follow redirects on federated downloads up to 5 hops.
 
 ### Changed
 
diff --git a/matrix/requests.go b/matrix/requests.go
index b14f8475..4b6c7dda 100644
--- a/matrix/requests.go
+++ b/matrix/requests.go
@@ -13,7 +13,6 @@ import (
 	"os"
 	"time"
 
-	"github.com/sirupsen/logrus"
 	"github.com/turt2live/matrix-media-repo/common/rcontext"
 )
 
@@ -80,7 +79,7 @@ func doRequest(ctx rcontext.RequestContext, method string, urlStr string, body i
 }
 
 func FederatedGet(url string, realHost string, ctx rcontext.RequestContext) (*http.Response, error) {
-	logrus.Debug("Doing federated GET to " + url + " with host " + realHost)
+	ctx.Log.Debug("Doing federated GET to " + url + " with host " + realHost)
 
 	cb := getFederationBreaker(realHost)
 
@@ -117,7 +116,6 @@ func FederatedGet(url string, realHost string, ctx rcontext.RequestContext) (*ht
 						ServerName: realHost,
 					},
 				},
-				Timeout: time.Duration(ctx.Config.TimeoutSeconds.Federation) * time.Second,
 			}
 		} else {
 			ctx.Log.Warn("Ignoring any certificate errors while making request")
@@ -143,10 +141,19 @@ func FederatedGet(url string, realHost string, ctx rcontext.RequestContext) (*ht
 			}
 			client = &http.Client{
 				Transport: tr,
-				Timeout:   time.Duration(ctx.Config.TimeoutSeconds.UrlPreviews) * time.Second,
 			}
 		}
 
+		client.Timeout = time.Duration(ctx.Config.TimeoutSeconds.UrlPreviews) * time.Second
+		client.CheckRedirect = func(req *http.Request, via []*http.Request) error {
+			if len(via) > 5 { // arbitrary
+				return errors.New("too many redirects")
+			}
+			ctx.Log.Debugf("Redirected to %s", req.URL.String())
+			client.Transport = nil // Clear our TLS handler as we're out of the Matrix certificate verification steps
+			return nil
+		}
+
 		resp, err = client.Do(req)
 		if err != nil {
 			return err
diff --git a/pipelines/_steps/download/try_download.go b/pipelines/_steps/download/try_download.go
index 54933388..35e4cbda 100644
--- a/pipelines/_steps/download/try_download.go
+++ b/pipelines/_steps/download/try_download.go
@@ -54,7 +54,7 @@ func TryDownload(ctx rcontext.RequestContext, origin string, mediaId string) (*d
 			return
 		}
 
-		downloadUrl := fmt.Sprintf("%s/_matrix/media/v3/download/%s/%s?allow_remote=false", baseUrl, url.PathEscape(origin), url.PathEscape(mediaId))
+		downloadUrl := fmt.Sprintf("%s/_matrix/media/v3/download/%s/%s?allow_remote=false&allow_redirect=true", baseUrl, url.PathEscape(origin), url.PathEscape(mediaId))
 		resp, err := matrix.FederatedGet(downloadUrl, realHost, ctx)
 		metrics.MediaDownloaded.With(prometheus.Labels{"origin": origin}).Inc()
 		if err != nil {