From ce7d447f16b6cf4c029c919d69df72fbdfb53292 Mon Sep 17 00:00:00 2001
From: Karmanyaah Malhotra <karmanyaah.git@malhotra.cc>
Date: Tue, 21 Feb 2023 22:40:15 -0600
Subject: [PATCH] limitRequestsWithTopic

---
 server/server.go            | 26 ++++++++++++--------------
 server/server_middleware.go | 26 +++++++++++++++++++++++++-
 2 files changed, 37 insertions(+), 15 deletions(-)

diff --git a/server/server.go b/server/server.go
index e0be8d41..642f9025 100644
--- a/server/server.go
+++ b/server/server.go
@@ -437,13 +437,14 @@ func (s *Server) handleInternal(w http.ResponseWriter, r *http.Request, v *visit
 	} else if r.Method == http.MethodOptions {
 		return s.limitRequests(s.handleOptions)(w, r, v) // Should work even if the web app is not enabled, see #598
 	} else if (r.Method == http.MethodPut || r.Method == http.MethodPost) && r.URL.Path == "/" {
-		return s.limitRequests(s.transformBodyJSON(s.authorizeTopicWrite(s.handlePublish)))(w, r, v)
+		// So I don't *really* have to switch this order, since this is unrelated to UP; But, making this and matrix inconsistent is just calling for confusion, no?
+		return s.transformBodyJSON(s.limitRequestsWithTopic(s.authorizeTopicWrite(s.handlePublish)))(w, r, v)
 	} else if r.Method == http.MethodPost && r.URL.Path == matrixPushPath {
-		return s.limitRequests(s.transformMatrixJSON(s.authorizeTopicWrite(s.handlePublishMatrix)))(w, r, v)
+		return s.transformMatrixJSON(s.limitRequestsWithTopic(s.authorizeTopicWrite(s.handlePublishMatrix)))(w, r, v)
 	} else if (r.Method == http.MethodPut || r.Method == http.MethodPost) && topicPathRegex.MatchString(r.URL.Path) {
-		return s.limitRequests(s.authorizeTopicWrite(s.handlePublish))(w, r, v)
+		return s.limitRequestsWithTopic(s.authorizeTopicWrite(s.handlePublish))(w, r, v)
 	} else if r.Method == http.MethodGet && publishPathRegex.MatchString(r.URL.Path) {
-		return s.limitRequests(s.authorizeTopicWrite(s.handlePublish))(w, r, v)
+		return s.limitRequestsWithTopic(s.authorizeTopicWrite(s.handlePublish))(w, r, v)
 	} else if r.Method == http.MethodGet && jsonPathRegex.MatchString(r.URL.Path) {
 		return s.limitRequests(s.authorizeTopicRead(s.handleSubscribeJSON))(w, r, v)
 	} else if r.Method == http.MethodGet && ssePathRegex.MatchString(r.URL.Path) {
@@ -602,20 +603,17 @@ func (s *Server) handleMatrixDiscovery(w http.ResponseWriter) error {
 }
 
 func (s *Server) handlePublishWithoutResponse(r *http.Request, v *visitor) (*message, error) {
-	t, err := s.topicFromPath(r.URL.Path)
-	if err != nil {
-		return nil, err
+	vRate, ok := r.Context().Value("vRate").(*visitor)
+	if !ok {
+		return nil, errHTTPInternalError
 	}
-	vRate := v
-	if topicCountsAgainst := t.Billee(); topicCountsAgainst != nil {
-		vRate = topicCountsAgainst
+	t, ok := r.Context().Value("topic").(*topic)
+	if !ok {
+		return nil, errHTTPInternalError
 	}
 
 	if !vRate.MessageAllowed() {
-		vRate = v
-		if !v.MessageAllowed() {
-			return nil, errHTTPTooManyRequestsLimitMessages
-		}
+		return nil, errHTTPTooManyRequestsLimitMessages
 	}
 	body, err := util.Peek(r.Body, s.config.MessageLimit)
 	if err != nil {
diff --git a/server/server_middleware.go b/server/server_middleware.go
index 684253ad..e4f658b5 100644
--- a/server/server_middleware.go
+++ b/server/server_middleware.go
@@ -1,8 +1,10 @@
 package server
 
 import (
-	"heckel.io/ntfy/util"
+	"context"
 	"net/http"
+
+	"heckel.io/ntfy/util"
 )
 
 func (s *Server) limitRequests(next handleFunc) handleFunc {
@@ -16,6 +18,28 @@ func (s *Server) limitRequests(next handleFunc) handleFunc {
 	}
 }
 
+// limitRequestsWithTopic limits requests with a topic and stores the rate-limiting-subscriber and topic into request.Context
+func (s *Server) limitRequestsWithTopic(next handleFunc) handleFunc {
+	return func(w http.ResponseWriter, r *http.Request, v *visitor) error {
+		t, err := s.topicFromPath(r.URL.Path)
+		if err != nil {
+			return err
+		}
+		vRate := v
+		if topicCountsAgainst := t.Billee(); topicCountsAgainst != nil {
+			vRate = topicCountsAgainst
+		}
+		r.WithContext(context.WithValue(context.WithValue(r.Context(), "vRate", vRate), "topic", t))
+
+		if util.ContainsIP(s.config.VisitorRequestExemptIPAddrs, v.ip) {
+			return next(w, r, v)
+		} else if !vRate.RequestAllowed() {
+			return errHTTPTooManyRequestsLimitRequests
+		}
+		return next(w, r, v)
+	}
+}
+
 func (s *Server) ensureWebEnabled(next handleFunc) handleFunc {
 	return func(w http.ResponseWriter, r *http.Request, v *visitor) error {
 		if !s.config.EnableWeb {