From 8250c8ed5a06dce90a0199f879f3a2828570cedc Mon Sep 17 00:00:00 2001
From: Jonas Kaninda <me@jonaskaninda.com>
Date: Mon, 28 Oct 2024 02:26:02 +0100
Subject: [PATCH] feat: add forward client real IP

---
 pkg/helpers.go                 | 10 ++++++++++
 pkg/middleware/rate_limiter.go | 17 ++++++++++-------
 pkg/proxy.go                   |  7 ++++---
 3 files changed, 24 insertions(+), 10 deletions(-)

diff --git a/pkg/helpers.go b/pkg/helpers.go
index 553f222..a73212f 100644
--- a/pkg/helpers.go
+++ b/pkg/helpers.go
@@ -14,6 +14,7 @@ import (
 	"github.com/common-nighthawk/go-figure"
 	"github.com/jedib0t/go-pretty/v6/table"
 	"github.com/jkaninda/goma-gateway/util"
+	"net/http"
 )
 
 func Intro() {
@@ -31,3 +32,12 @@ func printRoute(routes []Route) {
 	}
 	fmt.Println(t.Render())
 }
+func getRealIP(r *http.Request) string {
+	if ip := r.Header.Get("X-Real-IP"); ip != "" {
+		return ip
+	}
+	if ip := r.Header.Get("X-Forwarded-For"); ip != "" {
+		return ip
+	}
+	return r.RemoteAddr
+}
diff --git a/pkg/middleware/rate_limiter.go b/pkg/middleware/rate_limiter.go
index 8508cb8..38bf896 100644
--- a/pkg/middleware/rate_limiter.go
+++ b/pkg/middleware/rate_limiter.go
@@ -18,7 +18,6 @@ limitations under the License.
 import (
 	"encoding/json"
 	"github.com/gorilla/mux"
-	"github.com/jkaninda/goma-gateway/internal/logger"
 	"net/http"
 	"time"
 )
@@ -52,11 +51,7 @@ func (rl *TokenRateLimiter) RateLimitMiddleware() mux.MiddlewareFunc {
 func (rl *RateLimiter) RateLimitMiddleware() mux.MiddlewareFunc {
 	return func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-
-			//TODO:
-			clientID := r.RemoteAddr
-			logger.Info(clientID)
-
+			clientID := getRealIP(r)
 			rl.mu.Lock()
 			client, exists := rl.ClientMap[clientID]
 			if !exists || time.Now().After(client.ExpiresAt) {
@@ -82,9 +77,17 @@ func (rl *RateLimiter) RateLimitMiddleware() mux.MiddlewareFunc {
 				}
 				return
 			}
-
 			// Proceed to the next handler if rate limit is not exceeded
 			next.ServeHTTP(w, r)
 		})
 	}
 }
+func getRealIP(r *http.Request) string {
+	if ip := r.Header.Get("X-Real-IP"); ip != "" {
+		return ip
+	}
+	if ip := r.Header.Get("X-Forwarded-For"); ip != "" {
+		return ip
+	}
+	return r.RemoteAddr
+}
diff --git a/pkg/proxy.go b/pkg/proxy.go
index 235b7c0..97a91b9 100644
--- a/pkg/proxy.go
+++ b/pkg/proxy.go
@@ -36,7 +36,8 @@ type ProxyRoute struct {
 // ProxyHandler proxies requests to the backend
 func (proxyRoute ProxyRoute) ProxyHandler() http.HandlerFunc {
 	return func(w http.ResponseWriter, r *http.Request) {
-		logger.Info("%s %s %s %s", r.Method, r.RemoteAddr, r.URL, r.UserAgent())
+		realIP := getRealIP(r)
+		logger.Info("%s %s %s %s", r.Method, realIP, r.URL, r.UserAgent())
 		// Set CORS headers from the cors config
 		//Update Cors Headers
 		for k, v := range proxyRoute.cors.Headers {
@@ -76,8 +77,8 @@ func (proxyRoute ProxyRoute) ProxyHandler() http.HandlerFunc {
 			r.URL.Host = targetURL.Host
 			r.URL.Scheme = targetURL.Scheme
 			r.Header.Set("X-Forwarded-Host", r.Header.Get("Host"))
-			r.Header.Set("X-Forwarded-For", r.RemoteAddr)
-			r.Header.Set("X-Real-IP", r.RemoteAddr)
+			r.Header.Set("X-Forwarded-For", realIP)
+			r.Header.Set("X-Real-IP", realIP)
 			r.Host = targetURL.Host
 		}
 		// Create proxy