From aea36d463bda866f9823d163f1dd5baa77fb9ba6 Mon Sep 17 00:00:00 2001 From: Roi Feng <37480123+Rayzggz@users.noreply.github.com> Date: Thu, 13 Feb 2025 22:20:53 -0500 Subject: [PATCH] feat: new action process --- internal/action/action.go | 29 ++++++++----- internal/check/IPAllow.go | 23 +++++++++++ internal/check/IPBlock.go | 22 ++++++++++ internal/check/URLAllow.go | 17 ++++++++ internal/check/URLBlock.go | 17 ++++++++ internal/dataType/type.go | 6 +++ internal/server/server.go | 83 +++++++++----------------------------- 7 files changed, 123 insertions(+), 74 deletions(-) create mode 100644 internal/check/IPAllow.go create mode 100644 internal/check/IPBlock.go create mode 100644 internal/check/URLAllow.go create mode 100644 internal/check/URLBlock.go create mode 100644 internal/dataType/type.go diff --git a/internal/action/action.go b/internal/action/action.go index e960e1b..58c7fad 100644 --- a/internal/action/action.go +++ b/internal/action/action.go @@ -1,26 +1,35 @@ package action -type Action int +type checkState int const ( - Undecided Action = iota // 0：Undecided - Allow // 1：Pass - Block // 2：Deny + Continue checkState = iota + Done + Jump ) // Decision saves the result of the decision type Decision struct { - result Action + HTTPCode string + State checkState + JumpIndex int } func NewDecision() *Decision { - return &Decision{result: Undecided} + return &Decision{HTTPCode: "200", State: Continue, JumpIndex: -1} } -func (d *Decision) Get() Action { - return d.result +func (d *Decision) Set(state checkState) { + d.State = state } -func (d *Decision) Set(new Action) { - d.result = new +func (d *Decision) SetCode(state checkState, httpCode string) { + d.State = state + d.HTTPCode = httpCode +} + +func (d *Decision) SetJump(state checkState, httpCode string, jumpIndex int) { + d.State = state + d.HTTPCode = httpCode + d.JumpIndex = jumpIndex } diff --git a/internal/check/IPAllow.go b/internal/check/IPAllow.go new file mode 100644 index 0000000..262c915 --- /dev/null +++ b/internal/check/IPAllow.go @@ -0,0 +1,23 @@ +package check + +import ( + "net" + "server_torii/internal/action" + "server_torii/internal/config" + "server_torii/internal/dataType" +) + +func IPAllowList(reqData dataType.UserRequest, ruleSet *config.RuleSet, decision *action.Decision) { + remoteIP := reqData.RemoteIP + trie := ruleSet.IPAllowTrie + + ip := net.ParseIP(remoteIP) + if ip == nil { + return + } + if trie.Search(ip) { + decision.SetCode(action.Done, "200") + } else { + decision.Set(action.Continue) + } +} diff --git a/internal/check/IPBlock.go b/internal/check/IPBlock.go new file mode 100644 index 0000000..12911b3 --- /dev/null +++ b/internal/check/IPBlock.go @@ -0,0 +1,22 @@ +package check + +import ( + "net" + "server_torii/internal/action" + "server_torii/internal/config" + "server_torii/internal/dataType" +) + +func IPBlockList(reqData dataType.UserRequest, ruleSet *config.RuleSet, decision *action.Decision) { + remoteIP := reqData.RemoteIP + trie := ruleSet.IPBlockTrie + ip := net.ParseIP(remoteIP) + if ip == nil { + return + } + if trie.Search(ip) { + decision.SetCode(action.Done, "403") + } else { + decision.Set(action.Continue) + } +} diff --git a/internal/check/URLAllow.go b/internal/check/URLAllow.go new file mode 100644 index 0000000..6be1f5d --- /dev/null +++ b/internal/check/URLAllow.go @@ -0,0 +1,17 @@ +package check + +import ( + "server_torii/internal/action" + "server_torii/internal/config" + "server_torii/internal/dataType" +) + +func URLAllowList(reqData dataType.UserRequest, ruleSet *config.RuleSet, decision *action.Decision) { + url := reqData.Uri + list := ruleSet.URLAllowList + if list.Match(url) { + decision.SetCode(action.Done, "200") + } else { + decision.Set(action.Continue) + } +} diff --git a/internal/check/URLBlock.go b/internal/check/URLBlock.go new file mode 100644 index 0000000..2f4b9ce --- /dev/null +++ b/internal/check/URLBlock.go @@ -0,0 +1,17 @@ +package check + +import ( + "server_torii/internal/action" + "server_torii/internal/config" + "server_torii/internal/dataType" +) + +func URLBlockList(reqData dataType.UserRequest, ruleSet *config.RuleSet, decision *action.Decision) { + url := reqData.Uri + list := ruleSet.URLBlockList + if list.Match(url) { + decision.SetCode(action.Done, "403") + } else { + decision.Set(action.Continue) + } +} diff --git a/internal/dataType/type.go b/internal/dataType/type.go new file mode 100644 index 0000000..af30942 --- /dev/null +++ b/internal/dataType/type.go @@ -0,0 +1,6 @@ +package dataType + +type UserRequest struct { + RemoteIP string + Uri string +} diff --git a/internal/server/server.go b/internal/server/server.go index be8a206..cfbbfac 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -5,15 +5,13 @@ import ( "net" "net/http" "server_torii/internal/action" + "server_torii/internal/check" "server_torii/internal/config" "server_torii/internal/dataType" "strings" ) -type userRequest struct { - remoteIP string - uri string -} +type CheckFunc func(dataType.UserRequest, *config.RuleSet, *action.Decision) // StartServer starts the HTTP server func StartServer(cfg *config.MainConfig, ruleSet *config.RuleSet) error { @@ -23,22 +21,23 @@ func StartServer(cfg *config.MainConfig, ruleSet *config.RuleSet) error { 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) + checkFuncs := make([]CheckFunc, 0) + checkFuncs = append(checkFuncs, check.IPAllowList) + checkFuncs = append(checkFuncs, check.IPBlockList) + checkFuncs = append(checkFuncs, check.URLAllowList) + checkFuncs = append(checkFuncs, check.URLBlockList) - // if still undecided, allow - if decision.Get() == action.Undecided { - decision.Set(action.Allow) + for _, checkFunc := range checkFuncs { + checkFunc(userRequestData, ruleSet, decision) + if decision.State == action.Done { + break + } } - log.Printf("clientIP: %s, decision: %s, Headers: %v", userRequestData.remoteIP, decision.Get(), r.Header) - // return response - if decision.Get() == action.Allow { + + if decision.HTTPCode == "200" { w.WriteHeader(http.StatusOK) w.Write([]byte("Allowed")) - } else if decision.Get() == action.Block { + } else if decision.HTTPCode == "403" { w.WriteHeader(http.StatusForbidden) w.Write([]byte("Blocked")) } else { @@ -51,7 +50,7 @@ func StartServer(cfg *config.MainConfig, ruleSet *config.RuleSet) error { return http.ListenAndServe(":"+cfg.Port, nil) } -func processRequestData(cfg *config.MainConfig, r *http.Request) userRequest { +func processRequestData(cfg *config.MainConfig, r *http.Request) dataType.UserRequest { var clientIP string for _, headerName := range cfg.ConnectingIPHeaders { @@ -87,53 +86,9 @@ func processRequestData(cfg *config.MainConfig, r *http.Request) userRequest { clientURI = r.RequestURI } - userRequest := userRequest{ - remoteIP: clientIP, - uri: clientURI, + userRequest := dataType.UserRequest{ + RemoteIP: clientIP, + Uri: clientURI, } return userRequest } - -func checkIPAllow(remoteIP string, trie *dataType.TrieNode, decision *action.Decision) { - if decision.Get() != action.Undecided { - return - } - ip := net.ParseIP(remoteIP) - if ip == nil { - return - } - if trie.Search(ip) { - decision.Set(action.Allow) - } -} - -func checkIPBlock(remoteIP string, trie *dataType.TrieNode, decision *action.Decision) { - if decision.Get() != action.Undecided { - return - } - ip := net.ParseIP(remoteIP) - if ip == nil { - return - } - if trie.Search(ip) { - decision.Set(action.Block) - } -} - -func checkURLAllow(url string, list *dataType.URLRuleList, decision *action.Decision) { - if decision.Get() != action.Undecided { - return - } - if list.Match(url) { - decision.Set(action.Allow) - } -} - -func checkURLBlock(url string, list *dataType.URLRuleList, decision *action.Decision) { - if decision.Get() != action.Undecided { - return - } - if list.Match(url) { - decision.Set(action.Block) - } -}