From dcbfa277222047ef83e9bddf521f92208713c960 Mon Sep 17 00:00:00 2001 From: Roi Feng <37480123+Rayzggz@users.noreply.github.com> Date: Thu, 6 Feb 2025 01:16:56 -0500 Subject: [PATCH] feat Main HTTP Server; IP,URL Allowlist and Blocklist --- config/rules/IP_AllowList.conf | 0 config/rules/IP_BlockList.conf | 0 config/rules/URL_AllowList.conf | 0 config/rules/URL_BlockList.conf | 0 config/torii.yml | 4 + go.mod | 5 + internal/action/action.go | 26 +++++ internal/config/config.go | 162 ++++++++++++++++++++++++++++++++ internal/dataType/ip_trie.go | 46 +++++++++ internal/dataType/url_list.go | 47 +++++++++ internal/server/server.go | 112 ++++++++++++++++++++++ main.go | 51 ++++++++++ 12 files changed, 453 insertions(+) create mode 100644 config/rules/IP_AllowList.conf create mode 100644 config/rules/IP_BlockList.conf create mode 100644 config/rules/URL_AllowList.conf create mode 100644 config/rules/URL_BlockList.conf create mode 100644 config/torii.yml create mode 100644 go.mod create mode 100644 internal/action/action.go create mode 100644 internal/config/config.go create mode 100644 internal/dataType/ip_trie.go create mode 100644 internal/dataType/url_list.go create mode 100644 internal/server/server.go create mode 100644 main.go diff --git a/config/rules/IP_AllowList.conf b/config/rules/IP_AllowList.conf new file mode 100644 index 0000000..e69de29 diff --git a/config/rules/IP_BlockList.conf b/config/rules/IP_BlockList.conf new file mode 100644 index 0000000..e69de29 diff --git a/config/rules/URL_AllowList.conf b/config/rules/URL_AllowList.conf new file mode 100644 index 0000000..e69de29 diff --git a/config/rules/URL_BlockList.conf b/config/rules/URL_BlockList.conf new file mode 100644 index 0000000..e69de29 diff --git a/config/torii.yml b/config/torii.yml new file mode 100644 index 0000000..f95c328 --- /dev/null +++ b/config/torii.yml @@ -0,0 +1,4 @@ +port: "2550" +rule_path: "/www/dev/server_torii/config/rules" +connecting_ip_headers: + - "X-Real-IP" \ No newline at end of file diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..28f63ae --- /dev/null +++ b/go.mod @@ -0,0 +1,5 @@ +module server_torii + +go 1.23.5 + +require gopkg.in/yaml.v3 v3.0.1 diff --git a/internal/action/action.go b/internal/action/action.go new file mode 100644 index 0000000..e960e1b --- /dev/null +++ b/internal/action/action.go @@ -0,0 +1,26 @@ +package action + +type Action int + +const ( + Undecided Action = iota // 0：Undecided + Allow // 1：Pass + Block // 2：Deny +) + +// Decision saves the result of the decision +type Decision struct { + result Action +} + +func NewDecision() *Decision { + return &Decision{result: Undecided} +} + +func (d *Decision) Get() Action { + return d.result +} + +func (d *Decision) Set(new Action) { + d.result = new +} diff --git a/internal/config/config.go b/internal/config/config.go new file mode 100644 index 0000000..fb5fe4f --- /dev/null +++ b/internal/config/config.go @@ -0,0 +1,162 @@ +package config + +import ( + "bufio" + "gopkg.in/yaml.v3" + "net" + "os" + "path/filepath" + "regexp" + "server_torii/internal/dataType" + "strings" +) + +type MainConfig struct { + Port string `yaml:"port"` + RulePath string `yaml:"rule_path"` + ConnectingIPHeaders []string `yaml:"connecting_ip_headers"` +} + +// LoadMainConfig Read the configuration file and return the configuration object +func LoadMainConfig(basePath string) (*MainConfig, error) { + exePath, err := os.Executable() + if err != nil { + return nil, err + } + if basePath == "" { + basePath = filepath.Dir(exePath) + } + configPath := filepath.Join(basePath, "config", "torii.yml") + + data, err := os.ReadFile(configPath) + if err != nil { + return nil, err + } + + var cfg MainConfig + if err := yaml.Unmarshal(data, &cfg); err != nil { + return nil, err + } + + return &cfg, nil +} + +// RuleSet stores all rules +type RuleSet struct { + IPAllowTrie *dataType.TrieNode + IPBlockTrie *dataType.TrieNode + URLAllowList *dataType.URLRuleList + URLBlockList *dataType.URLRuleList +} + +// LoadRules Load all rules from the specified path +func LoadRules(rulePath string) (*RuleSet, error) { + rs := RuleSet{ + IPAllowTrie: &dataType.TrieNode{}, + IPBlockTrie: &dataType.TrieNode{}, + URLAllowList: &dataType.URLRuleList{}, + URLBlockList: &dataType.URLRuleList{}, + } + + // Load IP Allow List + ipAllowFile := rulePath + "/IP_AllowList.conf" + if err := loadIPRules(ipAllowFile, rs.IPAllowTrie); err != nil { + return nil, err + } + + // Load IP Block List + ipBlockFile := rulePath + "/IP_BlockList.conf" + if err := loadIPRules(ipBlockFile, rs.IPBlockTrie); err != nil { + return nil, err + } + + // Load URL Allow List + urlAllowFile := rulePath + "/URL_AllowList.conf" + if err := loadURLRules(urlAllowFile, rs.URLAllowList); err != nil { + return nil, err + } + + // Load URL Block List + urlBlockFile := rulePath + "/URL_BlockList.conf" + if err := loadURLRules(urlBlockFile, rs.URLBlockList); err != nil { + return nil, err + } + + return &rs, nil +} + +// loadIPRules read the IP rule file and insert the rules into the trie +func loadIPRules(filePath string, trie *dataType.TrieNode) error { + file, err := os.Open(filePath) + if err != nil { + return err + } + defer func(file *os.File) { + var err = file.Close() + if err != nil { + + } + }(file) + + scanner := bufio.NewScanner(file) + for scanner.Scan() { + line := strings.TrimSpace(scanner.Text()) + if line == "" { + continue + } + + if !strings.Contains(line, "/") { + line = line + "/32" + } + _, ipNet, err := net.ParseCIDR(line) + if err != nil { + continue + } + trie.Insert(ipNet) + } + + return scanner.Err() +} + +// loadURLRules Load URL rules from the specified file +func loadURLRules(filePath string, list *dataType.URLRuleList) error { + file, err := os.Open(filePath) + if err != nil { + return err + } + defer func(file *os.File) { + var err = file.Close() + if err != nil { + + } + }(file) + + scanner := bufio.NewScanner(file) + for scanner.Scan() { + line := strings.TrimSpace(scanner.Text()) + if line == "" { + continue + } + // check if the line is a regex + isRegex := false + if strings.HasPrefix(line, "^") || strings.HasSuffix(line, "$") || strings.ContainsAny(line, ".*+?()[]{}|\\") { + isRegex = true + } + var compiled *regexp.Regexp + if isRegex { + compiled, err = regexp.Compile(line) + if err != nil { + // skip invalid regex + continue + } + } + rule := &dataType.URLRule{ + Pattern: line, + IsRegex: isRegex, + Regex: compiled, + } + list.Append(rule) + } + + return scanner.Err() +} diff --git a/internal/dataType/ip_trie.go b/internal/dataType/ip_trie.go new file mode 100644 index 0000000..161a913 --- /dev/null +++ b/internal/dataType/ip_trie.go @@ -0,0 +1,46 @@ +package dataType + +import "net" + +type TrieNode struct { + children [2]*TrieNode + isEnd bool +} + +// Insert IP or CIDR rule into trie, prefixLength represents the prefix length +func (node *TrieNode) Insert(ipNet *net.IPNet) { + ones, _ := ipNet.Mask.Size() + ip := ipNet.IP.To4() + if ip == nil { + return + } + current := node + for i := 0; i < ones; i++ { + bit := (ip[i/8] >> (7 - uint(i%8))) & 1 + if current.children[bit] == nil { + current.children[bit] = &TrieNode{} + } + current = current.children[bit] + } + current.isEnd = true +} + +// Search if the ip is in the trie +func (node *TrieNode) Search(ip net.IP) bool { + ip = ip.To4() + if ip == nil { + return false + } + current := node + for i := 0; i < 32; i++ { + if current.isEnd { + return true + } + bit := (ip[i/8] >> (7 - uint(i%8))) & 1 + if current.children[bit] == nil { + return false + } + current = current.children[bit] + } + return current.isEnd +} diff --git a/internal/dataType/url_list.go b/internal/dataType/url_list.go new file mode 100644 index 0000000..a26038a --- /dev/null +++ b/internal/dataType/url_list.go @@ -0,0 +1,47 @@ +package dataType + +import "regexp" + +// URLRule struct for a URL rule +type URLRule struct { + Pattern string + IsRegex bool + Regex *regexp.Regexp + Next *URLRule +} + +// URLRuleList struct LinkedList +type URLRuleList struct { + Head *URLRule +} + +// Append add a rule to the end of the list +func (l *URLRuleList) Append(rule *URLRule) { + if l.Head == nil { + l.Head = rule + return + } + current := l.Head + for current.Next != nil { + current = current.Next + } + current.Next = rule +} + +// Match check if the URL matches any rule in the list +func (l *URLRuleList) Match(url string) bool { + current := l.Head + for current != nil { + if current.IsRegex { + if current.Regex.MatchString(url) { + return true + } + } else { + if current.Pattern == url { + return true + } + } + current = current.Next + } + return false +} diff --git a/internal/server/server.go b/internal/server/server.go new file mode 100644 index 0000000..1268aad --- /dev/null +++ b/internal/server/server.go @@ -0,0 +1,112 @@ +// internal/server/server.go +package server + +import ( + "log" + "net" + "net/http" + "server_torii/internal/action" + "server_torii/internal/config" + "server_torii/internal/dataType" + "strings" +) + +// StartServer starts the HTTP server +func StartServer(port string, ruleSet *config.RuleSet, ipHeaders []string) error { + http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { + var clientIP string + for _, headerName := range ipHeaders { + if ipVal := r.Header.Get(headerName); ipVal != "" { + if strings.Contains(clientIP, ",") { + parts := strings.Split(ipVal, ",") + clientIP = strings.TrimSpace(parts[0]) + } + clientIP = ipVal + break + } + } + + if clientIP == "" { + remoteAddr := r.RemoteAddr + ipStr, _, err := net.SplitHostPort(remoteAddr) + if err != nil { + //TODO: log error + clientIP = remoteAddr + } else { + clientIP = ipStr + } + } + + 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) + } + + // 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 ...", port) + return http.ListenAndServe(":"+port, nil) +} + +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) + } +} diff --git a/main.go b/main.go new file mode 100644 index 0000000..31a5f22 --- /dev/null +++ b/main.go @@ -0,0 +1,51 @@ +package main + +import ( + "flag" + "log" + "os" + "os/signal" + "server_torii/internal/config" + "server_torii/internal/server" + "syscall" +) + +func main() { + var basePath string + flag.StringVar(&basePath, "prefix", "", "Config file base path") + flag.Parse() + + // Load MainConfig + cfg, err := config.LoadMainConfig(basePath) + if err != nil { + log.Fatalf("Load config failed: %v", err) + } + + // Load rules + ruleSet, err := config.LoadRules(cfg.RulePath) + if err != nil { + log.Fatalf("Load rules failed: %v", err) + } + + log.Printf("Ready to start server on port %s", cfg.Port) + + // Start server + stop := make(chan os.Signal, 1) + signal.Notify(stop, syscall.SIGINT, syscall.SIGTERM) + + serverErr := make(chan error, 1) + go func() { + serverErr <- server.StartServer(cfg.Port, ruleSet, cfg.ConnectingIPHeaders) + }() + + select { + case <-stop: + log.Println("Stopping server...") + case err := <-serverErr: + if err != nil { + log.Fatalf("Failed to start server: %v", err) + } + } + + log.Println("Server stopped") +}