From 0f39ead97cbef8446448b5712e27a0c815fd442f Mon Sep 17 00:00:00 2001
From: Jonas Kaninda <jonaskaninda@gmail.com>
Date: Sat, 9 Nov 2024 10:59:17 +0100
Subject: [PATCH] feat: add routes loadblancing

---
 internal/handler.go     |  2 +-
 internal/healthCheck.go |  2 +-
 internal/proxy.go       | 15 ++++++++++++++-
 internal/route.go       |  5 +++++
 internal/types.go       | 12 ++++++++++++
 internal/var.go         |  2 ++
 6 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/internal/handler.go b/internal/handler.go
index 0fe9cdd..33fec61 100644
--- a/internal/handler.go
+++ b/internal/handler.go
@@ -73,7 +73,7 @@ func (heathRoute HealthCheckRoute) HealthCheckHandler(w http.ResponseWriter, r *
 		go func() {
 			defer wg.Done()
 			if route.HealthCheck != "" {
-				err := HealthCheck(route.Destination + route.HealthCheck)
+				err := healthCheck(route.Destination + route.HealthCheck)
 				if err != nil {
 					if heathRoute.DisableRouteHealthCheckError {
 						routes = append(routes, HealthCheckRouteResponse{Name: route.Name, Status: "unhealthy", Error: "Route healthcheck errors disabled"})
diff --git a/internal/healthCheck.go b/internal/healthCheck.go
index 9d0ecd8..931f1a4 100644
--- a/internal/healthCheck.go
+++ b/internal/healthCheck.go
@@ -23,7 +23,7 @@ import (
 	"net/url"
 )
 
-func HealthCheck(healthURL string) error {
+func healthCheck(healthURL string) error {
 	healthCheckURL, err := url.Parse(healthURL)
 	if err != nil {
 		return fmt.Errorf("error parsing HealthCheck URL: %v ", err)
diff --git a/internal/proxy.go b/internal/proxy.go
index e114f25..92d82c3 100644
--- a/internal/proxy.go
+++ b/internal/proxy.go
@@ -23,6 +23,7 @@ import (
 	"net/url"
 	"slices"
 	"strings"
+	"sync/atomic"
 )
 
 // ProxyHandler proxies requests to the backend
@@ -76,8 +77,13 @@ func (proxyRoute ProxyRoute) ProxyHandler() http.HandlerFunc {
 			r.Header.Set("X-Real-IP", getRealIP(r))
 			r.Host = targetURL.Host
 		}
+		backendURL, _ := url.Parse(proxyRoute.destination)
+		if len(proxyRoute.backends) > 0 {
+			// Select the next backend URL
+			backendURL = getNextBackend(proxyRoute.backends)
+		}
 		// Create proxy
-		proxy := httputil.NewSingleHostReverseProxy(targetURL)
+		proxy := httputil.NewSingleHostReverseProxy(backendURL)
 		// Rewrite
 		if proxyRoute.path != "" && proxyRoute.rewrite != "" {
 			// Rewrite the path
@@ -92,3 +98,10 @@ func (proxyRoute ProxyRoute) ProxyHandler() http.HandlerFunc {
 		proxy.ServeHTTP(w, r)
 	}
 }
+
+// getNextBackend selects the next backend in a round-robin fashion
+func getNextBackend(backendURLs []string) *url.URL {
+	idx := atomic.AddUint32(&counter, 1) % uint32(len(backendURLs))
+	backendURL, _ := url.Parse(backendURLs[idx])
+	return backendURL
+}
diff --git a/internal/route.go b/internal/route.go
index ecfdfac..3f227f3 100644
--- a/internal/route.go
+++ b/internal/route.go
@@ -53,7 +53,10 @@ func (gatewayServer GatewayServer) Initialize() *mux.Router {
 	}
 	for _, route := range gateway.Routes {
 		if route.Path != "" {
+			if route.Destination == "" && len(route.Backends) == 0 {
+				logger.Fatal("Route %s : destination or backends should not be empty", route.Name)
 
+			}
 			// Apply middlewares to route
 			for _, mid := range route.Middlewares {
 				if mid != "" {
@@ -84,6 +87,7 @@ func (gatewayServer GatewayServer) Initialize() *mux.Router {
 								path:            route.Path,
 								rewrite:         route.Rewrite,
 								destination:     route.Destination,
+								backends:        route.Backends,
 								disableXForward: route.DisableHeaderXForward,
 								methods:         route.Methods,
 								cors:            route.Cors,
@@ -190,6 +194,7 @@ func (gatewayServer GatewayServer) Initialize() *mux.Router {
 				path:            route.Path,
 				rewrite:         route.Rewrite,
 				destination:     route.Destination,
+				backends:        route.Backends,
 				methods:         route.Methods,
 				disableXForward: route.DisableHeaderXForward,
 				cors:            route.Cors,
diff --git a/internal/types.go b/internal/types.go
index 26963a5..671161d 100644
--- a/internal/types.go
+++ b/internal/types.go
@@ -143,6 +143,8 @@ type Route struct {
 	Rewrite string `yaml:"rewrite"`
 	// Destination Defines backend URL
 	Destination string `yaml:"destination"`
+	//
+	Backends []string `yaml:"backends"`
 	// Cors contains the route cors headers
 	Cors Cors `yaml:"cors"`
 	//RateLimit int      `yaml:"rateLimit"`
@@ -197,6 +199,14 @@ type Gateway struct {
 	// Routes holds proxy routes
 	Routes []Route `yaml:"routes"`
 }
+
+type RouteHealthCheck struct {
+	Path              string `yaml:"path"`
+	Interval          int    `yaml:"interval"`
+	Timeout           int    `yaml:"timeout"`
+	HealthyStatuses   []int  `yaml:"healthyStatuses"`
+	UnhealthyStatuses []int  `yaml:"unhealthyStatuses"`
+}
 type GatewayConfig struct {
 	Version string `yaml:"version"`
 	// GatewayConfig holds Gateway config
@@ -220,6 +230,8 @@ type ProxyRoute struct {
 	path            string
 	rewrite         string
 	destination     string
+	backends        []string
+	healthCheck     RouteHealthCheck
 	methods         []string
 	cors            Cors
 	disableXForward bool
diff --git a/internal/var.go b/internal/var.go
index 204fb80..5c87d6f 100644
--- a/internal/var.go
+++ b/internal/var.go
@@ -9,3 +9,5 @@ const AccessMiddleware = "access" // access middleware
 const BasicAuth = "basic"         // basic authentication middleware
 const JWTAuth = "jwt"             // JWT authentication middleware
 const OAuth = "oauth"             // OAuth authentication middleware
+// Round-robin counter
+var counter uint32