fix: Obtain User request URI

config/torii.yml

@@ -2,3 +2,5 @@ port: "25555"
 rule_path: "/www/dev/server_torii/config/rules"
 connecting_ip_headers:
   - "X-Real-IP"
+connecting_uri_headers:
+  - "X-Original-URI"

internal/config/config.go

@@ -15,6 +15,7 @@ type MainConfig struct {
 	Port                 string   `yaml:"port"`
 	RulePath             string   `yaml:"rule_path"`
 	ConnectingIPHeaders  []string `yaml:"connecting_ip_headers"`
+	ConnectingURIHeaders []string `yaml:"connecting_uri_headers"`
 }
 
 // LoadMainConfig Read the configuration file and return the configuration object

internal/server/server.go

@@ -10,11 +10,51 @@ import (
 	"strings"
 )
 
+type userRequest struct {
+	remoteIP string
+	uri      string
+}
+
 // StartServer starts the HTTP server
-func StartServer(port string, ruleSet *config.RuleSet, ipHeaders []string) error {
+func StartServer(cfg *config.MainConfig, ruleSet *config.RuleSet) error {
 	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
+
+		userRequestData := processRequestData(cfg, r)
+
+		decision := action.NewDecision()
+
+		// run main check logic
+		checkIPAllow(userRequestData.remoteIP, ruleSet.IPAllowTrie, decision)
+		checkIPBlock(userRequestData.remoteIP, ruleSet.IPBlockTrie, decision)
+		checkURLAllow(userRequestData.uri, ruleSet.URLAllowList, decision)
+		checkURLBlock(userRequestData.uri, ruleSet.URLBlockList, decision)
+
+		// if still undecided, allow
+		if decision.Get() == action.Undecided {
+			decision.Set(action.Allow)
+		}
+		log.Printf("clientIP: %s, decision: %s, Headers: %v", userRequestData.remoteIP, decision.Get(), r.Header)
+		// return response
+		if decision.Get() == action.Allow {
+			w.WriteHeader(http.StatusOK)
+			w.Write([]byte("Allowed"))
+		} else if decision.Get() == action.Block {
+			w.WriteHeader(http.StatusForbidden)
+			w.Write([]byte("Blocked"))
+		} else {
+			// should not reach here
+			w.WriteHeader(http.StatusInternalServerError)
+		}
+	})
+
+	log.Printf("HTTP Server listening on :%s ...", cfg.Port)
+	return http.ListenAndServe(":"+cfg.Port, nil)
+}
+
+func processRequestData(cfg *config.MainConfig, r *http.Request) userRequest {
+
 	var clientIP string
-		for _, headerName := range ipHeaders {
+	for _, headerName := range cfg.ConnectingIPHeaders {
 		if ipVal := r.Header.Get(headerName); ipVal != "" {
 			if strings.Contains(clientIP, ",") {
 				parts := strings.Split(ipVal, ",")
@@ -36,34 +76,22 @@ func StartServer(port string, ruleSet *config.RuleSet, ipHeaders []string) error
 		}
 	}
 
-		decision := action.NewDecision()
-
-		// run main check logic
-		checkIPAllow(clientIP, ruleSet.IPAllowTrie, decision)
-		checkIPBlock(clientIP, ruleSet.IPBlockTrie, decision)
-		checkURLAllow(r.RequestURI, ruleSet.URLAllowList, decision)
-		checkURLBlock(r.RequestURI, ruleSet.URLBlockList, decision)
-
-		// if still undecided, allow
-		if decision.Get() == action.Undecided {
-			decision.Set(action.Allow)
+	var clientURI string
+	for _, headerName := range cfg.ConnectingURIHeaders {
+		if uriVal := r.Header.Get(headerName); uriVal != "" {
+			clientURI = uriVal
+			break
+		}
+	}
+	if clientURI == "" {
+		clientURI = r.RequestURI
 	}
 
-		// return response
-		if decision.Get() == action.Allow {
-			w.WriteHeader(http.StatusOK)
-			w.Write([]byte("Allowed"))
-		} else if decision.Get() == action.Block {
-			w.WriteHeader(http.StatusForbidden)
-			w.Write([]byte("Blocked"))
-		} else {
-			// should not reach here
-			w.WriteHeader(http.StatusInternalServerError)
+	userRequest := userRequest{
+		remoteIP: clientIP,
+		uri:      clientURI,
 	}
-	})
-
-	log.Printf("HTTP Server listening on :%s ...", port)
-	return http.ListenAndServe(":"+port, nil)
+	return userRequest
 }
 
 func checkIPAllow(remoteIP string, trie *dataType.TrieNode, decision *action.Decision) {

main.go

@@ -35,7 +35,7 @@ func main() {
 
 	serverErr := make(chan error, 1)
 	go func() {
-		serverErr <- server.StartServer(cfg.Port, ruleSet, cfg.ConnectingIPHeaders)
+		serverErr <- server.StartServer(cfg, ruleSet)
 	}()
 
 	select {