feat Main HTTP Server; IP,URL Allowlist and Blocklist

config/torii.yml

@@ -0,0 +1,4 @@
+port: "2550"
+rule_path: "/www/dev/server_torii/config/rules"
+connecting_ip_headers:
+  - "X-Real-IP"

go.mod

@@ -0,0 +1,5 @@
+module server_torii
+
+go 1.23.5
+
+require gopkg.in/yaml.v3 v3.0.1

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
+}

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()
+}

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
+}

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
+}

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)
+	}
+}

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")
+}