From 23478caa6e73e85eb6a3fb738d36074515551597 Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Tue, 8 Oct 2024 23:47:59 +0300
Subject: [PATCH] media: cache errors

---
 go.mod   |  2 +-
 go.sum   |  4 ++--
 media.go | 46 ++++++++++++++++++++++++++++++++++++++++++++--
 3 files changed, 47 insertions(+), 5 deletions(-)

diff --git a/go.mod b/go.mod
index fed43bd..3ffe7da 100644
--- a/go.mod
+++ b/go.mod
@@ -11,7 +11,7 @@ require (
 	go.mau.fi/util v0.8.1-0.20241003092848-3b49d3e0b9ee
 	go.mau.fi/zeroconfig v0.1.3
 	maunium.net/go/mauflag v1.0.0
-	maunium.net/go/mautrix v0.21.1-0.20241007191346-1d8891fdb40e
+	maunium.net/go/mautrix v0.21.1-0.20241008204406-1f7f489fa97d
 )
 
 require (
diff --git a/go.sum b/go.sum
index 545d4b2..26af3bc 100644
--- a/go.sum
+++ b/go.sum
@@ -58,5 +58,5 @@ 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.1-0.20241007191346-1d8891fdb40e h1:/R+hxB3+o+6iWxhBP/ibzooMF3OU0r3WRCu5m7FeKNs=
-maunium.net/go/mautrix v0.21.1-0.20241007191346-1d8891fdb40e/go.mod h1:+fF5qsmXRCEXQZgW5ececC0PI3c7gISHTLcyftP4Bh0=
+maunium.net/go/mautrix v0.21.1-0.20241008204406-1f7f489fa97d h1:sQ0dEq7Xvak1MPey+n72sMb2Npf3jfDtsjzdDlHSLx4=
+maunium.net/go/mautrix v0.21.1-0.20241008204406-1f7f489fa97d/go.mod h1:+fF5qsmXRCEXQZgW5ececC0PI3c7gISHTLcyftP4Bh0=
diff --git a/media.go b/media.go
index 4924eb2..65ea80c 100644
--- a/media.go
+++ b/media.go
@@ -14,6 +14,9 @@ import (
 	"strconv"
 
 	"github.com/rs/zerolog"
+	"go.mau.fi/util/jsontime"
+	"go.mau.fi/util/ptr"
+
 	"maunium.net/go/mautrix"
 	"maunium.net/go/mautrix/hicli/database"
 	"maunium.net/go/mautrix/id"
@@ -25,13 +28,18 @@ var ErrBadGateway = mautrix.RespError{
 }
 
 func (gmx *Gomuks) downloadMediaFromCache(ctx context.Context, w http.ResponseWriter, entry *database.CachedMedia, force bool) bool {
-	if entry == nil || entry.Hash == nil {
+	if !entry.UseCache() {
 		if force {
 			mautrix.MNotFound.WithMessage("Media not found in cache").Write(w)
 			return true
 		}
 		return false
 	}
+	if entry.Error != nil {
+		w.Header().Set("Mau-Cached-Error", "true")
+		entry.Error.Write(w)
+		return true
+	}
 	log := zerolog.Ctx(ctx)
 	cacheFile, err := os.Open(gmx.cacheEntryToPath(entry))
 	if err != nil {
@@ -112,7 +120,40 @@ func (gmx *Gomuks) DownloadMedia(w http.ResponseWriter, r *http.Request) {
 	resp, err := gmx.Client.Client.Download(ctx, mxc)
 	if err != nil {
 		log.Err(err).Msg("Failed to download media")
-		ErrBadGateway.WithMessage(err.Error()).Write(w)
+		var httpErr mautrix.HTTPError
+		if cacheEntry == nil {
+			cacheEntry = &database.CachedMedia{
+				MXC: mxc,
+			}
+		}
+		if cacheEntry.Error == nil {
+			cacheEntry.Error = &database.MediaError{
+				ReceivedAt: jsontime.UnixMilliNow(),
+			}
+		} else {
+			cacheEntry.Error.Attempts++
+			cacheEntry.Error.ReceivedAt = jsontime.UnixMilliNow()
+		}
+		if errors.As(err, &httpErr) {
+			if httpErr.WrappedError != nil {
+				cacheEntry.Error.Matrix = ptr.Ptr(ErrBadGateway.WithMessage(httpErr.WrappedError.Error()))
+				cacheEntry.Error.StatusCode = http.StatusBadGateway
+			} else if httpErr.RespError != nil {
+				cacheEntry.Error.Matrix = httpErr.RespError
+				cacheEntry.Error.StatusCode = httpErr.Response.StatusCode
+			} else {
+				cacheEntry.Error.Matrix = ptr.Ptr(mautrix.MUnknown.WithMessage("Server returned non-JSON error with status %d", httpErr.Response.StatusCode))
+				cacheEntry.Error.StatusCode = httpErr.Response.StatusCode
+			}
+		} else {
+			cacheEntry.Error.Matrix = ptr.Ptr(ErrBadGateway.WithMessage(err.Error()))
+			cacheEntry.Error.StatusCode = http.StatusBadGateway
+		}
+		err = gmx.Client.DB.CachedMedia.Put(ctx, cacheEntry)
+		if err != nil {
+			log.Err(err).Msg("Failed to save errored cache entry")
+		}
+		cacheEntry.Error.Write(w)
 		return
 	}
 	defer func() {
@@ -159,6 +200,7 @@ func (gmx *Gomuks) DownloadMedia(w http.ResponseWriter, r *http.Request) {
 		cacheEntry.MimeType = resp.Header.Get("Content-Type")
 	}
 	cacheEntry.Hash = (*[32]byte)(fileHasher.Sum(nil))
+	cacheEntry.Error = nil
 	err = gmx.Client.DB.CachedMedia.Put(ctx, cacheEntry)
 	if err != nil {
 		log.Err(err).Msg("Failed to save cache entry")