add rate limit

main.go

@@ -6,7 +6,9 @@ import (
 	"log"
 	"net/http"
 	"os"
+	"time"
 
+	"git.1e99.eu/1e99/passed/middlewares"
 	"git.1e99.eu/1e99/passed/routes"
 	"git.1e99.eu/1e99/passed/storage"
 )
@@ -24,7 +26,14 @@ func run() error {
 
 	mux := http.NewServeMux()
 	mux.Handle("GET /", routes.ServeFiles(embedFS, "static"))
-	mux.Handle("POST /api/password", routes.CreatePassword(storage, 12*1024, base64.StdEncoding))
+	mux.Handle(
+		"POST /api/password",
+		middlewares.RateLimiter(
+			routes.CreatePassword(storage, 12*1024, base64.StdEncoding),
+			1*time.Minute,
+			5,
+		),
+	)
 	mux.Handle("GET /api/password/{id}", routes.GetPassword(storage, base64.StdEncoding))
 	mux.Handle("HEAD /api/password/{id}", routes.HasPassword(storage))
 
@@ -34,7 +43,7 @@ func run() error {
 		address = ":3000"
 	}
 
-	err = http.ListenAndServe(address, routes.Logger(mux))
+	err = http.ListenAndServe(address, middlewares.Logger(mux))
 	if err != nil {
 		return err
 	}

middlewares/logger.go

@@ -1,4 +1,4 @@
-package routes
+package middlewares
 
 import (
 	"log"
@@ -7,7 +7,7 @@ import (
 
 func Logger(handler http.Handler) http.HandlerFunc {
 	return func(res http.ResponseWriter, req *http.Request) {
-		log.Printf("%-80s", req.URL.Path)
+		log.Printf("%-30s %-80s", req.RemoteAddr, req.URL.Path)
 		handler.ServeHTTP(res, req)
 	}
 }

middlewares/rate_limiter.go

@@ -0,0 +1,44 @@
+package middlewares
+
+import (
+	"net/http"
+	"sync"
+	"time"
+)
+
+func RateLimiter(handler http.Handler, clearInterval time.Duration, maxRequests int) http.HandlerFunc {
+	requests := make(map[string]int)
+	lock := sync.Mutex{}
+
+	ticker := time.NewTicker(clearInterval)
+	go func() {
+		for {
+			<-ticker.C
+
+			lock.Lock()
+			clear(requests)
+			lock.Unlock()
+		}
+	}()
+
+	return func(res http.ResponseWriter, req *http.Request) {
+		addr := req.RemoteAddr
+
+		lock.Lock()
+		count, found := requests[addr]
+		if !found {
+			count = 0
+		}
+
+		count += 1
+		requests[addr] = count
+		lock.Unlock()
+
+		if count > maxRequests {
+			http.Error(res, "Too many requests", http.StatusTooManyRequests)
+			return
+		}
+
+		handler.ServeHTTP(res, req)
+	}
+}