Skip to content

Commit fb417e2

Browse files
ThisSeanZhangThisSeanZhang
committed
feat: support TXT geo IP imports with configurable key
1 parent 3d08568 commit fb417e2

File tree

10 files changed

+347
-59
lines changed

10 files changed

+347
-59
lines changed

‎landscape-common/src/config/geo.rs‎

Lines changed: 32 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,27 @@ pub enum GeoIpError {
5050
#[error("Geo IP file read error")]
5151
#[api_error(id = "geo_ip.file_read_error", status = 400)]
5252
FileReadError,
53+
54+
#[error("Geo IP config '{0}' not found")]
55+
#[api_error(id = "geo_ip.config_not_found", status = 404)]
56+
ConfigNotFound(String),
57+
58+
#[error("Geo IP DAT decode error")]
59+
#[api_error(id = "geo_ip.dat_decode_error", status = 400)]
60+
DatDecodeError,
61+
62+
#[error("Geo IP TXT file contains no valid CIDR entries")]
63+
#[api_error(id = "geo_ip.no_valid_cidr", status = 400)]
64+
NoValidCidrFound,
65+
}
66+
67+
#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, Default)]
68+
#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
69+
#[serde(rename_all = "snake_case")]
70+
pub enum GeoIpFileFormat {
71+
#[default]
72+
Dat,
73+
Txt,
5374
}
5475

5576
#[derive(Serialize, Deserialize, Clone, Debug)]
@@ -187,8 +208,17 @@ pub struct GeoIpSourceConfig {
187208
#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
188209
#[serde(tag = "t", rename_all = "snake_case")]
189210
pub enum GeoIpSource {
190-
Url { url: String, next_update_at: f64 },
191-
Direct { data: Vec<GeoIpDirectItem> },
211+
Url {
212+
url: String,
213+
next_update_at: f64,
214+
#[serde(default)]
215+
format: GeoIpFileFormat,
216+
#[serde(default)]
217+
txt_key: Option<String>,
218+
},
219+
Direct {
220+
data: Vec<GeoIpDirectItem>,
221+
},
192222
}
193223

194224
#[derive(Serialize, Deserialize, Clone, Debug)]

‎landscape-protobuf/src/lib.rs‎

Lines changed: 135 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,13 +5,24 @@ use std::{
55
};
66

77
use landscape_common::{
8-
config::{dns::DomainMatchType, geo::GeoSiteFileConfig},
8+
config::{
9+
dns::DomainMatchType,
10+
geo::{GeoIpError, GeoIpFileFormat, GeoSiteFileConfig},
11+
},
912
ip_mark::IpConfig,
1013
};
1114
use protos::geo::{mod_Domain::Type, Domain, GeoIPListOwned, GeoSiteListOwned};
1215

1316
mod protos;
1417

18+
pub const DEFAULT_TXT_GEO_KEY: &str = "DEFAULT";
19+
20+
pub struct GeoIpParseResult {
21+
pub entries: HashMap<String, Vec<IpConfig>>,
22+
pub valid_lines: usize,
23+
pub skipped_lines: usize,
24+
}
25+
1526
pub async fn read_geo_sites_from_bytes(
1627
contents: impl Into<Vec<u8>>,
1728
) -> HashMap<String, Vec<GeoSiteFileConfig>> {
@@ -69,6 +80,73 @@ pub async fn read_geo_ips_from_bytes(
6980
result
7081
}
7182

83+
pub async fn read_geo_ips_from_bytes_dat(
84+
contents: impl Into<Vec<u8>>,
85+
) -> Result<HashMap<String, Vec<IpConfig>>, GeoIpError> {
86+
let mut result = HashMap::new();
87+
let list = GeoIPListOwned::try_from(contents.into()).map_err(|_| GeoIpError::DatDecodeError)?;
88+
89+
for entry in list.proto().entry.iter() {
90+
let domains = entry.cidr.iter().filter_map(convert_ipconfig_from_proto).collect();
91+
result.insert(entry.country_code.to_string(), domains);
92+
}
93+
94+
Ok(result)
95+
}
96+
97+
pub fn read_geo_ips_from_bytes_txt(
98+
contents: impl AsRef<[u8]>,
99+
txt_key: Option<&str>,
100+
) -> Result<GeoIpParseResult, GeoIpError> {
101+
let key = txt_key
102+
.map(str::trim)
103+
.filter(|value| !value.is_empty())
104+
.unwrap_or(DEFAULT_TXT_GEO_KEY)
105+
.to_ascii_uppercase();
106+
107+
let text = String::from_utf8_lossy(contents.as_ref());
108+
let mut values = Vec::new();
109+
let mut skipped_lines = 0;
110+
111+
for line in text.lines() {
112+
let line = line.trim();
113+
if line.is_empty() || line.starts_with('#') {
114+
continue;
115+
}
116+
117+
if let Some(cidr) = parse_txt_cidr_line(line) {
118+
values.push(cidr);
119+
} else {
120+
skipped_lines += 1;
121+
}
122+
}
123+
124+
if values.is_empty() {
125+
return Err(GeoIpError::NoValidCidrFound);
126+
}
127+
128+
let valid_lines = values.len();
129+
let mut entries = HashMap::new();
130+
entries.insert(key, values);
131+
132+
Ok(GeoIpParseResult { entries, valid_lines, skipped_lines })
133+
}
134+
135+
pub async fn read_geo_ips_from_bytes_by_format(
136+
contents: impl Into<Vec<u8>>,
137+
format: &GeoIpFileFormat,
138+
txt_key: Option<&str>,
139+
) -> Result<GeoIpParseResult, GeoIpError> {
140+
let contents = contents.into();
141+
match format {
142+
GeoIpFileFormat::Dat => {
143+
let entries = read_geo_ips_from_bytes_dat(contents).await?;
144+
Ok(GeoIpParseResult { entries, valid_lines: 0, skipped_lines: 0 })
145+
}
146+
GeoIpFileFormat::Txt => read_geo_ips_from_bytes_txt(&contents, txt_key),
147+
}
148+
}
149+
72150
pub async fn read_geo_ips<T: AsRef<Path>>(geo_file_path: T) -> HashMap<String, Vec<IpConfig>> {
73151
let mut result = HashMap::new();
74152
let data = tokio::fs::read(geo_file_path).await.unwrap();
@@ -99,18 +177,36 @@ pub fn convert_ipconfig_from_proto(value: &crate::protos::geo::CIDR) -> Option<I
99177
result.map(|ip| IpConfig { ip, prefix: value.prefix })
100178
}
101179

180+
fn parse_txt_cidr_line(line: &str) -> Option<IpConfig> {
181+
let (ip, prefix) = line.split_once('/')?;
182+
let ip: IpAddr = ip.trim().parse().ok()?;
183+
let prefix: u32 = prefix.trim().parse().ok()?;
184+
185+
let max_prefix = match ip {
186+
IpAddr::V4(_) => 32,
187+
IpAddr::V6(_) => 128,
188+
};
189+
190+
if prefix > max_prefix {
191+
return None;
192+
}
193+
194+
Some(IpConfig { ip, prefix })
195+
}
196+
102197
#[cfg(test)]
103198
#[global_allocator]
104199
static GLOBAL: jemallocator::Jemalloc = jemallocator::Jemalloc;
105200

106201
#[cfg(test)]
107202
mod tests {
203+
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
108204

109205
use jemalloc_ctl::{epoch, stats};
110206

111207
use crate::{
112208
protos::geo::{GeoIPListOwned, GeoSiteListOwned},
113-
read_geo_sites,
209+
read_geo_ips_from_bytes_txt, read_geo_sites,
114210
};
115211

116212
fn test_memory_usage() {
@@ -179,4 +275,41 @@ mod tests {
179275
println!("other count: {sum:?}");
180276
test_memory_usage();
181277
}
278+
279+
#[test]
280+
fn parse_txt_geo_ips_skips_invalid_lines() {
281+
let result = read_geo_ips_from_bytes_txt(
282+
b"\n# comment\n1.1.1.0/24\ninvalid\n2001:db8::/32\n10.0.0.1/33\n",
283+
Some("custom"),
284+
)
285+
.unwrap();
286+
287+
assert_eq!(result.valid_lines, 2);
288+
assert_eq!(result.skipped_lines, 2);
289+
assert_eq!(result.entries.len(), 1);
290+
291+
let values = result.entries.get("CUSTOM").unwrap();
292+
assert_eq!(
293+
values[0],
294+
landscape_common::ip_mark::IpConfig {
295+
ip: IpAddr::V4(Ipv4Addr::new(1, 1, 1, 0)),
296+
prefix: 24,
297+
}
298+
);
299+
assert_eq!(
300+
values[1],
301+
landscape_common::ip_mark::IpConfig {
302+
ip: IpAddr::V6(Ipv6Addr::from([
303+
0x20, 0x01, 0x0d, 0xb8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
304+
])),
305+
prefix: 32,
306+
}
307+
);
308+
}
309+
310+
#[test]
311+
fn parse_txt_geo_ips_uses_default_key() {
312+
let result = read_geo_ips_from_bytes_txt(b"1.1.1.0/24\n", Some(" ")).unwrap();
313+
assert!(result.entries.contains_key("DEFAULT"));
314+
}
182315
}

‎landscape-webserver/src/geo/ips.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -222,7 +222,7 @@ async fn update_by_upload(
222222
return Err(GeoIpError::FileReadError)?;
223223
};
224224

225-
state.geo_ip_service.update_geo_config_by_bytes(name, bytes).await;
225+
state.geo_ip_service.update_geo_config_by_bytes(name, bytes).await?;
226226

227227
LandscapeApiResp::success(())
228228
}

‎landscape-webui/src/components/geo/ip/config/GeoIpEditModal.vue‎

Lines changed: 38 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,10 @@ const sourceType = ref<"url" | "direct">("url");
2626
async function enter() {
2727
if (props.id !== null) {
2828
rule.value = await get_geo_ip_config(props.id);
29+
if (rule.value.source.t === "url") {
30+
rule.value.source.format = rule.value.source.format || "dat";
31+
rule.value.source.txt_key = rule.value.source.txt_key || null;
32+
}
2933
sourceType.value = rule.value.source.t;
3034
} else {
3135
sourceType.value = "url";
@@ -34,7 +38,13 @@ async function enter() {
3438
update_at: new Date().getTime(),
3539
name: "",
3640
enable: true,
37-
source: { t: "url", url: "", next_update_at: 0 },
41+
source: {
42+
t: "url",
43+
url: "",
44+
next_update_at: 0,
45+
format: "dat",
46+
txt_key: null,
47+
},
3848
};
3949
}
4050
rule_json.value = JSON.stringify(rule.value);
@@ -43,7 +53,13 @@ async function enter() {
4353
function switchSourceType(t: "url" | "direct") {
4454
if (!rule.value) return;
4555
if (t === "url") {
46-
rule.value.source = { t: "url", url: "", next_update_at: 0 };
56+
rule.value.source = {
57+
t: "url",
58+
url: "",
59+
next_update_at: 0,
60+
format: "dat",
61+
txt_key: null,
62+
};
4763
} else {
4864
rule.value.source = { t: "direct", data: [] };
4965
}
@@ -173,6 +189,26 @@ const rules: FormRules = {
173189
<n-form-item-gi :label="t('geo_editor.common.source_url')" :span="5">
174190
<n-input v-model:value="rule.source.url" clearable />
175191
</n-form-item-gi>
192+
<n-form-item-gi
193+
:label="t('geo_editor.common.source_format')"
194+
:span="5"
195+
>
196+
<n-radio-group v-model:value="rule.source.format">
197+
<n-radio value="dat">DAT</n-radio>
198+
<n-radio value="txt">TXT</n-radio>
199+
</n-radio-group>
200+
</n-form-item-gi>
201+
<n-form-item-gi
202+
v-if="rule.source.format === 'txt'"
203+
:label="t('geo_editor.common.txt_key')"
204+
:span="5"
205+
>
206+
<n-input
207+
v-model:value="rule.source.txt_key"
208+
clearable
209+
:placeholder="t('geo_editor.common.txt_key_placeholder')"
210+
/>
211+
</n-form-item-gi>
176212
</template>
177213

178214
<!-- Direct mode -->

‎landscape-webui/src/components/geo/ip/config/GeoIpItemCard.vue‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,19 @@ const onGeoUpload = async (formData: FormData) => {
5656
</n-tag>
5757
</n-descriptions-item>
5858
<template v-if="geo_ip_source.source.t === 'url'">
59+
<n-descriptions-item :label="t('geo_editor.item_card.source_format')">
60+
<n-tag
61+
:bordered="false"
62+
:type="
63+
(geo_ip_source.source.format || 'dat') === 'txt'
64+
? 'warning'
65+
: 'info'
66+
"
67+
size="small"
68+
>
69+
{{ (geo_ip_source.source.format || "dat").toUpperCase() }}
70+
</n-tag>
71+
</n-descriptions-item>
5972
<n-descriptions-item label="URL">
6073
<n-ellipsis style="max-width: 200px">
6174
{{
@@ -74,6 +87,12 @@ const onGeoUpload = async (formData: FormData) => {
7487
:time-zone="prefStore.timezone"
7588
/>
7689
</n-descriptions-item>
90+
<n-descriptions-item
91+
v-if="(geo_ip_source.source.format || 'dat') === 'txt'"
92+
:label="t('geo_editor.item_card.txt_key')"
93+
>
94+
{{ (geo_ip_source.source.txt_key || "DEFAULT").toUpperCase() }}
95+
</n-descriptions-item>
7796
</template>
7897
<template v-if="geo_ip_source.source.t === 'direct'">
7998
<n-descriptions-item :label="t('geo_editor.item_card.key_count')">

‎landscape-webui/src/i18n/en/api_errors.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,9 @@ export default {
1818
"geo_ip.cache_not_found": "GeoIP cache not found (key: {0})",
1919
"geo_ip.file_not_found": "GeoIP file not found in upload",
2020
"geo_ip.file_read_error": "GeoIP file read error",
21+
"geo_ip.config_not_found": "GeoIP config not found ({0})",
22+
"geo_ip.dat_decode_error": "GeoIP DAT file decode error",
23+
"geo_ip.no_valid_cidr": "GeoIP TXT file contains no valid CIDR entries",
2124
"static_nat.not_found": "Static NAT mapping not found (ID: {0})",
2225
"dst_ip_rule.not_found": "Destination IP rule not found (ID: {0})",
2326
"enrolled_device.invalid": "Invalid enrolled device data: {0}",

‎landscape-webui/src/i18n/en/geo_editor.ts‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,9 @@ export default {
44
source_url: "Download URL",
55
source_url_mode: "Download from URL",
66
source_direct_mode: "Define Directly",
7+
source_format: "File Format",
8+
txt_key: "TXT Key",
9+
txt_key_placeholder: "Defaults to DEFAULT",
710
name_unique: "Name (must be unique for config distinction)",
811
name_required: "Name is required",
912
name_invalid:
@@ -35,8 +38,10 @@ export default {
3538
},
3639
item_card: {
3740
source_type: "Source Type",
41+
source_format: "File Format",
3842
next_update_time: "Next Update Time",
3943
key_count: "Key Count",
44+
txt_key: "TXT Key",
4045
update_with_file: "Update Using File",
4146
},
4247
detail_drawer: {

‎landscape-webui/src/i18n/zh/api_errors.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,9 @@ export default {
1818
"geo_ip.cache_not_found": "找不到 GeoIP 缓存 (key: {0})",
1919
"geo_ip.file_not_found": "上传中未找到 GeoIP 文件",
2020
"geo_ip.file_read_error": "GeoIP 文件读取错误",
21+
"geo_ip.config_not_found": "找不到 GeoIP 配置 ({0})",
22+
"geo_ip.dat_decode_error": "GeoIP DAT 文件解析错误",
23+
"geo_ip.no_valid_cidr": "GeoIP TXT 文件中没有合法的 CIDR 条目",
2124
"static_nat.not_found": "找不到静态 NAT 映射 (ID: {0})",
2225
"dst_ip_rule.not_found": "找不到目标 IP 规则 (ID: {0})",
2326
"enrolled_device.invalid": "设备数据无效: {0}",

‎landscape-webui/src/i18n/zh/geo_editor.ts‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,9 @@ export default {
44
source_url: "下载 URL",
55
source_url_mode: "URL 下载",
66
source_direct_mode: "直接定义",
7+
source_format: "文件格式",
8+
txt_key: "TXT Key",
9+
txt_key_placeholder: "留空则使用 DEFAULT",
710
name_unique: "名称 (与其他配置区分， 需要唯一)",
811
name_required: "名称不能为空",
912
name_invalid: "名称只能包含字母、数字、点、下划线和中划线",
@@ -34,8 +37,10 @@ export default {
3437
},
3538
item_card: {
3639
source_type: "来源类型",
40+
source_format: "文件格式",
3741
next_update_time: "下次更新时间",
3842
key_count: "Key 数量",
43+
txt_key: "TXT Key",
3944
update_with_file: "使用文件更新",
4045
},
4146
detail_drawer: {

0 commit comments

Comments
 (0)