Optimized rate limiter implementation

- Moved rate limiter scope into proxy router - Give IpTable a better name following clean code guideline - Optimized client IP retrieval method - Added stop channel for request counter ticker - Fixed #199 - Optimized UI for rate limit
2025-06-23 16:03:04 +02:00 · 2024-06-14 23:42:52 +08:00 · 2024-06-14 23:42:52 +08:00 · 10048150bb
commit 10048150bb
parent 85f9b297c4
11 changed files with 160 additions and 81 deletions
--- a/src/main.go
+++ b/src/main.go
@ -52,9 +52,9 @@ var logOutputToFile = flag.Bool("log", true, "Log terminal output to file")

 var (
 	name        = "Zoraxy"
-	version     = "3.0.6"
+	version     = "3.0.7"
 	nodeUUID    = "generic"
-	development = false //Set this to false to use embedded web fs
+	development = true //Set this to false to use embedded web fs
 	bootTime    = time.Now().Unix()

 	/*
--- a/src/mod/dynamicproxy/Server.go
+++ b/src/mod/dynamicproxy/Server.go
@ -72,7 +72,7 @@ func (h *ProxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 			return
 		}

-		// Rate Limit Check
+		// Rate Limit
 		if sep.RequireRateLimit {
 			err := h.handleRateLimitRouting(w, r, sep)
 			if err != nil {
--- a/src/mod/dynamicproxy/dpcore/header.go
+++ b/src/mod/dynamicproxy/dpcore/header.go
@ -91,7 +91,6 @@ func addXForwardedForHeader(req *http.Request) {
 					req.Header.Set("X-Real-Ip", strings.TrimSpace(ips[0]))
 				}
 			}
-
 		}

 	}
--- a/src/mod/dynamicproxy/dynamicproxy.go
+++ b/src/mod/dynamicproxy/dynamicproxy.go
@ -23,12 +23,12 @@ import (
 func NewDynamicProxy(option RouterOption) (*Router, error) {
 	proxyMap := sync.Map{}
 	thisRouter := Router{
-		Option:         &option,
-		ProxyEndpoints: &proxyMap,
-		Running:        false,
-		server:         nil,
-		routingRules:   []*RoutingRule{},
-		tldMap:         map[string]int{},
+		Option:           &option,
+		ProxyEndpoints:   &proxyMap,
+		Running:          false,
+		server:           nil,
+		routingRules:     []*RoutingRule{},
+		rateLimitCounter: RequestCountPerIpTable{},
 	}

 	thisRouter.mux = &ProxyHandler{
@ -85,6 +85,12 @@ func (router *Router) StartProxyService() error {
 		MinVersion:     uint16(minVersion),
 	}

+	//Start rate limitor
+	err := router.startRateLimterCounterResetTicker()
+	if err != nil {
+		return err
+	}
+
 	if router.Option.UseTls {
 		router.server = &http.Server{
 			Addr:      ":" + strconv.Itoa(router.Option.Port),
@ -129,12 +135,12 @@ func (router *Router) StartProxyService() error {
 							}
 						}

-						// Rate Limit Check
-						// if sep.RequireBasicAuth {
-						if err := handleRateLimit(w, r, sep); err != nil {
-							return
+						// Rate Limit
+						if sep.RequireRateLimit {
+							if err := router.handleRateLimit(w, r, sep); err != nil {
+								return
+							}
 						}
-						// }

 						//Validate basic auth
 						if sep.RequireBasicAuth {
@ -239,10 +245,23 @@ func (router *Router) StopProxyService() error {
 		return err
 	}

+	//Stop TLS listener
 	if router.tlsListener != nil {
 		router.tlsListener.Close()
 	}

+	//Stop rate limiter
+	if router.rateLimterStop != nil {
+		go func() {
+			// As the rate timer loop has a 1 sec ticker
+			// stop the rate limiter in go routine can prevent
+			// front end from freezing for 1 sec
+			router.rateLimterStop <- true
+		}()
+
+	}
+
+	//Stop TLS redirection (from port 80)
 	if router.tlsRedirectStop != nil {
 		router.tlsRedirectStop <- true
 	}
--- a/src/mod/dynamicproxy/ratelimit.go
+++ b/src/mod/dynamicproxy/ratelimit.go
@ -2,27 +2,27 @@ package dynamicproxy

 import (
 	"errors"
-	"log"
 	"net"
 	"net/http"
+	"strings"
 	"sync"
 	"sync/atomic"
 	"time"
 )

 // IpTable is a rate limiter implementation using sync.Map with atomic int64
-type IpTable struct {
+type RequestCountPerIpTable struct {
 	table sync.Map
 }

 // Increment the count of requests for a given IP
-func (t *IpTable) Increment(ip string) {
+func (t *RequestCountPerIpTable) Increment(ip string) {
 	v, _ := t.table.LoadOrStore(ip, new(int64))
 	atomic.AddInt64(v.(*int64), 1)
 }

 // Check if the IP is in the table and if it is, check if the count is less than the limit
-func (t *IpTable) Exceeded(ip string, limit int64) bool {
+func (t *RequestCountPerIpTable) Exceeded(ip string, limit int64) bool {
 	v, ok := t.table.Load(ip)
 	if !ok {
 		return false
@ -32,7 +32,7 @@ func (t *IpTable) Exceeded(ip string, limit int64) bool {
 }

 // Get the count of requests for a given IP
-func (t *IpTable) GetCount(ip string) int64 {
+func (t *RequestCountPerIpTable) GetCount(ip string) int64 {
 	v, ok := t.table.Load(ip)
 	if !ok {
 		return 0
@ -41,34 +41,50 @@ func (t *IpTable) GetCount(ip string) int64 {
 }

 // Clear the IP table
-func (t *IpTable) Clear() {
+func (t *RequestCountPerIpTable) Clear() {
 	t.table.Range(func(key, value interface{}) bool {
 		t.table.Delete(key)
 		return true
 	})
 }

-var ipTable = IpTable{}
-
 func (h *ProxyHandler) handleRateLimitRouting(w http.ResponseWriter, r *http.Request, pe *ProxyEndpoint) error {
-	err := handleRateLimit(w, r, pe)
+	err := h.Parent.handleRateLimit(w, r, pe)
 	if err != nil {
 		h.logRequest(r, false, 429, "ratelimit", pe.Domain)
 	}
 	return err
 }

-func handleRateLimit(w http.ResponseWriter, r *http.Request, pe *ProxyEndpoint) error {
-	ip, _, err := net.SplitHostPort(r.RemoteAddr)
-	if err != nil {
-		w.WriteHeader(500)
-		log.Println("Error resolving remote address", r.RemoteAddr, err)
-		return errors.New("internal server error")
+func (router *Router) handleRateLimit(w http.ResponseWriter, r *http.Request, pe *ProxyEndpoint) error {
+	//Get the real client-ip from request header
+	clientIP := r.RemoteAddr
+	if r.Header.Get("X-Real-Ip") == "" {
+		CF_Connecting_IP := r.Header.Get("CF-Connecting-IP")
+		Fastly_Client_IP := r.Header.Get("Fastly-Client-IP")
+		if CF_Connecting_IP != "" {
+			//Use CF Connecting IP
+			clientIP = CF_Connecting_IP
+		} else if Fastly_Client_IP != "" {
+			//Use Fastly Client IP
+			clientIP = Fastly_Client_IP
+		} else {
+			ips := strings.Split(clientIP, ",")
+			if len(ips) > 0 {
+				clientIP = strings.TrimSpace(ips[0])
+			}
+		}
 	}

-	ipTable.Increment(ip)
+	ip, _, err := net.SplitHostPort(clientIP)
+	if err != nil {
+		//Default allow passthrough on error
+		return nil
+	}

-	if ipTable.Exceeded(ip, int64(pe.RateLimit)) {
+	router.rateLimitCounter.Increment(ip)
+
+	if router.rateLimitCounter.Exceeded(ip, int64(pe.RateLimit)) {
 		w.WriteHeader(429)
 		return errors.New("rate limit exceeded")
 	}
@ -78,9 +94,26 @@ func handleRateLimit(w http.ResponseWriter, r *http.Request, pe *ProxyEndpoint)
 	return nil
 }

-func InitRateLimit() {
-	for {
-		ipTable.Clear()
-		time.Sleep(time.Second)
+// Start the ticker routine for reseting the rate limit counter every seconds
+func (r *Router) startRateLimterCounterResetTicker() error {
+	if r.rateLimterStop != nil {
+		return errors.New("another rate limiter ticker already running")
 	}
+	tickerStopChan := make(chan bool)
+	r.rateLimterStop = tickerStopChan
+
+	counterResetTicker := time.NewTicker(1 * time.Second)
+	go func() {
+		for {
+			select {
+			case <-tickerStopChan:
+				r.rateLimterStop = nil
+				return
+			case <-counterResetTicker.C:
+				r.rateLimitCounter.Clear()
+			}
+		}
+	}()
+
+	return nil
 }
--- a/src/mod/dynamicproxy/typedef.go
+++ b/src/mod/dynamicproxy/typedef.go
@ -51,8 +51,9 @@ type Router struct {
 	tlsListener    net.Listener
 	routingRules   []*RoutingRule

-	tlsRedirectStop chan bool      //Stop channel for tls redirection server
-	tldMap          map[string]int //Top level domain map, see tld.json
+	tlsRedirectStop  chan bool              //Stop channel for tls redirection server
+	rateLimterStop   chan bool              //Stop channel for rate limiter
+	rateLimitCounter RequestCountPerIpTable //Request counter for rate limter
 }

 // Auth credential for basic auth on certain endpoints
--- a/src/redirect.go
+++ b/src/redirect.go
@ -91,7 +91,7 @@ func handleToggleRedirectRegexpSupport(w http.ResponseWriter, r *http.Request) {
 	//Update the current regex support rule enable state
 	enableRegexSupport := strings.EqualFold(strings.TrimSpace(enabled), "true")
 	redirectTable.AllowRegex = enableRegexSupport
-	err = sysdb.Write("Redirect", "regex", enableRegexSupport)
+	err = sysdb.Write("redirect", "regex", enableRegexSupport)

 	if enableRegexSupport {
 		SystemWideLogger.PrintAndLog("redirect", "Regex redirect rule enabled", nil)
--- a/src/reverseproxy.go
+++ b/src/reverseproxy.go
@ -145,11 +145,6 @@ func ReverseProxtInit() {
 		})
 		SystemWideLogger.Println("Uptime Monitor background service started")
 	}()
-
-	// Init Rate Limit
-	go func() {
-		dynamicproxy.InitRateLimit()
-	}()
 }

 func ReverseProxyHandleOnOff(w http.ResponseWriter, r *http.Request) {
--- a/src/web/components/httprp.html
+++ b/src/web/components/httprp.html
@ -19,8 +19,7 @@
                    <th>Host</th>
                    <th>Destination</th>
                    <th>Virtual Directory</th>
-                    <th>Basic Auth</th>
-                    <th>Rate Limit</th>
+                    <th>Advanced Settings</th>
                    <th class="no-sort" style="min-width:150px;">Actions</th>
                </tr>
            </thead>
@ -105,11 +104,9 @@
                        </td>
                        <td data-label="" editable="true" datatype="domain">${subd.Domain} ${tlsIcon}</td>
                        <td data-label="" editable="true" datatype="vdir">${vdList}</td>
-                        <td data-label="" editable="true" datatype="basicauth">
-                            ${subd.RequireBasicAuth?`<i class="ui green check icon"></i>`:`<i class="ui grey remove icon"></i>`}
-                        </td>
-                        <td data-label="" editable="true" datatype="ratelimit">
-                            ${subd.RequireRateLimit?`<i class="ui green check icon"></i> ${subd.RateLimit}req/s`:`<i class="ui grey remove icon"></i>`}
+                        <td data-label="" editable="true" datatype="advanced">
+                            ${subd.RequireBasicAuth?`<i class="ui green check icon"></i> Basic Auth`:`<i class="ui grey remove icon"></i> Basic Auth`}<br>
+                            ${subd.RequireRateLimit?`<i class="ui green check icon"></i> Rate Limit @ ${subd.RateLimit} req/s`:`<i class="ui grey remove icon"></i> Rate Limit`}
                        </td>
                        <td class="center aligned ignoremw" editable="true" datatype="action" data-label="">
                            <div class="ui toggle tiny fitted checkbox" style="margin-bottom: -0.5em; margin-right: 0.4em;" title="Enable / Disable Rule">
@ -267,11 +264,11 @@
                    <i class="ui yellow folder icon"></i> Edit Virtual Directories
                </button>`);

-            }else if (datatype == "basicauth"){
+            }else if (datatype == "advanced"){
                let requireBasicAuth = payload.RequireBasicAuth;
-                let checkstate = "";
+                let basicAuthCheckstate = "";
                if (requireBasicAuth){
-                    checkstate = "checked";
+                    basicAuthCheckstate = "checked";
                }

                let skipWebSocketOriginCheck = payload.SkipWebSocketOriginCheck;
@ -280,16 +277,36 @@
                    wsCheckstate = "checked";
                }

+                let requireRateLimit = payload.RequireRateLimit;
+                let rateLimitCheckState = "";
+                if (requireRateLimit){
+                    rateLimitCheckState = "checked";
+                }
+                let rateLimit = payload.RateLimit;
+                if (rateLimit == 0){
+                    //This value is not set. Make it default to 100
+                    rateLimit = 100;
+                }
+                let rateLimitDisableState = "";
+                if (!payload.RequireRateLimit){
+                    rateLimitDisableState = "disabled";
+                }
+
                column.empty().append(`<div class="ui checkbox" style="margin-top: 0.4em;">
-                    <input type="checkbox" class="RequireBasicAuth" ${checkstate}>
+                    <input type="checkbox" class="RequireBasicAuth" ${basicAuthCheckstate}>
                        <label>Require Basic Auth</label>
                    </div>
-                    <button class="ui basic tiny button" style="margin-left: 0.4em; margin-top: 0.4em;" onclick="editBasicAuthCredentials('${uuid}');"><i class="ui blue user circle icon"></i> Edit Credentials</button>
+                    <br>
+                    <button class="ui basic compact tiny button" style="margin-left: 0.4em; margin-top: 0.4em;" onclick="editBasicAuthCredentials('${uuid}');"><i class="ui blue user circle icon"></i> Edit Credentials</button>
+                    <br>
+                    <button class="ui basic compact tiny button" style="margin-left: 0.4em; margin-top: 0.4em;" onclick="editCustomHeaders('${uuid}');"><i class="heading icon"></i> Custom Headers</button>
+                    <!-- <button class="ui basic compact tiny button" style="margin-left: 0.4em; margin-top: 0.4em;" onclick="editLoadBalanceOptions('${uuid}');"><i class="blue server icon"></i> Load Balance</button> -->
+
                    <div class="ui basic advance segment" style="padding: 0.4em !important; border-radius: 0.4em;">
                        <div class="ui endpointAdvanceConfig accordion" style="padding-right: 0.6em;">
                            <div class="title">
                                <i class="dropdown icon"></i>
-                                Advance Configs
+                                Security Options
                            </div>
                            <div class="content">
                                <div class="ui checkbox" style="margin-top: 0.4em;">
@ -298,23 +315,26 @@
                                    <small>Check this to allow cross-origin websocket requests</small></label>
                                </div>
                                <br>
-                                <button class="ui basic compact tiny button" style="margin-left: 0.4em; margin-top: 0.4em;" onclick="editCustomHeaders('${uuid}');"><i class="heading icon"></i> Custom Headers</button>
-                                <!-- <button class="ui basic compact tiny button" style="margin-left: 0.4em; margin-top: 0.4em;" onclick="editLoadBalanceOptions('${uuid}');"><i class="blue server icon"></i> Load Balance</button> -->
+                                <div class="ui checkbox" style="margin-top: 0.4em;">
+                                    <input type="checkbox" onchange="handleToggleRateLimitInput();" class="RequireRateLimit" ${rateLimitCheckState}>
+                                    <label>Require Rate Limit</label>
+                                </div><br>
+                                <div class="ui mini right labeled fluid input ${rateLimitDisableState}" style="margin-top: 0.4em;">
+                                    <input type="number" class="RateLimit" value="${rateLimit}" min="1" >
+                                    <label class="ui basic label">
+                                        req / sec / IP
+                                    </label>
+                                </div>
                            </div>
                        </div>
                    <div>
                `);

            } else if (datatype == "ratelimit"){
-                let requireRateLimit = payload.RequireRateLimit;
-                let checkstate = "";
-                if (requireRateLimit){
-                    checkstate = "checked";
-                }
-                let rateLimit = payload.RateLimit;
-
-                column.empty().append(`<div class="ui checkbox" style="margin-top: 0.4em;">
-                    <input type="checkbox" class="RequireRateLimit" ${checkstate}>
+               
+                column.empty().append(`
+                    <div class="ui checkbox" style="margin-top: 0.4em;">
+                        <input type="checkbox" class="RequireRateLimit" ${checkstate}>
                        <label>Require Rate Limit</label>
                    </div>
                    <div class="ui mini fluid input">
@ -352,6 +372,17 @@
        $("#httpProxyList").find(".editBtn").addClass("disabled");
    }

+    //handleToggleRateLimitInput will get trigger if the "require rate limit" checkbox
+    // is changed and toggle the disable state of the rate limit input field
+    function handleToggleRateLimitInput(){
+        let isRateLimitEnabled = $("#httpProxyList input.RequireRateLimit")[0].checked;
+        if (isRateLimitEnabled){
+            $("#httpProxyList input.RateLimit").parent().removeClass("disabled");
+        }else{
+            $("#httpProxyList input.RateLimit").parent().addClass("disabled");
+        }
+    }
+
    function exitProxyInlineEdit(){
        listProxyEndpoints();
        $("#httpProxyList").find(".editBtn").removeClass("disabled");
@ -463,10 +494,6 @@
        })
    }

-    /* Access List handling */
-    
-    
-
    //Bind on tab switch events
    tabSwitchEventBind["httprp"] = function(){
        listProxyEndpoints();
--- a/src/web/components/rules.html
+++ b/src/web/components/rules.html
@ -81,8 +81,13 @@
                                </div>
                                <div class="field">
                                    <label>Rate Limit</label>
-                                    <input type="number" id="proxyRateLimit" placeholder="100" min="1" max="1000" value="100">
-                                    <small>The Rate Limit is applied to the whole proxy endpoint. If the number of requests exceeds the limit, the proxy will return a 429 error code.</small>
+                                    <div class="ui fluid right labeled input">
+                                        <input type="number" id="proxyRateLimit" placeholder="100" min="1" max="1000" value="100">
+                                        <div class="ui basic label">
+                                           req / sec / IP
+                                        </div>
+                                    </div>
+                                    <small>Return a 429 error code if request rate exceed the rate limit.</small>
                                </div>
                                <div class="field">
                                    <div class="ui checkbox">
@ -282,9 +287,9 @@
    
    function toggleRateLimit() {
        if ($("#requireRateLimit").parent().checkbox("is checked")) {
-            $("#proxyRateLimit").parent().removeClass("disabled");
+            $("#proxyRateLimit").parent().parent().removeClass("disabled");
        } else {
-            $("#proxyRateLimit").parent().addClass("disabled");
+            $("#proxyRateLimit").parent().parent().addClass("disabled");
        }
    }
    $("#requireRateLimit").on('change', toggleRateLimit);
@ -422,9 +427,9 @@
        initNewProxyRuleAccessDropdownList();
    }

-    $(document).ready(function(){
-        $("#advanceProxyRules").accordion();
-        $("#newProxyRuleAccessFilter").parent().dropdown();
-    });
+   
+    $("#advanceProxyRules").accordion();
+    $("#newProxyRuleAccessFilter").parent().dropdown();
+   
    
 </script>
--- a/src/web/components/status.html
+++ b/src/web/components/status.html
@ -74,7 +74,7 @@
        <input type="text" id="incomingPort" placeholder="Incoming Port" value="80">
        <button class="ui green notloopbackOnly button" style="background: linear-gradient(60deg, #27e7ff, #00ca52);" onclick="handlePortChange();"><i class="ui checkmark icon"></i> Apply</button>
    </div>
-    <br><br>
+    <br>
    <div id="tls" class="ui toggle notloopbackOnly checkbox">
        <input type="checkbox">
        <label>Use TLS to serve proxy request</label>