From ce7d447f16b6cf4c029c919d69df72fbdfb53292 Mon Sep 17 00:00:00 2001 From: Karmanyaah Malhotra 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 {