Merge pull request #7 from karthikbhandary2/karthik

Ratelimiting added

Author: karthikbhandary2

bin/build-errors.log

@@ -1,8 +1,12 @@
 package main
 
 import (
+	"errors"
 	"fmt"
 	"net/http"
+	"os"
+	"os/signal"
+	"syscall"
 	"time"
 
 	"github.com/go-chi/chi/v5"
@@ -12,10 +16,12 @@ import (
 	_ "github.com/karthikbhandary2/Social/docs"
 	"github.com/karthikbhandary2/Social/internal/auth"
 	"github.com/karthikbhandary2/Social/internal/mailer"
+	"github.com/karthikbhandary2/Social/internal/ratelimiter"
 	"github.com/karthikbhandary2/Social/internal/store"
 	"github.com/karthikbhandary2/Social/internal/store/cache"
 	"github.com/swaggo/http-swagger/v2"
 	"go.uber.org/zap"
+	"golang.org/x/net/context"
 )
 
 type application struct {
@@ -25,6 +31,7 @@ type application struct {
 	logger        *zap.SugaredLogger
 	mailer        mailer.Client
 	authenticator auth.Authenticator
+	rateLimiter   ratelimiter.Limiter
 }
 
 type config struct {
@@ -36,6 +43,7 @@ type config struct {
 	frontendURL string
 	auth        authConfig
 	redisCfg    redisConfig
+	rateLimiter ratelimiter.Config
 }
 
 type redisConfig struct {
@@ -94,10 +102,10 @@ func (app *application) mount() http.Handler {
 	r.Use(middleware.Logger)
 	r.Use(middleware.RealIP)
 	r.Use(middleware.RequestID)
-	r.Use(middleware.Timeout(60 * time.Second))
+	r.Use(app.RateLimiterMiddleware)
 
 	r.Route("/v1", func(r chi.Router) {
-		r.With(app.BasicAuthMiddleware()).Get("/health", app.healthCheckHandler)
+		r.Get("/health", app.healthCheckHandler)
 
 		docsURL := fmt.Sprintf("%s/swagger/doc.json", app.config.addr)
 		r.Get("/swagger/*", httpSwagger.Handler(httpSwagger.URL(docsURL)))
@@ -149,6 +157,30 @@ func (app *application) run(mux http.Handler) error {
 		ReadTimeout:  time.Second * 10, // It should be less than writes
 		IdleTimeout:  time.Minute,
 	}
+
+	shutdown := make(chan error)
+	go func() {
+		quit := make(chan os.Signal, 1)
+		signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
+		s := <-quit
+
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+
+		app.logger.Infow("signal caught", "signal", s.String())
+
+		shutdown <- srv.Shutdown(ctx)
+	}()
 	app.logger.Infow("server has started", "addr", app.config.addr, "env", app.config.env)
-	return srv.ListenAndServe()
+
+	err := srv.ListenAndServe()
+	if !errors.Is(err, http.ErrServerClosed) {
+		return err
+	}
+	err = <-shutdown
+	if err != nil {
+		return err
+	}
+	app.logger.Infow("server has stopped", "addr", app.config.addr, "env", app.config.env)
+	return nil
 }

bin/main

@@ -0,0 +1,54 @@
+package main
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/karthikbhandary2/Social/internal/ratelimiter"
+)
+
+func TestRateLimiterMiddleware(t *testing.T) {
+	cfg := config{
+		rateLimiter: ratelimiter.Config{
+			RequestsPerTimeFrame: 20,
+			TimeFrame:            time.Second * 5,
+			Enabled:              true,
+		},
+		addr: ":8080",
+	}
+
+	app := newTestApplication(t, cfg)
+	ts := httptest.NewServer(app.mount())
+	defer ts.Close()
+
+	client := &http.Client{}
+	mockIP := "192.168.1.1"
+	marginOfError := 2
+
+	for i := 0; i < cfg.rateLimiter.RequestsPerTimeFrame+marginOfError; i++ {
+		req, err := http.NewRequest("GET", ts.URL+"/v1/health", nil)
+		if err != nil {
+			t.Fatalf("could not create request: %v", err)
+		}
+
+		req.Header.Set("X-Forwarded-For", mockIP)
+
+		resp, err := client.Do(req)
+		if err != nil {
+			t.Fatalf("could not send request: %v", err)
+		}
+		defer resp.Body.Close()
+
+		if i < cfg.rateLimiter.RequestsPerTimeFrame {
+			if resp.StatusCode != http.StatusOK {
+				t.Errorf("expected status OK; got %v", resp.Status)
+			}
+		} else {
+			if resp.StatusCode != http.StatusTooManyRequests {
+				t.Errorf("expected status Too Many Requests; got %v", resp.Status)
+			}
+		}
+	}
+}

cmd/api/api.go

@@ -10,8 +10,8 @@ import (
 
 	"github.com/golang-jwt/jwt/v5"
 	"github.com/google/uuid"
-	"github.com/karthikbhandary2/Social/internal/store"
 	"github.com/karthikbhandary2/Social/internal/mailer"
+	"github.com/karthikbhandary2/Social/internal/store"
 )
 
 type RegisterUserPayload struct {

cmd/api/api_test.go

@@ -39,3 +39,11 @@ func (app *application) notFound(w http.ResponseWriter, r *http.Request, err err
 	app.logger.Warnf("not found error", "method", r.Method,"path", r.URL.Path,"error",err)
 	writeJSON(w, http.StatusNotFound, err.Error())
 }
+
+func (app *application) rateLimitExceeded(w http.ResponseWriter, r *http.Request, retryAfter string) {
+	app.logger.Warnw("rate limit exceeded", "method", r.Method, "path", r.URL.Path)
+
+	w.Header().Set("Retry-After", retryAfter)
+
+	writeJSONError(w, http.StatusTooManyRequests, "rate limit exceeded, retry after: "+retryAfter)
+}

cmd/api/auth.go

@@ -8,6 +8,7 @@ import (
 	"github.com/karthikbhandary2/Social/internal/db"
 	"github.com/karthikbhandary2/Social/internal/env"
 	"github.com/karthikbhandary2/Social/internal/mailer"
+	"github.com/karthikbhandary2/Social/internal/ratelimiter"
 	"github.com/karthikbhandary2/Social/internal/store"
 	"github.com/karthikbhandary2/Social/internal/store/cache"
 	"go.uber.org/zap"
@@ -72,6 +73,11 @@ func main() {
 				iss:    "social",
 			},
 		},
+		rateLimiter: ratelimiter.Config{
+			RequestsPerTimeFrame: env.GetInt("RATELIMITER_REQUESTS_COUNT", 20),
+			TimeFrame:            time.Second * 5,
+			Enabled:              env.GetBool("RATE_LIMITER_ENABLED", true),
+		},
 	}
 
 	//logger
@@ -93,6 +99,13 @@ func main() {
 		rdb = cache.NewRedisClient(cfg.redisCfg.addr, cfg.redisCfg.pw, cfg.redisCfg.db)
 		logger.Info("Redis connection pool established")
 	}
+
+	//ratelimiter
+	rateLimiter := ratelimiter.NewFixedWindowLimiter(
+		cfg.rateLimiter.RequestsPerTimeFrame,
+		cfg.rateLimiter.TimeFrame,
+	)
+
 	store := store.NewStorage(db)
 	cacheStorage := cache.NewRedisStorage(rdb)
 	mailer := mailer.NewSendgrid(cfg.mail.sendGrid.apiKey, cfg.mail.fromEmail)
@@ -105,6 +118,7 @@ func main() {
 		logger:        logger,
 		mailer:        mailer,
 		authenticator: jwtAuthenticator,
+		rateLimiter:   rateLimiter,
 	}
 
 	mux := app.mount()

cmd/api/errors.go

@@ -144,4 +144,16 @@ func (app *application) getUser(ctx context.Context, userID int64) (*store.User,
 	return user, nil
 
 
+}
+
+func (app *application) RateLimiterMiddleware(next http.Handler) http.Handler {
+	return http.HandlerFunc(func (w http.ResponseWriter, r *http.Request)  {
+		if app.config.rateLimiter.Enabled {
+			if allow, retryAfter := app.rateLimiter.Allow(r.RemoteAddr); !allow {
+				app.rateLimitExceeded(w, r, retryAfter.String())
+				return
+			}
+		}
+		next.ServeHTTP(w, r)
+	})
 }

cmd/api/main.go

@@ -6,27 +6,40 @@ import (
 	"testing"
 
 	"github.com/karthikbhandary2/Social/internal/auth"
+	"github.com/karthikbhandary2/Social/internal/ratelimiter"
 	"github.com/karthikbhandary2/Social/internal/store"
 	"github.com/karthikbhandary2/Social/internal/store/cache"
 	"go.uber.org/zap"
 )
 
-func newTestApplication(t *testing.T) *application {
+func newTestApplication(t *testing.T, cfg config) *application {
 	t.Helper()
 
+	logger := zap.NewNop().Sugar()
+	// Uncomment to enable logs
 	// logger := zap.Must(zap.NewProduction()).Sugar()
-	logger := zap.Must(zap.NewProduction()).Sugar()
 	mockStore := store.NewMockStore()
 	mockCacheStore := cache.NewMockStore()
+
 	testAuth := &auth.TestAuthenticator{}
+
+	// Rate limiter
+	rateLimiter := ratelimiter.NewFixedWindowLimiter(
+		cfg.rateLimiter.RequestsPerTimeFrame,
+		cfg.rateLimiter.TimeFrame,
+	)
+
 	return &application{
-		logger: logger,
-		store: mockStore,
-		cacheStorage: mockCacheStore,
+		logger:        logger,
+		store:         mockStore,
+		cacheStorage:  mockCacheStore,
 		authenticator: testAuth,
+		config:        cfg,
+		rateLimiter:   rateLimiter,
 	}
 }
 
+
 func executeRequest(req *http.Request, mux http.Handler) *httptest.ResponseRecorder {
 	rr := httptest.NewRecorder()
 	mux.ServeHTTP(rr, req)

cmd/api/middleware.go

@@ -7,7 +7,12 @@ import (
 )
 
 func TestGetUser(t *testing.T) {
-	app := newTestApplication(t)
+	withRedis := config{
+		redisCfg: redisConfig{
+			enabled: true,
+		},
+	}
+	app := newTestApplication(t, withRedis)
 	mux := app.mount()
 	testToken, err := app.authenticator.GenerateToken(nil)
 	if err != nil {