serverIp/internal/middleware/ip_whitelist.go

package middleware

import (
	"github.com/gin-gonic/gin"
	"github.com/spf13/viper"
	"go.uber.org/zap"
	"net"
	"net/http"
	"projectName/pkg/log"
	"strings"
	"sync"
)

// IPAllowlist 保存允许访问的IP列表
type IPAllowlist struct {
	allowedIPs map[string]bool
	enabled    bool
	mu         sync.RWMutex
	logger     *log.Logger
}

func getClientIP(r *http.Request) string {
	ip := r.Header.Get("X-Forwarded-For")
	if ip != "" {
		parts := strings.Split(ip, ",")
		return strings.TrimSpace(parts[0])
	}
	ip, _, _ = net.SplitHostPort(r.RemoteAddr)
	return ip
}

// NewIPAllowlist 创建一个新的IP白名单实例
func NewIPAllowlist(conf *viper.Viper, logger *log.Logger) *IPAllowlist {
	allowedIPs := conf.GetStringSlice("ip_allowlist.ips")
	enabled := conf.GetBool("ip_allowlist.enabled")

	allowlist := &IPAllowlist{
		allowedIPs: make(map[string]bool),
		enabled:    enabled,
		logger:     logger,
	}

	// 将配置文件中的IP添加到白名单
	for _, ip := range allowedIPs {
		allowlist.allowedIPs[ip] = true
	}

	return allowlist
}

// IPAllowlistMiddleware 创建一个IP白名单中间件
func (a *IPAllowlist) IPAllowlistMiddleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		// 如果白名单未启用，直接放行
		if !a.enabled {
			c.Next()
			return
		}

		clientIP := getClientIP(c.Request)
		a.mu.RLock()
		_, allowed := a.allowedIPs[clientIP]
		a.mu.RUnlock()

		if allowed {
			c.Next()
		} else {
			a.logger.WithContext(c).Warn("拒绝未授权的IP访问: %s", zap.String("ip", clientIP))
			c.AbortWithStatusJSON(http.StatusForbidden, gin.H{
				"code": 403,
				"msg":  "IP访问受限，您的IP没有权限访问此资源",
			})
		}
	}
}

// AddIP 添加IP到白名单
func (a *IPAllowlist) AddIP(ip string) {
	a.mu.Lock()
	defer a.mu.Unlock()
	a.allowedIPs[ip] = true
}

// RemoveIP 从白名单中移除IP
func (a *IPAllowlist) RemoveIP(ip string) {
	a.mu.Lock()
	defer a.mu.Unlock()
	delete(a.allowedIPs, ip)
}

// IsIPAllowed 检查IP是否在白名单中
func (a *IPAllowlist) IsIPAllowed(ip string) bool {
	a.mu.RLock()
	defer a.mu.RUnlock()
	return a.allowedIPs[ip]
}

// EnableAllowlist 启用IP白名单
func (a *IPAllowlist) EnableAllowlist() {
	a.mu.Lock()
	defer a.mu.Unlock()
	a.enabled = true
}

// DisableAllowlist 禁用IP白名单
func (a *IPAllowlist) DisableAllowlist() {
	a.mu.Lock()
	defer a.mu.Unlock()
	a.enabled = false
}

// IsEnabled 检查白名单是否启用
func (a *IPAllowlist) IsEnabled() bool {
	a.mu.RLock()
	defer a.mu.RUnlock()
	return a.enabled
}