actr_framework/util/
geoip.rs

1//! GeoIP 地理位置查询工具
2//!
3//! 基于 MaxMind GeoLite2 数据库提供 IP 地址到地理坐标的转换
4//!
5//! # 使用示例
6//!
7//! ```rust,ignore
8//! use actr_framework::util::geoip::GeoIpService;
9//! use std::net::IpAddr;
10//!
11//! // 初始化 GeoIP 服务
12//! let geoip = GeoIpService::new("data/geoip/GeoLite2-City.mmdb")?;
13//!
14//! // 查询 IP 地址的坐标
15//! let ip: IpAddr = "8.8.8.8".parse()?;
16//! if let Some((lat, lon)) = geoip.lookup(ip) {
17//!     println!("IP {} 位于坐标: ({}, {})", ip, lat, lon);
18//! }
19//! ```
20//!
21//! # 获取 GeoLite2 数据库
22//!
23//! **自动下载(推荐):**
24//! 1. 访问 https://www.maxmind.com/en/geolite2/signup 获取 License Key
25//! 2. 设置环境变量:`export MAXMIND_LICENSE_KEY="your-key-here"`
26//! 3. 首次调用 `GeoIpService::new()` 时自动下载
27//!
28//! **手动下载(生产环境):**
29//! ```bash
30//! curl -o GeoLite2-City.tar.gz \
31//!   "https://download.maxmind.com/app/geoip_download?edition_id=GeoLite2-City&license_key=YOUR_KEY&suffix=tar.gz"
32//! tar -xzf GeoLite2-City.tar.gz --strip-components=1 -C data/geoip/ "*/GeoLite2-City.mmdb"
33//! ```
34
35#[cfg(feature = "geoip")]
36use anyhow::{Context, Result};
37#[cfg(feature = "geoip")]
38use maxminddb::{MaxMindDBError, Reader, geoip2::City};
39#[cfg(feature = "geoip")]
40use std::net::IpAddr;
41#[cfg(feature = "geoip")]
42use std::path::Path;
43#[cfg(feature = "geoip")]
44use tracing::{debug, info, warn};
45
46/// GeoIP 查询服务
47///
48/// 提供 IP 地址到地理坐标的转换功能
49#[cfg(feature = "geoip")]
50#[derive(Debug)]
51pub struct GeoIpService {
52    reader: Reader<Vec<u8>>,
53}
54
55#[cfg(feature = "geoip")]
56impl GeoIpService {
57    /// 初始化 GeoIP 服务(支持自动下载)
58    ///
59    /// # Arguments
60    /// * `db_path` - GeoLite2-City.mmdb 数据库文件路径
61    ///
62    /// # Errors
63    /// 如果数据库文件不存在或格式错误,返回错误
64    ///
65    /// # 自动下载
66    /// 如果数据库文件不存在,且设置了 `MAXMIND_LICENSE_KEY` 环境变量,
67    /// 将自动从 MaxMind 下载 GeoLite2-City 数据库。
68    pub fn new<P: AsRef<Path>>(db_path: P) -> Result<Self> {
69        let path = db_path.as_ref();
70
71        // 如果数据库不存在,尝试自动下载
72        if !path.exists() {
73            info!("GeoIP database not found at {:?}", path);
74
75            if let Ok(license_key) = std::env::var("MAXMIND_LICENSE_KEY") {
76                info!("MAXMIND_LICENSE_KEY found, attempting auto-download...");
77                Self::download_database(path, &license_key)?;
78            } else {
79                anyhow::bail!(
80                    "GeoIP database not found at {:?}\n\
81                     \n\
82                     To auto-download:\n\
83                     1. Get License Key: https://www.maxmind.com/en/geolite2/signup\n\
84                     2. export MAXMIND_LICENSE_KEY=\"your-key\"\n\
85                     3. Retry\n\
86                     \n\
87                     Or manually download:\n\
88                     curl -o GeoLite2-City.tar.gz \\\n\
89                       'https://download.maxmind.com/app/geoip_download?edition_id=GeoLite2-City&license_key=YOUR_KEY&suffix=tar.gz'\n\
90                     tar -xzf GeoLite2-City.tar.gz --strip-components=1 -C {:?}/ '*/GeoLite2-City.mmdb'",
91                    path,
92                    path.parent().unwrap_or(Path::new("."))
93                );
94            }
95        }
96
97        info!("Loading GeoIP database from: {:?}", path);
98        let reader = Reader::open_readfile(path)
99            .context(format!("Failed to open GeoIP database at {path:?}"))?;
100
101        info!(
102            "✅ GeoIP service initialized (build epoch: {})",
103            reader.metadata.build_epoch
104        );
105        Ok(Self { reader })
106    }
107
108    /// 自动下载 GeoLite2-City 数据库
109    fn download_database(db_path: &Path, license_key: &str) -> Result<()> {
110        use reqwest::blocking::Client;
111        use std::io::Read;
112
113        info!("📥 Downloading GeoLite2-City database (~70MB)...");
114
115        // 构建下载 URL
116        let url = format!(
117            "https://download.maxmind.com/app/geoip_download?edition_id=GeoLite2-City&license_key={}&suffix=tar.gz",
118            license_key
119        );
120
121        // 下载
122        let client = Client::builder()
123            .timeout(std::time::Duration::from_secs(300)) // 5 分钟超时
124            .build()?;
125
126        let response = client
127            .get(&url)
128            .send()
129            .context("Failed to download GeoLite2 database")?;
130
131        if !response.status().is_success() {
132            anyhow::bail!(
133                "Download failed with status: {} - Check your MAXMIND_LICENSE_KEY",
134                response.status()
135            );
136        }
137
138        info!("📦 Download complete, extracting...");
139
140        // 解压 tar.gz
141        let tar_gz_data = response.bytes()?;
142        let tar_decoder = flate2::read::GzDecoder::new(&tar_gz_data[..]);
143        let mut archive = tar::Archive::new(tar_decoder);
144
145        // 查找并提取 .mmdb 文件
146        for entry in archive.entries()? {
147            let mut entry = entry?;
148            let path_in_archive = entry.path()?;
149
150            if path_in_archive.extension() == Some(std::ffi::OsStr::new("mmdb"))
151                && path_in_archive.to_string_lossy().contains("GeoLite2-City")
152            {
153                // 创建父目录
154                if let Some(parent) = db_path.parent() {
155                    std::fs::create_dir_all(parent)?;
156                }
157
158                // 解压到目标位置
159                let mut output = std::fs::File::create(db_path)?;
160                std::io::copy(&mut entry, &mut output)?;
161
162                let size = std::fs::metadata(db_path)?.len();
163                info!(
164                    "✅ GeoIP database downloaded to {:?} ({:.1} MB)",
165                    db_path,
166                    size as f64 / 1_048_576.0
167                );
168                return Ok(());
169            }
170        }
171
172        anyhow::bail!("GeoLite2-City.mmdb not found in downloaded archive");
173    }
174
175    /// 查询 IP 地址的地理坐标
176    ///
177    /// # Arguments
178    /// * `ip` - 要查询的 IP 地址
179    ///
180    /// # Returns
181    /// * `Some((latitude, longitude))` - 成功找到坐标
182    /// * `None` - IP 地址不在数据库中或无坐标信息
183    pub fn lookup(&self, ip: IpAddr) -> Option<(f64, f64)> {
184        match self.reader.lookup::<City>(ip) {
185            Ok(city) => {
186                if let Some(location) = city.location {
187                    if let (Some(lat), Some(lon)) = (location.latitude, location.longitude) {
188                        debug!("GeoIP lookup: {} -> ({}, {})", ip, lat, lon);
189                        return Some((lat, lon));
190                    }
191                }
192                debug!("GeoIP lookup: {} found but no coordinates", ip);
193                None
194            }
195            Err(MaxMindDBError::AddressNotFoundError(_)) => {
196                debug!("GeoIP lookup: {} not in database", ip);
197                None
198            }
199            Err(e) => {
200                warn!("GeoIP lookup error for {}: {}", ip, e);
201                None
202            }
203        }
204    }
205
206    /// 获取数据库元信息
207    pub fn metadata(&self) -> &maxminddb::Metadata {
208        &self.reader.metadata
209    }
210}
211
212/// 无 GeoIP 功能时的降级实现
213#[cfg(not(feature = "geoip"))]
214#[derive(Debug)]
215pub struct GeoIpService;
216
217#[cfg(not(feature = "geoip"))]
218impl GeoIpService {
219    /// 初始化失败(需要启用 geoip feature)
220    pub fn new<P>(_db_path: P) -> anyhow::Result<Self> {
221        anyhow::bail!("GeoIP feature is not enabled. Rebuild with --features geoip")
222    }
223
224    /// 总是返回 None(需要启用 geoip feature)
225    pub fn lookup(&self, _ip: std::net::IpAddr) -> Option<(f64, f64)> {
226        None
227    }
228}
229
230#[cfg(test)]
231mod tests {
232    use super::*;
233
234    #[test]
235    fn test_geoip_module_compiles() {
236        // 确保模块编译通过
237        assert!(true);
238    }
239
240    #[cfg(feature = "geoip")]
241    #[test]
242    fn test_geoip_lookup_requires_database() {
243        // 测试需要真实的数据库文件,这里只验证 API 可用性
244        let result = GeoIpService::new("/nonexistent/path.mmdb");
245        assert!(result.is_err());
246    }
247
248    #[cfg(not(feature = "geoip"))]
249    #[test]
250    fn test_geoip_feature_disabled() {
251        let result = GeoIpService::new("/any/path.mmdb");
252        assert!(result.is_err());
253        assert!(
254            result
255                .unwrap_err()
256                .to_string()
257                .contains("GeoIP feature is not enabled")
258        );
259    }
260}