feat Main HTTP Server; IP,URL Allowlist and Blocklist

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