125 lines
2.3 KiB
Go
125 lines
2.3 KiB
Go
package middleware
|
|
|
|
import (
|
|
"net/http"
|
|
"strings"
|
|
"sync"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
// IPBlacklist IP黑名单管理
|
|
var IPBlacklist = &ipBlacklist{
|
|
blockedIPs: make(map[string]bool),
|
|
allowedIPs: make(map[string]bool),
|
|
blockedCIDR: []string{},
|
|
}
|
|
|
|
type ipBlacklist struct {
|
|
mu sync.RWMutex
|
|
blockedIPs map[string]bool
|
|
allowedIPs map[string]bool
|
|
blockedCIDR []string // 预留CIDR支持
|
|
enabled bool
|
|
}
|
|
|
|
// SetEnabled 启用/禁用IP黑名单
|
|
func (bl *ipBlacklist) SetEnabled(enabled bool) {
|
|
bl.mu.Lock()
|
|
defer bl.mu.Unlock()
|
|
bl.enabled = enabled
|
|
}
|
|
|
|
// IsEnabled 是否启用
|
|
func (bl *ipBlacklist) IsEnabled() bool {
|
|
bl.mu.RLock()
|
|
defer bl.mu.RUnlock()
|
|
return bl.enabled
|
|
}
|
|
|
|
// AddBlockedIP 添加封禁IP
|
|
func (bl *ipBlacklist) AddBlockedIP(ip string) {
|
|
bl.mu.Lock()
|
|
defer bl.mu.Unlock()
|
|
bl.blockedIPs[ip] = true
|
|
}
|
|
|
|
// RemoveBlockedIP 移除封禁IP
|
|
func (bl *ipBlacklist) RemoveBlockedIP(ip string) {
|
|
bl.mu.Lock()
|
|
defer bl.mu.Unlock()
|
|
delete(bl.blockedIPs, ip)
|
|
}
|
|
|
|
// IsBlocked 检查IP是否被封禁
|
|
func (bl *ipBlacklist) IsBlocked(ip string) bool {
|
|
bl.mu.RLock()
|
|
defer bl.mu.RUnlock()
|
|
|
|
if !bl.enabled {
|
|
return false
|
|
}
|
|
|
|
// 白名单优先
|
|
if bl.allowedIPs[ip] {
|
|
return false
|
|
}
|
|
|
|
return bl.blockedIPs[ip]
|
|
}
|
|
|
|
// GetBlockedIPs 获取所有封禁IP
|
|
func (bl *ipBlacklist) GetBlockedIPs() []string {
|
|
bl.mu.RLock()
|
|
defer bl.mu.RUnlock()
|
|
|
|
ips := make([]string, 0, len(bl.blockedIPs))
|
|
for ip := range bl.blockedIPs {
|
|
ips = append(ips, ip)
|
|
}
|
|
return ips
|
|
}
|
|
|
|
// SetBlockedIPs 批量设置封禁IP
|
|
func (bl *ipBlacklist) SetBlockedIPs(ips []string) {
|
|
bl.mu.Lock()
|
|
defer bl.mu.Unlock()
|
|
|
|
bl.blockedIPs = make(map[string]bool)
|
|
for _, ip := range ips {
|
|
ip = strings.TrimSpace(ip)
|
|
if ip != "" {
|
|
bl.blockedIPs[ip] = true
|
|
}
|
|
}
|
|
}
|
|
|
|
// AddAllowedIP 添加白名单IP
|
|
func (bl *ipBlacklist) AddAllowedIP(ip string) {
|
|
bl.mu.Lock()
|
|
defer bl.mu.Unlock()
|
|
bl.allowedIPs[ip] = true
|
|
}
|
|
|
|
// IPBlacklistMiddleware IP黑名单中间件
|
|
func IPBlacklistMiddleware() gin.HandlerFunc {
|
|
return func(c *gin.Context) {
|
|
if !IPBlacklist.IsEnabled() {
|
|
c.Next()
|
|
return
|
|
}
|
|
|
|
clientIP := c.ClientIP()
|
|
|
|
if IPBlacklist.IsBlocked(clientIP) {
|
|
c.AbortWithStatusJSON(http.StatusForbidden, gin.H{
|
|
"code": 403,
|
|
"message": "您的IP已被封禁",
|
|
})
|
|
return
|
|
}
|
|
|
|
c.Next()
|
|
}
|
|
}
|