From 9b1b0426c302706950fb456ccaa371735fdd29e8 Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Sun, 20 Oct 2024 19:46:14 +0300
Subject: [PATCH] media: allow streaming unencrypted files

---
 cmd/gomuks/media.go | 38 ++++++++++++++++++++++++++++----------
 go.mod              |  2 +-
 go.sum              |  4 ++--
 3 files changed, 31 insertions(+), 13 deletions(-)

diff --git a/cmd/gomuks/media.go b/cmd/gomuks/media.go
index a6f0cc8..df56c6a 100644
--- a/cmd/gomuks/media.go
+++ b/cmd/gomuks/media.go
@@ -104,6 +104,15 @@ func cacheEntryToHeaders(w http.ResponseWriter, entry *database.Media) {
 	w.Header().Set("Content-Security-Policy", "sandbox; default-src 'none'; script-src 'none';")
 }
 
+type noErrorWriter struct {
+	io.Writer
+}
+
+func (new *noErrorWriter) Write(p []byte) (n int, err error) {
+	n, _ = new.Writer.Write(p)
+	return
+}
+
 func (gmx *Gomuks) DownloadMedia(w http.ResponseWriter, r *http.Request) {
 	mxc := id.ContentURI{
 		Homeserver: r.PathValue("server"),
@@ -207,9 +216,23 @@ func (gmx *Gomuks) DownloadMedia(w http.ResponseWriter, r *http.Request) {
 		}
 		reader = cacheEntry.EncFile.DecryptStream(reader)
 	}
+	if cacheEntry.FileName == "" {
+		_, params, _ := mime.ParseMediaType(resp.Header.Get("Content-Disposition"))
+		cacheEntry.FileName = params["filename"]
+	}
+	if cacheEntry.MimeType == "" {
+		cacheEntry.MimeType = resp.Header.Get("Content-Type")
+	}
+	cacheEntry.Size = resp.ContentLength
 	fileHasher := sha256.New()
-	hashReader := io.TeeReader(reader, fileHasher)
-	cacheEntry.Size, err = io.Copy(tempFile, hashReader)
+	wrappedReader := io.TeeReader(reader, fileHasher)
+	if cacheEntry.Size > 0 && cacheEntry.EncFile == nil {
+		cacheEntryToHeaders(w, cacheEntry)
+		w.WriteHeader(http.StatusOK)
+		wrappedReader = io.TeeReader(wrappedReader, &noErrorWriter{w})
+		w = nil
+	}
+	cacheEntry.Size, err = io.Copy(tempFile, wrappedReader)
 	if err != nil {
 		log.Err(err).Msg("Failed to copy media to temporary file")
 		mautrix.MUnknown.WithMessage(fmt.Sprintf("Failed to copy media to temp file: %v", err)).Write(w)
@@ -222,13 +245,6 @@ func (gmx *Gomuks) DownloadMedia(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 	_ = tempFile.Close()
-	if cacheEntry.FileName == "" {
-		_, params, _ := mime.ParseMediaType(resp.Header.Get("Content-Disposition"))
-		cacheEntry.FileName = params["filename"]
-	}
-	if cacheEntry.MimeType == "" {
-		cacheEntry.MimeType = resp.Header.Get("Content-Type")
-	}
 	cacheEntry.Hash = (*[32]byte)(fileHasher.Sum(nil))
 	cacheEntry.Error = nil
 	err = gmx.Client.DB.Media.Put(ctx, cacheEntry)
@@ -250,7 +266,9 @@ func (gmx *Gomuks) DownloadMedia(w http.ResponseWriter, r *http.Request) {
 		mautrix.MUnknown.WithMessage(fmt.Sprintf("Failed to rename temp file: %v", err)).Write(w)
 		return
 	}
-	gmx.downloadMediaFromCache(ctx, w, cacheEntry, true)
+	if w != nil {
+		gmx.downloadMediaFromCache(ctx, w, cacheEntry, true)
+	}
 }
 
 func (gmx *Gomuks) UploadMedia(w http.ResponseWriter, r *http.Request) {
diff --git a/go.mod b/go.mod
index 9c14635..0b3aff7 100644
--- a/go.mod
+++ b/go.mod
@@ -23,7 +23,7 @@ require (
 	golang.org/x/net v0.30.0
 	gopkg.in/yaml.v3 v3.0.1
 	maunium.net/go/mauflag v1.0.0
-	maunium.net/go/mautrix v0.21.2-0.20241018174905-3277c529a2e5
+	maunium.net/go/mautrix v0.21.2-0.20241020164413-e17cb8385518
 	mvdan.cc/xurls/v2 v2.5.0
 )
 
diff --git a/go.sum b/go.sum
index be078a3..5ccdaca 100644
--- a/go.sum
+++ b/go.sum
@@ -87,7 +87,7 @@ gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 maunium.net/go/mauflag v1.0.0 h1:YiaRc0tEI3toYtJMRIfjP+jklH45uDHtT80nUamyD4M=
 maunium.net/go/mauflag v1.0.0/go.mod h1:nLivPOpTpHnpzEh8jEdSL9UqO9+/KBJFmNRlwKfkPeA=
-maunium.net/go/mautrix v0.21.2-0.20241018174905-3277c529a2e5 h1:ITR15W0vCyMiX5ZO6JLCPmHrdSL2+4TLBuDaISwGsFE=
-maunium.net/go/mautrix v0.21.2-0.20241018174905-3277c529a2e5/go.mod h1:sjCZR1R/3NET/WjkcXPL6WpAHlWKku9HjRsdOkbM8Qw=
+maunium.net/go/mautrix v0.21.2-0.20241020164413-e17cb8385518 h1:ROFV2DnfWw6wzNGXJP+SVSsd+cHoyeqHvVvaW5PQyk0=
+maunium.net/go/mautrix v0.21.2-0.20241020164413-e17cb8385518/go.mod h1:sjCZR1R/3NET/WjkcXPL6WpAHlWKku9HjRsdOkbM8Qw=
 mvdan.cc/xurls/v2 v2.5.0 h1:lyBNOm8Wo71UknhUs4QTFUNNMyxy2JEIaKKo0RWOh+8=
 mvdan.cc/xurls/v2 v2.5.0/go.mod h1:yQgaGQ1rFtJUzkmKiHYSSfuQxqfYmd//X6PxvholpeE=