Ratelimit requests to /media/r0/download|upload (#2020)

* Add /media/r0/config handler

Signed-off-by: Till Faelligen <tfaelligen@gmail.com>

* Add rate limiting to media api

* Rename variable

* Add passing tests

* Don't send multiple headers

Co-authored-by: Neil Alexander <neilalexander@users.noreply.github.com>

2 files changed, 42 insertions, 6 deletions

diff --git a/mediaapi/mediaapi.go b/mediaapi/mediaapi.go
index 811d8e4a..c010981c 100644
--- a/mediaapi/mediaapi.go
+++ b/mediaapi/mediaapi.go
@@ -26,7 +26,9 @@ import (
 
 // AddPublicRoutes sets up and registers HTTP handlers for the MediaAPI component.
 func AddPublicRoutes(
-	router *mux.Router, cfg *config.MediaAPI,
+	router *mux.Router,
+	cfg *config.MediaAPI,
+	rateLimit *config.RateLimiting,
 	userAPI userapi.UserInternalAPI,
 	client *gomatrixserverlib.Client,
 ) {
@@ -36,6 +38,6 @@ func AddPublicRoutes(
 	}
 
 	routing.Setup(
-		router, cfg, mediaDB, userAPI, client,
+		router, cfg, rateLimit, mediaDB, userAPI, client,
 	)
 }
diff --git a/mediaapi/routing/routing.go b/mediaapi/routing/routing.go
index 917a8596..44f9a9d6 100644
--- a/mediaapi/routing/routing.go
+++ b/mediaapi/routing/routing.go
@@ -15,16 +15,16 @@
 package routing
 
 import (
+	"encoding/json"
 	"net/http"
 	"strings"
 
-	userapi "github.com/matrix-org/dendrite/userapi/api"
-
 	"github.com/gorilla/mux"
 	"github.com/matrix-org/dendrite/internal/httputil"
 	"github.com/matrix-org/dendrite/mediaapi/storage"
 	"github.com/matrix-org/dendrite/mediaapi/types"
 	"github.com/matrix-org/dendrite/setup/config"
+	userapi "github.com/matrix-org/dendrite/userapi/api"
 	"github.com/matrix-org/gomatrixserverlib"
 	"github.com/matrix-org/util"
 	"github.com/prometheus/client_golang/prometheus"
@@ -32,6 +32,12 @@ import (
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 )
 
+// configResponse is the response to GET /_matrix/media/r0/config
+// https://matrix.org/docs/spec/client_server/latest#get-matrix-media-r0-config
+type configResponse struct {
+	UploadSize config.FileSizeBytes `json:"m.upload.size"`
+}
+
 // Setup registers the media API HTTP handlers
 //
 // Due to Setup being used to call many other functions, a gocyclo nolint is
@@ -40,10 +46,13 @@ import (
 func Setup(
 	publicAPIMux *mux.Router,
 	cfg *config.MediaAPI,
+	rateLimit *config.RateLimiting,
 	db storage.Database,
 	userAPI userapi.UserInternalAPI,
 	client *gomatrixserverlib.Client,
 ) {
+	rateLimits := httputil.NewRateLimits(rateLimit)
+
 	r0mux := publicAPIMux.PathPrefix("/r0").Subrouter()
 	v1mux := publicAPIMux.PathPrefix("/v1").Subrouter()
 
@@ -54,31 +63,46 @@ func Setup(
 	uploadHandler := httputil.MakeAuthAPI(
 		"upload", userAPI,
 		func(req *http.Request, dev *userapi.Device) util.JSONResponse {
+			if r := rateLimits.Limit(req); r != nil {
+				return *r
+			}
 			return Upload(req, cfg, dev, db, activeThumbnailGeneration)
 		},
 	)
 
+	configHandler := httputil.MakeAuthAPI("config", userAPI, func(req *http.Request, device *userapi.Device) util.JSONResponse {
+		if r := rateLimits.Limit(req); r != nil {
+			return *r
+		}
+		return util.JSONResponse{
+			Code: http.StatusOK,
+			JSON: configResponse{UploadSize: *cfg.MaxFileSizeBytes},
+		}
+	})
+
 	r0mux.Handle("/upload", uploadHandler).Methods(http.MethodPost, http.MethodOptions)
+	r0mux.Handle("/config", configHandler).Methods(http.MethodGet, http.MethodOptions)
 	v1mux.Handle("/upload", uploadHandler).Methods(http.MethodPost, http.MethodOptions)
 
 	activeRemoteRequests := &types.ActiveRemoteRequests{
 		MXCToResult: map[string]*types.RemoteRequestResult{},
 	}
 
-	downloadHandler := makeDownloadAPI("download", cfg, db, client, activeRemoteRequests, activeThumbnailGeneration)
+	downloadHandler := makeDownloadAPI("download", cfg, rateLimits, db, client, activeRemoteRequests, activeThumbnailGeneration)
 	r0mux.Handle("/download/{serverName}/{mediaId}", downloadHandler).Methods(http.MethodGet, http.MethodOptions)
 	r0mux.Handle("/download/{serverName}/{mediaId}/{downloadName}", downloadHandler).Methods(http.MethodGet, http.MethodOptions)
 	v1mux.Handle("/download/{serverName}/{mediaId}", downloadHandler).Methods(http.MethodGet, http.MethodOptions)                // TODO: remove when synapse is fixed
 	v1mux.Handle("/download/{serverName}/{mediaId}/{downloadName}", downloadHandler).Methods(http.MethodGet, http.MethodOptions) // TODO: remove when synapse is fixed
 
 	r0mux.Handle("/thumbnail/{serverName}/{mediaId}",
-		makeDownloadAPI("thumbnail", cfg, db, client, activeRemoteRequests, activeThumbnailGeneration),
+		makeDownloadAPI("thumbnail", cfg, rateLimits, db, client, activeRemoteRequests, activeThumbnailGeneration),
 	).Methods(http.MethodGet, http.MethodOptions)
 }
 
 func makeDownloadAPI(
 	name string,
 	cfg *config.MediaAPI,
+	rateLimits *httputil.RateLimits,
 	db storage.Database,
 	client *gomatrixserverlib.Client,
 	activeRemoteRequests *types.ActiveRemoteRequests,
@@ -99,6 +123,16 @@ func makeDownloadAPI(
 		// Content-Type will be overridden in case of returning file data, else we respond with JSON-formatted errors
 		w.Header().Set("Content-Type", "application/json")
 
+		// Ratelimit requests
+		if r := rateLimits.Limit(req); r != nil {
+			if err := json.NewEncoder(w).Encode(r); err != nil {
+				w.WriteHeader(http.StatusInternalServerError)
+				return
+			}
+			w.WriteHeader(http.StatusTooManyRequests)
+			return
+		}
+
 		vars, _ := httputil.URLDecodeMapValues(mux.Vars(req))
 		serverName := gomatrixserverlib.ServerName(vars["serverName"])