Merge pull request #151 from bluwumeaway/feature/ratelimit-middleware

add rate limiting middleware

pkg/api/api.go

···
 	"os"
 	"slices"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/NYTimes/gziphandler"
···
 	"github.com/julienschmidt/httprouter"
 	"github.com/rs/cors"
 	sloghttp "github.com/samber/slog-http"
+	"golang.org/x/time/rate"
 
 	"stream.place/streamplace/js/app"
 	"stream.place/streamplace/pkg/atproto"
···
 	Bus      *bus.Bus
 	ATSync   *atproto.ATProtoSynchronizer
 	Director *director.Director
+
+	connTracker *WebsocketTracker
+
+	limiters   map[string]*rate.Limiter
+	limitersMu sync.Mutex
+}
+
+type WebsocketTracker struct {
+	connections   map[string]int
+	maxConnsPerIP int
+	mu            sync.RWMutex
 }
 
 func MakeStreamplaceAPI(cli *config.CLI, mod model.Model, signer *eip712.EIP712Signer, noter notifications.FirebaseNotifier, mm *media.MediaManager, ms media.MediaSigner, bus *bus.Bus, atsync *atproto.ATProtoSynchronizer, d *director.Director) (*StreamplaceAPI, error) {
···
 		Bus:              bus,
 		ATSync:           atsync,
 		Director:         d,
+		connTracker:      NewWebsocketTracker(5),
+		limiters:         make(map[string]*rate.Limiter),
 	}
 	a.Mimes, err = updater.GetMimes()
 	if err != nil {
···
 	handler := sloghttp.Recovery(router)
 	handler = cors.AllowAll().Handler(handler)
 	handler = sloghttp.New(slog.Default())(handler)
+	handler = a.RateLimitMiddleware(ctx)(handler)
 
 	return handler, nil
 }
···
 	}
 }
 
+func (a *StreamplaceAPI) RateLimitMiddleware(ctx context.Context) func(http.Handler) http.Handler {
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
+			ip, _, err := net.SplitHostPort(req.RemoteAddr)
+			if err != nil {
+				ip = req.RemoteAddr
+			}
+
+			limiter := a.getLimiter(ip)
+
+			if !limiter.Allow() {
+				log.Warn(ctx, "rate limit exceeded", "ip", ip, "path", req.URL.Path)
+				apierrors.WriteHTTPTooManyRequests(w, "rate limit exceeded")
+				return
+			}
+
+			next.ServeHTTP(w, req)
+		})
+	}
+}
+
 func (a *StreamplaceAPI) ServeHTTP(ctx context.Context) error {
 	handler, err := a.Handler(ctx)
 	if err != nil {
···
 		w.WriteHeader(200)
 	}
 }
+
+func (a *StreamplaceAPI) getLimiter(ip string) *rate.Limiter {
+	a.limitersMu.Lock()
+	defer a.limitersMu.Unlock()
+
+	limiter, exists := a.limiters[ip]
+	if !exists {
+		// 5 actions per second with a burst of 3
+		limiter = rate.NewLimiter(rate.Limit(10.0), 8)
+		a.limiters[ip] = limiter
+	}
+
+	return limiter
+}
+
+func NewWebsocketTracker(maxConns int) *WebsocketTracker {
+	return &WebsocketTracker{
+		connections:   make(map[string]int),
+		maxConnsPerIP: maxConns,
+	}
+}
+
+func (t *WebsocketTracker) AddConnection(ip string) bool {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+
+	count := t.connections[ip]
+
+	if count >= t.maxConnsPerIP {
+		return false
+	}
+
+	t.connections[ip] = count + 1
+	return true
+}
+
+func (t *WebsocketTracker) RemoveConnection(ip string) {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+
+	count := t.connections[ip]
+	if count > 0 {
+		t.connections[ip] = count - 1
+	}
+
+	if t.connections[ip] == 0 {
+		delete(t.connections, ip)
+	}
+}

pkg/api/websocket.go

···
 import (
 	"context"
 	"encoding/json"
+	"net"
 	"net/http"
 	"time"
 
 	"github.com/google/uuid"
 	"github.com/gorilla/websocket"
 	"github.com/julienschmidt/httprouter"
+
 	apierrors "stream.place/streamplace/pkg/errors"
 	"stream.place/streamplace/pkg/log"
 	"stream.place/streamplace/pkg/renditions"
···
 func (a *StreamplaceAPI) HandleWebsocket(ctx context.Context) httprouter.Handle {
 	ctx = log.WithLogValues(ctx, "func", "HandleWebsocket")
 	return func(w http.ResponseWriter, req *http.Request, params httprouter.Params) {
+		ip, _, err := net.SplitHostPort(req.RemoteAddr)
+		if err != nil {
+			ip = req.RemoteAddr
+		}
+
+		if !a.connTracker.AddConnection(ip) {
+			log.Warn(ctx, "rate limit exceeded", "ip", ip, "path", req.URL.Path)
+			apierrors.WriteHTTPTooManyRequests(w, "rate limit exceeded")
+			return
+		}
+
+		defer a.connTracker.RemoveConnection(ip)
+
 		uu, _ := uuid.NewV7()
-		ctx = log.WithLogValues(ctx, "uuid", uu.String(), "remoteAddr", req.RemoteAddr, "url", req.URL.String())
+		connID := uu.String()
+
+		ctx = log.WithLogValues(ctx, "uuid", connID, "remoteAddr", req.RemoteAddr, "url", req.URL.String())
 		log.Log(ctx, "websocket opened")
 		spmetrics.WebsocketsOpen.Inc()
 		defer spmetrics.WebsocketsOpen.Dec()
···
 		ctx, cancel := context.WithCancel(ctx)
 		defer cancel()
 		defer conn.Close()
+
 		initialBurst := make(chan any, 200)
 		err = conn.SetReadDeadline(time.Now().Add(30 * time.Second))
 		if err != nil {

pkg/errors/errors.go

···
 func WriteHTTPNotImplemented(w http.ResponseWriter, msg string, err error) APIError {
 	return writeHttpError(w, msg, http.StatusNotImplemented, err)
 }
+
+func WriteHTTPTooManyRequests(w http.ResponseWriter, msg string) APIError {
+	return writeHttpError(w, msg, http.StatusTooManyRequests, nil)
+}