Merge pull request #196 from Kirari04/main

[ENHANCEMENTS] Add Rate Limits Limits to Zoraxy

.gitignore

@@ -34,3 +34,9 @@ docker/ImagePublisher.sh
 src/mod/acme/test/stackoverflow.pem
 /tools/dns_challenge_update/code-gen/acmedns
 /tools/dns_challenge_update/code-gen/lego
+src/tmp/localhost.key
+src/tmp/localhost.pem
+src/www/html/index.html
+src/sys.uuid
+src/zoraxy
+src/log/

src/mod/dynamicproxy/Server.go

@@ -72,6 +72,14 @@ func (h *ProxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 			return
 		}
 
+		// Rate Limit Check
+		if sep.RequireRateLimit {
+			err := h.handleRateLimitRouting(w, r, sep)
+			if err != nil {
+				return
+			}
+		}
+
 		//Validate basic auth
 		if sep.RequireBasicAuth {
 			err := h.handleBasicAuthRouting(w, r, sep)

src/mod/dynamicproxy/dynamicproxy.go

@@ -129,6 +129,13 @@ func (router *Router) StartProxyService() error {
 							}
 						}
 
+						// Rate Limit Check
+						// if sep.RequireBasicAuth {
+						if err := handleRateLimit(w, r, sep); err != nil {
+							return
+						}
+						// }
+
 						//Validate basic auth
 						if sep.RequireBasicAuth {
 							err := handleBasicAuth(w, r, sep)

src/mod/dynamicproxy/ratelimit.go

@@ -0,0 +1,86 @@
+package dynamicproxy
+
+import (
+	"errors"
+	"log"
+	"net"
+	"net/http"
+	"sync"
+	"sync/atomic"
+	"time"
+)
+
+// IpTable is a rate limiter implementation using sync.Map with atomic int64
+type IpTable struct {
+	table sync.Map
+}
+
+// Increment the count of requests for a given IP
+func (t *IpTable) 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 {
+	v, ok := t.table.Load(ip)
+	if !ok {
+		return false
+	}
+	count := atomic.LoadInt64(v.(*int64))
+	return count >= limit
+}
+
+// Get the count of requests for a given IP
+func (t *IpTable) GetCount(ip string) int64 {
+	v, ok := t.table.Load(ip)
+	if !ok {
+		return 0
+	}
+	return atomic.LoadInt64(v.(*int64))
+}
+
+// Clear the IP table
+func (t *IpTable) 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)
+	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")
+	}
+
+	ipTable.Increment(ip)
+
+	if ipTable.Exceeded(ip, int64(pe.RateLimit)) {
+		w.WriteHeader(429)
+		return errors.New("rate limit exceeded")
+	}
+
+	// log.Println("Rate limit check", ip, ipTable.GetCount(ip))
+
+	return nil
+}
+
+func InitRateLimit() {
+	for {
+		ipTable.Clear()
+		time.Sleep(time.Second)
+	}
+}

src/mod/dynamicproxy/typedef.go

@@ -124,6 +124,10 @@ type ProxyEndpoint struct {
 	BasicAuthCredentials    []*BasicAuthCredentials   //Basic auth credentials
 	BasicAuthExceptionRules []*BasicAuthExceptionRule //Path to exclude in a basic auth enabled proxy target
 
+	// Rate Limiting
+	RequireRateLimit bool
+	RateLimit        int64 // Rate limit in requests per second
+
 	//Access Control
 	AccessFilterUUID string //Access filter ID
 

src/reverseproxy.go

@@ -146,6 +146,10 @@ func ReverseProxtInit() {
 		SystemWideLogger.Println("Uptime Monitor background service started")
 	}()
 
+	// Init Rate Limit
+	go func() {
+		dynamicproxy.InitRateLimit()
+	}()
 }
 
 func ReverseProxyHandleOnOff(w http.ResponseWriter, r *http.Request) {
@@ -229,6 +233,26 @@ func ReverseProxyHandleAddEndpoint(w http.ResponseWriter, r *http.Request) {
 
 	requireBasicAuth := (rba == "true")
 
+	// Require Rate Limiting?
+	rl, _ := utils.PostPara(r, "rate")
+	if rl == "" {
+		rl = "false"
+	}
+	requireRateLimit := (rl == "true")
+	rlnum, _ := utils.PostPara(r, "ratenum")
+	if rlnum == "" {
+		rlnum = "0"
+	}
+	proxyRateLimit, err := strconv.ParseInt(rlnum, 10, 64)
+	if err != nil {
+		utils.SendErrorResponse(w, "invalid rate limit number")
+		return
+	}
+	if proxyRateLimit <= 0 {
+		utils.SendErrorResponse(w, "rate limit number must be greater than 0")
+		return
+	}
+
 	// Bypass WebSocket Origin Check
 	strbpwsorg, _ := utils.PostPara(r, "bpwsorg")
 	if strbpwsorg == "" {
@@ -309,6 +333,9 @@ func ReverseProxyHandleAddEndpoint(w http.ResponseWriter, r *http.Request) {
 			BasicAuthExceptionRules: []*dynamicproxy.BasicAuthExceptionRule{},
 			DefaultSiteOption:       0,
 			DefaultSiteValue:        "",
+			// Rate Limit
+			RequireRateLimit: requireRateLimit,
+			RateLimit:        proxyRateLimit,
 		}
 
 		preparedEndpoint, err := dynamicProxyRouter.PrepareProxyRoute(&thisProxyEndpoint)
@@ -430,6 +457,26 @@ func ReverseProxyHandleEditEndpoint(w http.ResponseWriter, r *http.Request) {
 
 	requireBasicAuth := (rba == "true")
 
+	// Rate Limiting?
+	rl, _ := utils.PostPara(r, "rate")
+	if rl == "" {
+		rl = "false"
+	}
+	requireRateLimit := (rl == "true")
+	rlnum, _ := utils.PostPara(r, "ratenum")
+	if rlnum == "" {
+		rlnum = "0"
+	}
+	proxyRateLimit, err := strconv.ParseInt(rlnum, 10, 64)
+	if err != nil {
+		utils.SendErrorResponse(w, "invalid rate limit number")
+		return
+	}
+	if proxyRateLimit <= 0 {
+		utils.SendErrorResponse(w, "rate limit number must be greater than 0")
+		return
+	}
+
 	// Bypass WebSocket Origin Check
 	strbpwsorg, _ := utils.PostPara(r, "bpwsorg")
 	if strbpwsorg == "" {
@@ -451,6 +498,8 @@ func ReverseProxyHandleEditEndpoint(w http.ResponseWriter, r *http.Request) {
 	newProxyEndpoint.BypassGlobalTLS = bypassGlobalTLS
 	newProxyEndpoint.SkipCertValidations = skipTlsValidation
 	newProxyEndpoint.RequireBasicAuth = requireBasicAuth
+	newProxyEndpoint.RequireRateLimit = requireRateLimit
+	newProxyEndpoint.RateLimit = proxyRateLimit
 	newProxyEndpoint.SkipWebSocketOriginCheck = bypassWebsocketOriginCheck
 
 	//Prepare to replace the current routing rule

src/web/components/httprp.html

@@ -20,6 +20,7 @@
                     <th>Destination</th>
                     <th>Virtual Directory</th>
                     <th>Basic Auth</th>
+                    <th>Rate Limit</th>
                     <th class="no-sort" style="min-width:150px;">Actions</th>
                 </tr>
             </thead>
@@ -107,6 +108,9 @@
                         <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>
                         <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">
                                 <input type="checkbox" class="enableToggle" name="active" ${enableChecked} eptuuid="${subd.RootOrMatchingDomain}" onchange="handleProxyRuleToggle(this);">
@@ -301,6 +305,23 @@
                     <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}>
+                        <label>Require Rate Limit</label>
+                    </div>
+                    <div class="ui mini fluid input">
+                        <input type="number" class="RateLimit" value="${rateLimit}" placeholder="100" min="1" max="1000" >
+                    </div>
+                `);
+
             }else if (datatype == 'action'){
                 column.empty().append(`
                 <button title="Save" onclick="saveProxyInlineEdit('${uuid.hexEncode()}');" class="ui basic small icon circular button inlineEditActionBtn"><i class="ui green save icon"></i></button>
@@ -348,6 +369,8 @@
         let requireTLS = $(row).find(".RequireTLS")[0].checked;
         let skipCertValidations = $(row).find(".SkipCertValidations")[0].checked;
         let requireBasicAuth = $(row).find(".RequireBasicAuth")[0].checked;
+        let requireRateLimit = $(row).find(".RequireRateLimit")[0].checked;
+        let rateLimit = $(row).find(".RateLimit").val();
         let bypassGlobalTLS = $(row).find(".BypassGlobalTLS")[0].checked;
         let bypassWebsocketOrigin = $(row).find(".SkipWebSocketOriginCheck")[0].checked;
         console.log(newDomain, requireTLS, skipCertValidations, requireBasicAuth)
@@ -364,6 +387,8 @@
                 "tlsval": skipCertValidations,
                 "bpwsorg" : bypassWebsocketOrigin,
                 "bauth" :requireBasicAuth,
+                "rate" :requireRateLimit,
+                "ratenum" :rateLimit,
             },
             success: function(data){
                 if (data.error !== undefined){

src/web/components/rules.html

@@ -73,6 +73,17 @@
                                         <label>Allow plain HTTP access<br><small>Allow this subdomain to be connected without TLS (Require HTTP server enabled on port 80)</small></label>
                                     </div>
                                 </div>
+                                <div class="field">
+                                    <div class="ui checkbox">
+                                        <input type="checkbox" id="requireRateLimit">
+                                        <label>Require Rate Limit<br><small>This proxy endpoint will be rate limited.</small></label>
+                                    </div>
+                                </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>
                                 <div class="field">
                                     <div class="ui checkbox">
                                         <input type="checkbox" id="requireBasicAuth">
@@ -147,6 +158,8 @@
         var skipTLSValidation = $("#skipTLSValidation")[0].checked;
         var bypassGlobalTLS = $("#bypassGlobalTLS")[0].checked;
         var requireBasicAuth = $("#requireBasicAuth")[0].checked;
+        var proxyRateLimit = $("#proxyRateLimit").val();
+        var requireRateLimit = $("#requireRateLimit")[0].checked;
         var skipWebSocketOriginCheck = $("#skipWebsocketOriginCheck")[0].checked;
         var accessRuleToUse = $("#newProxyRuleAccessFilter").val();
         
@@ -176,6 +189,8 @@
                 bpwsorg: skipWebSocketOriginCheck,
                 bypassGlobalTLS: bypassGlobalTLS,
                 bauth: requireBasicAuth,
+                rate: requireRateLimit,
+                ratenum: proxyRateLimit,
                 cred: JSON.stringify(credentials),
                 access: accessRuleToUse,
             },
@@ -264,6 +279,16 @@
     }
     $("#requireBasicAuth").on('change', toggleBasicAuth);
     toggleBasicAuth();
+    
+    function toggleRateLimit() {
+        if ($("#requireRateLimit").parent().checkbox("is checked")) {
+            $("#proxyRateLimit").parent().removeClass("disabled");
+        } else {
+            $("#proxyRateLimit").parent().addClass("disabled");
+        }
+    }
+    $("#requireRateLimit").on('change', toggleRateLimit);
+    toggleRateLimit();
 
 
     /*