From d571e197fd04175320f287f0068eb08af21dbc0a Mon Sep 17 00:00:00 2001 From: samwaf Date: Wed, 3 Jun 2026 13:29:34 +0800 Subject: [PATCH] feat:support ipdb #809 --- api/waf_iplocation.go | 447 +++++++++++++++++++++++++++++---------- cmd/samwaf/main.go | 27 +++ go.mod | 1 + go.sum | 2 + iplocation/manager.go | 109 +++++++++- iplocation/types.go | 1 + router/waf_iplocation.go | 2 + 7 files changed, 472 insertions(+), 117 deletions(-) diff --git a/api/waf_iplocation.go b/api/waf_iplocation.go index e681af1..8cd33a8 100644 --- a/api/waf_iplocation.go +++ b/api/waf_iplocation.go @@ -4,9 +4,11 @@ import ( "SamWaf/global" "SamWaf/iplocation" "SamWaf/model/common/response" + "SamWaf/model/request" "SamWaf/utils" "fmt" "io/ioutil" + "net" "os" "path/filepath" @@ -16,6 +18,153 @@ import ( type WafIPLocationApi struct { } +// ipdbStatusResponse 在 DBStatus 基础上追加文件存在信息 +type ipdbStatusResponse struct { + iplocation.DBStatus + FileExists map[string]bool `json:"file_exists"` +} + +// IPDBConfigResp IP 数据库配置项响应 +type IPDBConfigResp struct { + Ipv4Source string `json:"ipv4_source"` + Ipv4Format string `json:"ipv4_format"` + Ipv6Source string `json:"ipv6_source"` + Ipv6Format string `json:"ipv6_format"` +} + +// IPDBConfigReq IP 数据库配置项保存入参 +type IPDBConfigReq struct { + Ipv4Source string `json:"ipv4_source"` + Ipv4Format string `json:"ipv4_format"` + Ipv6Source string `json:"ipv6_source"` + Ipv6Format string `json:"ipv6_format"` +} + +// ---- 内部 helpers ---- + +// applyIPConfig 原子地更新单个 IP 配置项：① 写 DB → ② 同步 global +func applyIPConfig(item, value string) error { + if err := wafSystemConfigService.ModifyByItemApi(request.WafSystemConfigEditByItemReq{ + Item: item, + Value: value, + }); err != nil { + return err + } + switch item { + case "ip_v4_source": + global.GCONFIG_IP_V4_SOURCE = value + case "ip_v6_source": + global.GCONFIG_IP_V6_SOURCE = value + case "ip_v4_format": + global.GCONFIG_IP_V4_FORMAT = value + case "ip_v6_format": + global.GCONFIG_IP_V6_FORMAT = value + } + return nil +} + +// getConfigOrDefault 从 sys_config 读取配置项，空值返回默认值 +func getConfigOrDefault(item, def string) string { + bean := wafSystemConfigService.GetDetailByItemApi(request.WafSystemConfigDetailByItemReq{Item: item}) + if bean.Value == "" { + return def + } + return bean.Value +} + +// sourceFileExists 检查指定 ip 类型与 source 对应的物理文件是否存在 +func sourceFileExists(ipType, source, dataDir string) bool { + var fileName string + switch source { + case "ip2region": + if ipType == "ipv4" { + fileName = "ip2region.xdb" + } else { + fileName = "ip2region_v6.xdb" + } + case "geolite2": + fileName = "GeoLite2-Country.mmdb" + case "ipdb": + fileName = "iplocation.ipdb" + default: + return false + } + _, err := os.Stat(filepath.Join(dataDir, fileName)) + return err == nil +} + +// reloadManagerByCurrentConfig 根据 global 中当前的 source/format 重新加载 manager +// 所有 source 切换、配置变更、手动 reload 都应通过此函数集中处理 +func reloadManagerByCurrentConfig() error { + if global.GIPLOCATION_MANAGER == nil { + return nil + } + dataDir := filepath.Join(utils.GetCurrentDir(), "data") + + // ipdb 双栈共用，优先处理 + if global.GCONFIG_IP_V4_SOURCE == "ipdb" || global.GCONFIG_IP_V6_SOURCE == "ipdb" { + ipdbPath := filepath.Join(dataDir, "iplocation.ipdb") + if _, err := os.Stat(ipdbPath); err == nil { + if err = global.GIPLOCATION_MANAGER.LoadIpdb(ipdbPath); err != nil { + return fmt.Errorf("重新加载 ipdb 数据库失败: %w", err) + } + global.GIPLOCATION_MANAGER.SetBothSourceIpdb() + } + } + + // IPv4（非 ipdb 来源） + switch global.GCONFIG_IP_V4_SOURCE { + case "ip2region": + ipv4Path := filepath.Join(dataDir, "ip2region.xdb") + if _, err := os.Stat(ipv4Path); err == nil { + data, err := ioutil.ReadFile(ipv4Path) + if err == nil { + if err = global.GIPLOCATION_MANAGER.LoadV4Ip2Region(data, iplocation.DBFormat(global.GCONFIG_IP_V4_FORMAT)); err != nil { + return fmt.Errorf("重新加载 IPv4 数据库失败: %w", err) + } + } + } + case "geolite2": + ipv4Path := filepath.Join(dataDir, "GeoLite2-Country.mmdb") + if _, err := os.Stat(ipv4Path); err == nil { + data, err := ioutil.ReadFile(ipv4Path) + if err == nil { + if err = global.GIPLOCATION_MANAGER.LoadV4GeoLite2(data); err != nil { + return fmt.Errorf("重新加载 IPv4 数据库失败: %w", err) + } + } + } + } + + // IPv6（非 ipdb 来源） + switch global.GCONFIG_IP_V6_SOURCE { + case "ip2region": + ipv6Path := filepath.Join(dataDir, "ip2region_v6.xdb") + if _, err := os.Stat(ipv6Path); err == nil { + data, err := ioutil.ReadFile(ipv6Path) + if err == nil { + if err = global.GIPLOCATION_MANAGER.LoadV6Ip2Region(data, iplocation.DBFormat(global.GCONFIG_IP_V6_FORMAT)); err != nil { + return fmt.Errorf("重新加载 IPv6 数据库失败: %w", err) + } + } + } + case "geolite2": + ipv6Path := filepath.Join(dataDir, "GeoLite2-Country.mmdb") + if _, err := os.Stat(ipv6Path); err == nil { + data, err := ioutil.ReadFile(ipv6Path) + if err == nil { + if err = global.GIPLOCATION_MANAGER.LoadV6GeoLite2(data); err != nil { + return fmt.Errorf("重新加载 IPv6 数据库失败: %w", err) + } + } + } + } + + return nil +} + +// ---- API Handlers ---- + // GetIPDBStatusApi 获取 IP 数据库状态 func (w *WafIPLocationApi) GetIPDBStatusApi(c *gin.Context) { if global.GIPLOCATION_MANAGER == nil { @@ -25,15 +174,17 @@ func (w *WafIPLocationApi) GetIPDBStatusApi(c *gin.Context) { status := global.GIPLOCATION_MANAGER.GetStatus() - // 获取文件的实际创建时间 dataDir := filepath.Join(utils.GetCurrentDir(), "data") // 获取 IPv4 文件创建时间 var ipv4FilePath string - if status.IPv4Source == "ip2region" { + switch status.IPv4Source { + case "ip2region": ipv4FilePath = filepath.Join(dataDir, "ip2region.xdb") - } else if status.IPv4Source == "geolite2" { + case "geolite2": ipv4FilePath = filepath.Join(dataDir, "GeoLite2-Country.mmdb") + case "ipdb": + ipv4FilePath = filepath.Join(dataDir, "iplocation.ipdb") } if ipv4FilePath != "" { if fileInfo, err := os.Stat(ipv4FilePath); err == nil { @@ -43,10 +194,13 @@ func (w *WafIPLocationApi) GetIPDBStatusApi(c *gin.Context) { // 获取 IPv6 文件创建时间 var ipv6FilePath string - if status.IPv6Source == "ip2region" { + switch status.IPv6Source { + case "ip2region": ipv6FilePath = filepath.Join(dataDir, "ip2region_v6.xdb") - } else if status.IPv6Source == "geolite2" { + case "geolite2": ipv6FilePath = filepath.Join(dataDir, "GeoLite2-Country.mmdb") + case "ipdb": + ipv6FilePath = filepath.Join(dataDir, "iplocation.ipdb") } if ipv6FilePath != "" { if fileInfo, err := os.Stat(ipv6FilePath); err == nil { @@ -54,45 +208,171 @@ func (w *WafIPLocationApi) GetIPDBStatusApi(c *gin.Context) { } } - response.OkWithDetailed(status, "获取成功", c) + // 检查各数据库文件是否存在于磁盘 + checkFile := func(name string) bool { + _, err := os.Stat(filepath.Join(dataDir, name)) + return err == nil + } + fileExists := map[string]bool{ + "ip2region_v4": checkFile("ip2region.xdb"), + "ip2region_v6": checkFile("ip2region_v6.xdb"), + "geolite2": checkFile("GeoLite2-Country.mmdb"), + "ipdb": checkFile("iplocation.ipdb"), + } + + response.OkWithDetailed(ipdbStatusResponse{ + DBStatus: *status, + FileExists: fileExists, + }, "获取成功", c) +} + +// GetIPDBConfigApi 一次性读取 IP 数据库的 source/format 配置项 +func (w *WafIPLocationApi) GetIPDBConfigApi(c *gin.Context) { + resp := IPDBConfigResp{ + Ipv4Source: getConfigOrDefault("ip_v4_source", "ip2region"), + Ipv4Format: getConfigOrDefault("ip_v4_format", "legacy"), + Ipv6Source: getConfigOrDefault("ip_v6_source", "geolite2"), + Ipv6Format: getConfigOrDefault("ip_v6_format", "legacy"), + } + response.OkWithDetailed(resp, "获取成功", c) +} + +// SaveIPDBConfigApi 保存 IP 数据库配置，原子更新 DB+global 并重载 manager +func (w *WafIPLocationApi) SaveIPDBConfigApi(c *gin.Context) { + var req IPDBConfigReq + if err := c.ShouldBindJSON(&req); err != nil { + response.FailWithMessage("参数解析失败: "+err.Error(), c) + return + } + + // 合法值校验 + validV4 := map[string]bool{"ip2region": true, "ipdb": true} + validV6 := map[string]bool{"ip2region": true, "geolite2": true, "ipdb": true} + if !validV4[req.Ipv4Source] { + response.FailWithMessage("IPv4 数据源非法: "+req.Ipv4Source, c) + return + } + if !validV6[req.Ipv6Source] { + response.FailWithMessage("IPv6 数据源非法: "+req.Ipv6Source, c) + return + } + + // 文件存在性兜底校验 + dataDir := filepath.Join(utils.GetCurrentDir(), "data") + if !sourceFileExists("ipv4", req.Ipv4Source, dataDir) { + response.FailWithMessage("IPv4 所选来源的数据库文件不存在，请先上传", c) + return + } + if !sourceFileExists("ipv6", req.Ipv6Source, dataDir) { + response.FailWithMessage("IPv6 所选来源的数据库文件不存在，请先上传", c) + return + } + + // 同步 DB + global + if err := applyIPConfig("ip_v4_source", req.Ipv4Source); err != nil { + response.FailWithMessage("保存 IPv4 数据源失败: "+err.Error(), c) + return + } + if req.Ipv4Format != "" { + if err := applyIPConfig("ip_v4_format", req.Ipv4Format); err != nil { + response.FailWithMessage("保存 IPv4 字段格式失败: "+err.Error(), c) + return + } + } + if err := applyIPConfig("ip_v6_source", req.Ipv6Source); err != nil { + response.FailWithMessage("保存 IPv6 数据源失败: "+err.Error(), c) + return + } + if req.Ipv6Format != "" { + if err := applyIPConfig("ip_v6_format", req.Ipv6Format); err != nil { + response.FailWithMessage("保存 IPv6 字段格式失败: "+err.Error(), c) + return + } + } + + // 按最新 source 重载 manager + if err := reloadManagerByCurrentConfig(); err != nil { + response.FailWithMessage("配置已保存但重载失败: "+err.Error(), c) + return + } + + response.OkWithMessage("配置已保存并生效", c) } // UploadIPDBFileApi 上传 IP 数据库文件 func (w *WafIPLocationApi) UploadIPDBFileApi(c *gin.Context) { - // 获取文件类型 (ipv4/ipv6) ipType := c.PostForm("type") - if ipType != "ipv4" && ipType != "ipv6" { - response.FailWithMessage("无效的类型参数，必须是 ipv4 或 ipv6", c) + if ipType != "ipv4" && ipType != "ipv6" && ipType != "ipdb" { + response.FailWithMessage("无效的类型参数，必须是 ipv4、ipv6 或 ipdb", c) return } - // 获取上传的文件 file, err := c.FormFile("file") if err != nil { response.FailWithMessage("文件上传失败: "+err.Error(), c) return } - // 检查文件扩展名 ext := filepath.Ext(file.Filename) - if ext != ".xdb" && ext != ".mmdb" { - response.FailWithMessage("不支持的文件类型，仅支持 .xdb 和 .mmdb 文件", c) + if ext != ".xdb" && ext != ".mmdb" && ext != ".ipdb" { + response.FailWithMessage("不支持的文件类型，仅支持 .xdb、.mmdb 和 .ipdb 文件", c) + return + } + if ext == ".ipdb" && ipType != "ipdb" { + response.FailWithMessage(".ipdb 文件请使用 type=ipdb 上传", c) + return + } + if ext != ".ipdb" && ipType == "ipdb" { + response.FailWithMessage("type=ipdb 仅支持 .ipdb 文件", c) return } - // 确定保存路径 - var finalPath string dataDir := filepath.Join(utils.GetCurrentDir(), "data") - - // 确保 data 目录存在 if _, err := os.Stat(dataDir); os.IsNotExist(err) { - err = os.MkdirAll(dataDir, 0755) - if err != nil { + if err = os.MkdirAll(dataDir, 0755); err != nil { response.FailWithMessage("创建数据目录失败: "+err.Error(), c) return } } + // ipdb 双栈共用一个文件，独立处理 + if ext == ".ipdb" { + finalPath := filepath.Join(dataDir, "iplocation.ipdb") + tempPath := finalPath + ".tmp" + + if err = c.SaveUploadedFile(file, tempPath); err != nil { + response.FailWithMessage("保存临时文件失败: "+err.Error(), c) + return + } + + if global.GIPLOCATION_MANAGER != nil { + if err = global.GIPLOCATION_MANAGER.LoadIpdb(tempPath); err != nil { + os.Remove(tempPath) + response.FailWithMessage("加载 ipdb 数据库失败: "+err.Error(), c) + return + } + } + + if err = os.Rename(tempPath, finalPath); err != nil { + os.Remove(tempPath) + response.FailWithMessage("替换文件失败: "+err.Error(), c) + return + } + + // 从最终路径重新加载，更新 manager 内部元数据 + if global.GIPLOCATION_MANAGER != nil { + _ = global.GIPLOCATION_MANAGER.LoadIpdb(finalPath) + global.GIPLOCATION_MANAGER.SetBothSourceIpdb() + _ = applyIPConfig("ip_v4_source", "ipdb") + _ = applyIPConfig("ip_v6_source", "ipdb") + } + + response.OkWithMessage("ipdb 文件上传成功并已重新加载（IPv4+IPv6）", c) + return + } + + // xdb / mmdb 保存路径 + var finalPath string if ipType == "ipv4" { if ext == ".xdb" { finalPath = filepath.Join(dataDir, "ip2region.xdb") @@ -107,63 +387,51 @@ func (w *WafIPLocationApi) UploadIPDBFileApi(c *gin.Context) { } } - // 先保存到临时文件 tempPath := finalPath + ".tmp" - err = c.SaveUploadedFile(file, tempPath) - if err != nil { + if err = c.SaveUploadedFile(file, tempPath); err != nil { response.FailWithMessage("保存临时文件失败: "+err.Error(), c) return } - // 读取临时文件内容 fileData, err := ioutil.ReadFile(tempPath) if err != nil { - os.Remove(tempPath) // 清理临时文件 + os.Remove(tempPath) response.FailWithMessage("读取文件失败: "+err.Error(), c) return } - // 先热加载到内存，验证文件有效性 if global.GIPLOCATION_MANAGER != nil { var reloadErr error - + var sourceItem, sourceValue string if ipType == "ipv4" { if ext == ".xdb" { reloadErr = global.GIPLOCATION_MANAGER.LoadV4Ip2Region(fileData, iplocation.DBFormat(global.GCONFIG_IP_V4_FORMAT)) - if reloadErr == nil { - global.GCONFIG_IP_V4_SOURCE = "ip2region" - } + sourceItem, sourceValue = "ip_v4_source", "ip2region" } else { reloadErr = global.GIPLOCATION_MANAGER.LoadV4GeoLite2(fileData) - if reloadErr == nil { - global.GCONFIG_IP_V4_SOURCE = "geolite2" - } + sourceItem, sourceValue = "ip_v4_source", "geolite2" } } else { if ext == ".xdb" { reloadErr = global.GIPLOCATION_MANAGER.LoadV6Ip2Region(fileData, iplocation.DBFormat(global.GCONFIG_IP_V6_FORMAT)) - if reloadErr == nil { - global.GCONFIG_IP_V6_SOURCE = "ip2region" - } + sourceItem, sourceValue = "ip_v6_source", "ip2region" } else { reloadErr = global.GIPLOCATION_MANAGER.LoadV6GeoLite2(fileData) - if reloadErr == nil { - global.GCONFIG_IP_V6_SOURCE = "geolite2" - } + sourceItem, sourceValue = "ip_v6_source", "geolite2" } } if reloadErr != nil { - os.Remove(tempPath) // 清理临时文件 + os.Remove(tempPath) response.FailWithMessage("加载数据库失败: "+reloadErr.Error(), c) return } + + _ = applyIPConfig(sourceItem, sourceValue) } - // 热加载成功后，原子替换正式文件 - err = os.Rename(tempPath, finalPath) - if err != nil { - os.Remove(tempPath) // 清理临时文件 + if err = os.Rename(tempPath, finalPath); err != nil { + os.Remove(tempPath) response.FailWithMessage("替换文件失败: "+err.Error(), c) return } @@ -178,60 +446,9 @@ func (w *WafIPLocationApi) ReloadIPDBApi(c *gin.Context) { return } - dataDir := filepath.Join(utils.GetCurrentDir(), "data") - - // 重新加载 IPv4 - if global.GCONFIG_IP_V4_SOURCE == "ip2region" { - ipv4Path := filepath.Join(dataDir, "ip2region.xdb") - if _, err := os.Stat(ipv4Path); err == nil { - data, err := ioutil.ReadFile(ipv4Path) - if err == nil { - err = global.GIPLOCATION_MANAGER.LoadV4Ip2Region(data, iplocation.DBFormat(global.GCONFIG_IP_V4_FORMAT)) - if err != nil { - response.FailWithMessage("重新加载 IPv4 数据库失败: "+err.Error(), c) - return - } - } - } - } else if global.GCONFIG_IP_V4_SOURCE == "geolite2" { - ipv4Path := filepath.Join(dataDir, "GeoLite2-Country.mmdb") - if _, err := os.Stat(ipv4Path); err == nil { - data, err := ioutil.ReadFile(ipv4Path) - if err == nil { - err = global.GIPLOCATION_MANAGER.LoadV4GeoLite2(data) - if err != nil { - response.FailWithMessage("重新加载 IPv4 数据库失败: "+err.Error(), c) - return - } - } - } - } - - // 重新加载 IPv6 - if global.GCONFIG_IP_V6_SOURCE == "ip2region" { - ipv6Path := filepath.Join(dataDir, "ip2region_v6.xdb") - if _, err := os.Stat(ipv6Path); err == nil { - data, err := ioutil.ReadFile(ipv6Path) - if err == nil { - err = global.GIPLOCATION_MANAGER.LoadV6Ip2Region(data, iplocation.DBFormat(global.GCONFIG_IP_V6_FORMAT)) - if err != nil { - response.FailWithMessage("重新加载 IPv6 数据库失败: "+err.Error(), c) - return - } - } - } - } else if global.GCONFIG_IP_V6_SOURCE == "geolite2" { - ipv6Path := filepath.Join(dataDir, "GeoLite2-Country.mmdb") - if _, err := os.Stat(ipv6Path); err == nil { - data, err := ioutil.ReadFile(ipv6Path) - if err == nil { - err = global.GIPLOCATION_MANAGER.LoadV6GeoLite2(data) - if err != nil { - response.FailWithMessage("重新加载 IPv6 数据库失败: "+err.Error(), c) - return - } - } - } + if err := reloadManagerByCurrentConfig(); err != nil { + response.FailWithMessage(err.Error(), c) + return } response.OkWithMessage("数据库重新加载成功", c) @@ -243,8 +460,7 @@ func (w *WafIPLocationApi) TestIPLookupApi(c *gin.Context) { IP string `json:"ip" binding:"required"` } - err := c.ShouldBindJSON(&req) - if err != nil { + if err := c.ShouldBindJSON(&req); err != nil { response.FailWithMessage("参数解析失败", c) return } @@ -256,15 +472,32 @@ func (w *WafIPLocationApi) TestIPLookupApi(c *gin.Context) { result := global.GIPLOCATION_MANAGER.Lookup(req.IP) + // 识别 IP 类型 + 使用的数据源/格式 + ipType, usedSource, usedFormat := "未知", "", "" + if parsed := net.ParseIP(req.IP); parsed != nil { + if parsed.To4() != nil { + ipType = "IPv4" + usedSource = global.GCONFIG_IP_V4_SOURCE + usedFormat = global.GCONFIG_IP_V4_FORMAT + } else { + ipType = "IPv6" + usedSource = global.GCONFIG_IP_V6_SOURCE + usedFormat = global.GCONFIG_IP_V6_FORMAT + } + } + resp := map[string]interface{}{ - "ip": req.IP, - "country": result.Country, - "province": result.Province, - "city": result.City, - "isp": result.ISP, - "region": result.Region, - "district": result.District, - "raw": fmt.Sprintf("%v", result.ToSlice()), + "ip": req.IP, + "ip_type": ipType, + "used_source": usedSource, + "used_format": usedFormat, // 仅 ip2region 时有意义 + "country": result.Country, + "province": result.Province, + "city": result.City, + "isp": result.ISP, + "region": result.Region, + "district": result.District, + "raw": fmt.Sprintf("%v", result.ToSlice()), } response.OkWithDetailed(resp, "查询成功", c) diff --git a/cmd/samwaf/main.go b/cmd/samwaf/main.go index 4b75e5a..c95b0f8 100644 --- a/cmd/samwaf/main.go +++ b/cmd/samwaf/main.go @@ -187,6 +187,17 @@ func (m *wafSystenService) run() { log.Fatalf("Failed to load IPv4 GeoLite2 database: %v", err) } zlog.Info("IPv4 GeoLite2 database loaded successfully") + } else if global.GCONFIG_IP_V4_SOURCE == "ipdb" { + ipdbPath := filepath.Join(utils.GetCurrentDir(), "data", "iplocation.ipdb") + if _, err := os.Stat(ipdbPath); err == nil { + if err2 := global.GIPLOCATION_MANAGER.LoadIpdb(ipdbPath); err2 != nil { + zlog.Warn("Failed to load ipdb database (v4): ", err2) + } else { + zlog.Info("ipdb database loaded successfully (v4 source)") + } + } else { + zlog.Warn("ipdb database file not found, please upload iplocation.ipdb") + } } // 加载 IPv6 数据库 @@ -231,6 +242,22 @@ func (m *wafSystenService) run() { log.Fatalf("Failed to load IPv6 GeoLite2 database: %v", err) } zlog.Info("IPv6 GeoLite2 database loaded successfully") + } else if global.GCONFIG_IP_V6_SOURCE == "ipdb" { + // 如果 v4 已经加载了 ipdb，跳过重复加载 + if !global.GIPLOCATION_MANAGER.IsIpdbLoaded() { + ipdbPath := filepath.Join(utils.GetCurrentDir(), "data", "iplocation.ipdb") + if _, err := os.Stat(ipdbPath); err == nil { + if err2 := global.GIPLOCATION_MANAGER.LoadIpdb(ipdbPath); err2 != nil { + zlog.Warn("Failed to load ipdb database (v6): ", err2) + } else { + zlog.Info("ipdb database loaded successfully (v6 source)") + } + } else { + zlog.Warn("ipdb database file not found, please upload iplocation.ipdb") + } + } else { + zlog.Info("ipdb database already loaded (shared with v4)") + } } global.GWAF_DLP_CONFIG = ldpConfig global.GWAF_REG_PUBLIC_KEY = publicKey diff --git a/go.mod b/go.mod index 6d8c8ee..9d63ffa 100644 --- a/go.mod +++ b/go.mod @@ -103,6 +103,7 @@ require ( github.com/golang/protobuf v1.5.4 // indirect github.com/hashicorp/yamux v0.1.2 // indirect github.com/huaweicloud/huaweicloud-sdk-go-v3 v0.1.180 // indirect + github.com/ipipdotnet/ipdb-go v1.3.3 // indirect github.com/jbenet/go-context v0.0.0-20150711004518-d14ea06fba99 // indirect github.com/jinzhu/inflection v1.0.0 // indirect github.com/jinzhu/now v1.1.5 // indirect diff --git a/go.sum b/go.sum index 984d316..3e54055 100644 --- a/go.sum +++ b/go.sum @@ -234,6 +234,8 @@ github.com/huaweicloud/huaweicloud-sdk-go-v3 v0.1.180 h1:uia+R3K1izQRGpxTV+bS4q3 github.com/huaweicloud/huaweicloud-sdk-go-v3 v0.1.180/go.mod h1:M+yna96Fx9o5GbIUnF3OvVvQGjgfVSyeJbV9Yb1z/wI= github.com/hyperjumptech/grule-rule-engine v1.15.0 h1:HqCjhZK+YsNC6udTR6/O90xRwxcefTwStheATUjYK34= github.com/hyperjumptech/grule-rule-engine v1.15.0/go.mod h1:K8HweZ21+ccFgIfXxyJbAuUZU2OAIapCWhZv1a7GP/8= +github.com/ipipdotnet/ipdb-go v1.3.3 h1:GLSAW9ypLUd6EF9QNK2Uhxew9Jzs4XMJ9gOZEFnJm7U= +github.com/ipipdotnet/ipdb-go v1.3.3/go.mod h1:yZ+8puwe3R37a/3qRftXo40nZVQbxYDLqls9o5foexs= github.com/jbenet/go-context v0.0.0-20150711004518-d14ea06fba99 h1:BQSFePA1RWJOlocH6Fxy8MmwDt+yVQYULKfN0RoTN8A= github.com/jbenet/go-context v0.0.0-20150711004518-d14ea06fba99/go.mod h1:1lJo3i6rXxKeerYnT8Nvf0QmHCRC1n8sfWVwXF2Frvo= github.com/jcchavezs/mergefs v0.1.0 h1:7oteO7Ocl/fnfFMkoVLJxTveCjrsd//UB0j89xmnpec= diff --git a/iplocation/manager.go b/iplocation/manager.go index 24a6283..afbe6a8 100644 --- a/iplocation/manager.go +++ b/iplocation/manager.go @@ -3,9 +3,11 @@ package iplocation import ( "fmt" "net" + "os" "sync" "time" + "github.com/ipipdotnet/ipdb-go" "github.com/lionsoul2014/ip2region/binding/golang/xdb" "github.com/oschwald/geoip2-golang" ) @@ -31,6 +33,12 @@ type Manager struct { v6LoadTime time.Time v6FileSize int64 v6CreateTime time.Time // 文件创建时间 + + // ipdb 后端（IPv4+IPv6 共用同一个文件） + ipdbReader *ipdb.City + ipdbLoadTime time.Time + ipdbFileSize int64 + ipdbCreateTime time.Time } // NewManager 创建新的 IP 地理位置查询管理器 @@ -86,6 +94,12 @@ func (m *Manager) lookupV4(ipStr string) *IPLocationResult { countryName = record.Country.Names["en"] } return &IPLocationResult{Country: countryName} + } else if m.v4Source == SourceIpdb && m.ipdbReader != nil { + info, err := m.ipdbReader.FindMap(ipStr, "CN") + if err != nil { + return &IPLocationResult{Country: "查询失败"} + } + return parseIpdbMap(info) } return &IPLocationResult{Country: "未配置"} @@ -116,6 +130,12 @@ func (m *Manager) lookupV6(ipStr string) *IPLocationResult { countryName = "内网" } return &IPLocationResult{Country: countryName} + } else if m.v6Source == SourceIpdb && m.ipdbReader != nil { + info, err := m.ipdbReader.FindMap(ipStr, "CN") + if err != nil { + return &IPLocationResult{Country: "查询失败"} + } + return parseIpdbMap(info) } return &IPLocationResult{Country: "未配置"} @@ -227,6 +247,31 @@ func (m *Manager) LoadV6GeoLite2(data []byte) error { return nil } +// LoadIpdb 加载 ipdb 数据库（ipdb 支持 IPv4+IPv6，共用同一文件） +func (m *Manager) LoadIpdb(filePath string) error { + m.mu.Lock() + defer m.mu.Unlock() + + reader, err := ipdb.NewCity(filePath) + if err != nil { + return fmt.Errorf("创建 ipdb reader 失败: %w", err) + } + m.ipdbReader = reader + m.ipdbLoadTime = time.Now() + if fi, err2 := os.Stat(filePath); err2 == nil { + m.ipdbFileSize = fi.Size() + m.ipdbCreateTime = fi.ModTime() + } + return nil +} + +// IsIpdbLoaded 返回 ipdb reader 是否已加载 +func (m *Manager) IsIpdbLoaded() bool { + m.mu.RLock() + defer m.mu.RUnlock() + return m.ipdbReader != nil +} + // ReloadV4 热加载 IPv4 数据库 func (m *Manager) ReloadV4(data []byte, source DBSource, format DBFormat) error { if source == SourceIp2Region { @@ -261,17 +306,39 @@ func (m *Manager) GetStatus() *DBStatus { IPv6FileSize: m.v6FileSize, } - if !m.v4LoadTime.IsZero() { - status.IPv4LoadTime = m.v4LoadTime.Format("2006-01-02 15:04:05") + // ipdb 来源使用 ipdb 的文件元数据覆盖对应槽位 + if m.v4Source == SourceIpdb { + status.IPv4FileSize = m.ipdbFileSize + if !m.ipdbLoadTime.IsZero() { + status.IPv4LoadTime = m.ipdbLoadTime.Format("2006-01-02 15:04:05") + } + if !m.ipdbCreateTime.IsZero() { + status.IPv4CreateTime = m.ipdbCreateTime.Format("2006-01-02 15:04:05") + } + } else { + if !m.v4LoadTime.IsZero() { + status.IPv4LoadTime = m.v4LoadTime.Format("2006-01-02 15:04:05") + } + if !m.v4CreateTime.IsZero() { + status.IPv4CreateTime = m.v4CreateTime.Format("2006-01-02 15:04:05") + } } - if !m.v6LoadTime.IsZero() { - status.IPv6LoadTime = m.v6LoadTime.Format("2006-01-02 15:04:05") - } - if !m.v4CreateTime.IsZero() { - status.IPv4CreateTime = m.v4CreateTime.Format("2006-01-02 15:04:05") - } - if !m.v6CreateTime.IsZero() { - status.IPv6CreateTime = m.v6CreateTime.Format("2006-01-02 15:04:05") + + if m.v6Source == SourceIpdb { + status.IPv6FileSize = m.ipdbFileSize + if !m.ipdbLoadTime.IsZero() { + status.IPv6LoadTime = m.ipdbLoadTime.Format("2006-01-02 15:04:05") + } + if !m.ipdbCreateTime.IsZero() { + status.IPv6CreateTime = m.ipdbCreateTime.Format("2006-01-02 15:04:05") + } + } else { + if !m.v6LoadTime.IsZero() { + status.IPv6LoadTime = m.v6LoadTime.Format("2006-01-02 15:04:05") + } + if !m.v6CreateTime.IsZero() { + status.IPv6CreateTime = m.v6CreateTime.Format("2006-01-02 15:04:05") + } } return status @@ -298,6 +365,20 @@ func (m *Manager) Close() { m.v6GeoReader.Close() m.v6GeoReader = nil } + // ipdb.City 无 Close 方法，置 nil 由 GC 回收 + m.ipdbReader = nil +} + +// parseIpdbMap 将 ipdb FindMap 返回的字段映射转换为统一结果结构 +func parseIpdbMap(info map[string]string) *IPLocationResult { + get := func(k string) string { return info[k] } + return &IPLocationResult{ + Country: get("country_name"), + Province: get("region_name"), + City: get("city_name"), + ISP: get("isp_domain"), + Region: get("continent_code"), + } } // SetV4Source 设置 IPv4 数据源 @@ -314,6 +395,14 @@ func (m *Manager) SetV6Source(source DBSource) { m.v6Source = source } +// SetBothSourceIpdb 将 IPv4 和 IPv6 数据源同时设置为 ipdb +func (m *Manager) SetBothSourceIpdb() { + m.mu.Lock() + defer m.mu.Unlock() + m.v4Source = SourceIpdb + m.v6Source = SourceIpdb +} + // SetV4Format 设置 IPv4 数据格式 func (m *Manager) SetV4Format(format DBFormat) { m.mu.Lock() diff --git a/iplocation/types.go b/iplocation/types.go index 92a3a45..1c658e6 100644 --- a/iplocation/types.go +++ b/iplocation/types.go @@ -32,6 +32,7 @@ type DBSource string const ( SourceIp2Region DBSource = "ip2region" SourceGeoLite2 DBSource = "geolite2" + SourceIpdb DBSource = "ipdb" ) // DBStatus 数据库状态信息 diff --git a/router/waf_iplocation.go b/router/waf_iplocation.go index 381e726..27d38c2 100644 --- a/router/waf_iplocation.go +++ b/router/waf_iplocation.go @@ -13,6 +13,8 @@ func (receiver *IPLocationRouter) InitIPLocationRouter(group *gin.RouterGroup) { router := group.Group("/api/v1/iplocation") { router.GET("/status", apiInstance.GetIPDBStatusApi) + router.GET("/config", apiInstance.GetIPDBConfigApi) + router.POST("/config/save", apiInstance.SaveIPDBConfigApi) router.POST("/upload", apiInstance.UploadIPDBFileApi) router.POST("/reload", apiInstance.ReloadIPDBApi) router.POST("/test", apiInstance.TestIPLookupApi)