wormhole-server/pkg/socks5/rules.go

package socks5

import (
	"net"
	"strconv"
	"strings"
)

// SimpleRule 简单访问规则
type SimpleRule struct {
	action   string // allow, deny
	networks []*net.IPNet
	ports    []int
}

// NewRule 创建访问规则
func NewRule(config RuleConfig) Rule {
	rule := &SimpleRule{
		action: strings.ToLower(config.Action),
		ports:  config.Ports,
	}

	// 解析IP网段
	for _, ipStr := range config.IPs {
		if !strings.Contains(ipStr, "/") {
			// 单个IP，添加适当的子网掩码
			if strings.Contains(ipStr, ":") {
				ipStr += "/128" // IPv6
			} else {
				ipStr += "/32" // IPv4
			}
		}

		_, network, err := net.ParseCIDR(ipStr)
		if err != nil {
			// 如果解析失败，尝试作为单个IP处理
			ip := net.ParseIP(ipStr)
			if ip != nil {
				if ip.To4() != nil {
					_, network, _ = net.ParseCIDR(ip.String() + "/32")
				} else {
					_, network, _ = net.ParseCIDR(ip.String() + "/128")
				}
			}
		}

		if network != nil {
			rule.networks = append(rule.networks, network)
		}
	}

	return rule
}

// Allow 检查是否允许访问
func (r *SimpleRule) Allow(addr net.Addr) bool {
	// 解析地址
	var ip net.IP
	var port int

	switch a := addr.(type) {
	case *net.IPAddr:
		ip = a.IP
	case *net.TCPAddr:
		ip = a.IP
		port = a.Port
	case *net.UDPAddr:
		ip = a.IP
		port = a.Port
	default:
		// 尝试从字符串解析
		host, portStr, err := net.SplitHostPort(addr.String())
		if err != nil {
			return r.action == "allow" // 默认策略
		}
		ip = net.ParseIP(host)
		if p, err := strconv.Atoi(portStr); err == nil {
			port = p
		}
	}

	if ip == nil {
		return r.action == "allow" // 默认策略
	}

	// 检查IP是否匹配
	ipMatches := len(r.networks) == 0 // 如果没有指定网段，默认匹配
	for _, network := range r.networks {
		if network.Contains(ip) {
			ipMatches = true
			break
		}
	}

	// 检查端口是否匹配
	portMatches := len(r.ports) == 0 // 如果没有指定端口，默认匹配
	for _, p := range r.ports {
		if p == port {
			portMatches = true
			break
		}
	}

	// 根据规则类型返回结果
	matches := ipMatches && portMatches
	switch r.action {
	case "allow":
		return matches
	case "deny":
		return !matches
	default:
		return true // 默认允许
	}
}