gru/RustProxy
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

v0.28.0

Browse Source
This commit is contained in:
Mateusz Gruszczyński committed 2026-10-08 09:42:10 +02:00
1 parent c24bc4279b
commit 3b3f0db377
84 files changed
+2686 -392
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/geoip/mmdb.rs
+334
View File
@@ -0,0 +1,334 @@
use anyhow::{bail, ensure, Context, Result};
use std::{fs, net::IpAddr, path::Path};
const METADATA_MARKER: &[u8] = b"\xab\xcd\xefMaxMind.com";
const MAX_METADATA_BYTES: usize = 128 * 1024;
const MAX_VALUES: usize = 65_536;
const MAX_DEPTH: usize = 512;
const MAX_STRING_BYTES: usize = 4096;
#[derive(Debug)]
pub(super) struct Database {
bytes: Vec<u8>,
node_count: u32,
record_size: u16,
ip_version: u16,
search_tree_size: usize,
data_section_start: usize,
metadata_offset: usize,
database_type: String,
}
#[derive(Clone, Copy)]
struct Header {
kind: u8,
size: usize,
payload: usize,
next: usize,
pointer: Option<usize>,
}
#[derive(Default)]
struct Budget {
values: usize,
}
impl Budget {
fn charge(&mut self) -> Result<()> {
self.values = self.values.saturating_add(1);
ensure!(self.values <= MAX_VALUES, "MMDB value traversal limit exceeded");
Ok(())
}
}
impl Database {
pub(super) fn load(path: &Path, max_bytes: u64) -> Result<Self> {
let metadata = fs::symlink_metadata(path)
.with_context(|| format!("open GeoIP MMDB {}", path.display()))?;
ensure!(metadata.is_file() && !metadata.file_type().is_symlink(), "GeoIP MMDB must be a regular file");
ensure!(metadata.len() <= max_bytes, "GeoIP MMDB exceeds configured maximum size");
let bytes = fs::read(path).with_context(|| format!("read GeoIP MMDB {}", path.display()))?;
let marker = metadata_marker(&bytes).context("MMDB metadata marker not found")?;
let metadata_offset = marker + METADATA_MARKER.len();
let base = metadata_offset;
let node_count = u32::try_from(read_uint_path(&bytes, metadata_offset, base, &["node_count"])?).context("MMDB node_count is too large")?;
let record_size = u16::try_from(read_uint_path(&bytes, metadata_offset, base, &["record_size"])?).context("MMDB record_size is too large")?;
let ip_version = u16::try_from(read_uint_path(&bytes, metadata_offset, base, &["ip_version"])?).context("MMDB ip_version is too large")?;
let major = read_uint_path(&bytes, metadata_offset, base, &["binary_format_major_version"])?;
let database_type = read_string_path(&bytes, metadata_offset, base, &["database_type"])?;
ensure!(node_count > 0, "MMDB node_count must be positive");
ensure!(matches!(record_size, 24 | 28 | 32), "unsupported MMDB record_size {record_size}");
ensure!(matches!(ip_version, 4 | 6), "invalid MMDB ip_version {ip_version}");
ensure!(major == 2, "unsupported MMDB major format version {major}");
let node_bytes = usize::from(record_size)
.checked_mul(2)
.and_then(|bits| bits.checked_div(8))
.context("MMDB node size overflow")?;
let search_tree_size = (node_count as usize)
.checked_mul(node_bytes)
.context("MMDB search tree size overflow")?;
let data_section_start = search_tree_size.checked_add(16).context("MMDB data section overflow")?;
ensure!(data_section_start <= marker, "MMDB search tree overlaps metadata");
ensure!(bytes.get(search_tree_size..data_section_start) == Some(&[0u8; 16][..]), "MMDB data-section separator is invalid");
Ok(Self {
bytes,
node_count,
record_size,
ip_version,
search_tree_size,
data_section_start,
metadata_offset,
database_type,
})
}
pub(super) fn database_type(&self) -> &str {
&self.database_type
}
pub(super) fn country(&self, ip: IpAddr) -> Option<String> {
let offset = self.lookup_offset(ip).ok().flatten()?;
for path in [
&["country", "iso_code"][..],
&["registered_country", "iso_code"][..],
&["country_code"][..],
] {
if let Ok(code) = read_string_path(&self.bytes, offset, self.data_section_start, path) {
if let Some(code) = super::normalize_country_code(&code) {
return Some(String::from_utf8_lossy(&code).into_owned());
}
}
}
let mut budget = Budget::default();
if let Ok(code) = read_string(&self.bytes, offset, self.data_section_start, 0, &mut budget) {
if let Some(code) = super::normalize_country_code(&code) {
return Some(String::from_utf8_lossy(&code).into_owned());
}
}
None
}
fn lookup_offset(&self, ip: IpAddr) -> Result<Option<usize>> {
let (bytes, bits): ([u8; 16], usize) = match (self.ip_version, ip) {
(4, IpAddr::V4(value)) => {
let mut out = [0u8; 16];
out[..4].copy_from_slice(&value.octets());
(out, 32)
}
(4, IpAddr::V6(_)) => return Ok(None),
(6, IpAddr::V4(value)) => {
let mut out = [0u8; 16];
out[12..].copy_from_slice(&value.octets());
(out, 128)
}
(6, IpAddr::V6(value)) => (value.octets(), 128),
_ => bail!("invalid MMDB IP version"),
};
let mut node = 0u32;
for bit_index in 0..bits {
if node >= self.node_count {
break;
}
let bit = (bytes[bit_index / 8] >> (7 - (bit_index % 8))) & 1;
node = self.read_node(node, bit as usize)?;
}
if node == self.node_count {
return Ok(None);
}
ensure!(node >= self.node_count.saturating_add(16), "invalid MMDB data pointer");
let offset = (node as usize)
.checked_sub(self.node_count as usize)
.and_then(|value| value.checked_add(self.search_tree_size))
.context("MMDB data pointer overflow")?;
ensure!(offset >= self.data_section_start && offset < self.metadata_offset - METADATA_MARKER.len(), "MMDB data pointer is out of bounds");
Ok(Some(offset))
}
fn read_node(&self, node: u32, side: usize) -> Result<u32> {
ensure!((node as usize) < self.node_count as usize, "MMDB node is out of range");
let base = match self.record_size {
24 => node as usize * 6,
28 => node as usize * 7,
32 => node as usize * 8,
_ => unreachable!(),
};
let value = match self.record_size {
24 => read_be(&self.bytes, base + side * 3, 3)? as u32,
28 => {
let middle = *self.bytes.get(base + 3).context("MMDB node is truncated")?;
if side == 0 {
(((middle & 0xf0) as u32) << 20) | read_be(&self.bytes, base, 3)? as u32
} else {
(((middle & 0x0f) as u32) << 24) | read_be(&self.bytes, base + 4, 3)? as u32
}
}
32 => read_be(&self.bytes, base + side * 4, 4)? as u32,
_ => unreachable!(),
};
Ok(value)
}
}
fn metadata_marker(bytes: &[u8]) -> Option<usize> {
let start = bytes.len().saturating_sub(MAX_METADATA_BYTES + METADATA_MARKER.len());
bytes[start..]
.windows(METADATA_MARKER.len())
.rposition(|window| window == METADATA_MARKER)
.map(|offset| start + offset)
}
fn read_uint_path(bytes: &[u8], offset: usize, base: usize, path: &[&str]) -> Result<u64> {
let mut budget = Budget::default();
let value = find_path(bytes, offset, base, path, 0, &mut budget)?.context("MMDB metadata field missing")?;
read_uint(bytes, value, base, 0, &mut budget)
}
fn read_string_path(bytes: &[u8], offset: usize, base: usize, path: &[&str]) -> Result<String> {
let mut budget = Budget::default();
let value = find_path(bytes, offset, base, path, 0, &mut budget)?.context("MMDB string field missing")?;
read_string(bytes, value, base, 0, &mut budget)
}
fn find_path(bytes: &[u8], offset: usize, base: usize, path: &[&str], depth: usize, budget: &mut Budget) -> Result<Option<usize>> {
ensure!(depth <= MAX_DEPTH, "MMDB nesting limit exceeded");
budget.charge()?;
let header = header(bytes, offset, base)?;
if let Some(pointer) = header.pointer {
return find_path(bytes, pointer, base, path, depth + 1, budget);
}
if path.is_empty() {
return Ok(Some(offset));
}
ensure!(header.kind == 7, "MMDB selected value is not a map");
let mut cursor = header.payload;
for _ in 0..header.size {
let (key, next) = inline_string(bytes, cursor, base)?;
cursor = next;
let value_offset = cursor;
if key == path[0] {
if path.len() == 1 {
return Ok(Some(value_offset));
}
if let Some(found) = find_path(bytes, value_offset, base, &path[1..], depth + 1, budget)? {
return Ok(Some(found));
}
}
cursor = skip(bytes, value_offset, base, depth + 1, budget)?;
}
Ok(None)
}
fn read_string(bytes: &[u8], offset: usize, base: usize, depth: usize, budget: &mut Budget) -> Result<String> {
ensure!(depth <= MAX_DEPTH, "MMDB nesting limit exceeded");
budget.charge()?;
let h = header(bytes, offset, base)?;
if let Some(pointer) = h.pointer {
return read_string(bytes, pointer, base, depth + 1, budget);
}
ensure!(h.kind == 2, "MMDB value is not a string");
ensure!(h.size <= MAX_STRING_BYTES, "MMDB string exceeds the decode limit");
let data = bytes.get(h.payload..h.next).context("MMDB string is truncated")?;
Ok(std::str::from_utf8(data).context("MMDB string is not UTF-8")?.to_owned())
}
fn read_uint(bytes: &[u8], offset: usize, base: usize, depth: usize, budget: &mut Budget) -> Result<u64> {
ensure!(depth <= MAX_DEPTH, "MMDB nesting limit exceeded");
budget.charge()?;
let h = header(bytes, offset, base)?;
if let Some(pointer) = h.pointer {
return read_uint(bytes, pointer, base, depth + 1, budget);
}
ensure!(matches!(h.kind, 5 | 6 | 9), "MMDB value is not an unsigned integer");
ensure!(h.size <= 8, "MMDB integer is too wide");
read_be(bytes, h.payload, h.size)
}
fn inline_string<'a>(bytes: &'a [u8], offset: usize, base: usize) -> Result<(&'a str, usize)> {
let h = header(bytes, offset, base)?;
ensure!(h.pointer.is_none() && h.kind == 2, "MMDB map key is not an inline string");
let value = bytes.get(h.payload..h.next).context("MMDB map key is truncated")?;
Ok((std::str::from_utf8(value).context("MMDB map key is not UTF-8")?, h.next))
}
fn skip(bytes: &[u8], offset: usize, base: usize, depth: usize, budget: &mut Budget) -> Result<usize> {
ensure!(depth <= MAX_DEPTH, "MMDB nesting limit exceeded");
budget.charge()?;
let h = header(bytes, offset, base)?;
if h.pointer.is_some() {
return Ok(h.next);
}
let mut cursor = h.payload;
match h.kind {
7 => {
for _ in 0..h.size {
cursor = skip(bytes, cursor, base, depth + 1, budget)?;
cursor = skip(bytes, cursor, base, depth + 1, budget)?;
}
Ok(cursor)
}
11 => {
for _ in 0..h.size {
cursor = skip(bytes, cursor, base, depth + 1, budget)?;
}
Ok(cursor)
}
14 => Ok(h.payload),
12 | 13 => bail!("unsupported deprecated MMDB data type {}", h.kind),
_ => Ok(h.next),
}
}
fn header(bytes: &[u8], offset: usize, base: usize) -> Result<Header> {
let control = *bytes.get(offset).context("MMDB value is truncated")?;
let primary = control >> 5;
let size_bits = (control & 0x1f) as usize;
if primary == 1 {
let class = (size_bits >> 3) & 0x03;
let low = size_bits & 0x07;
let count = class + 1;
let payload = offset + 1;
let raw = read_be(bytes, payload, count)? as usize;
let pointer = match class {
0 => (low << 8) | raw,
1 => ((low << 16) | raw).checked_add(2048).context("MMDB pointer overflow")?,
2 => ((low << 24) | raw).checked_add(526_336).context("MMDB pointer overflow")?,
3 => raw,
_ => unreachable!(),
};
let target = base.checked_add(pointer).context("MMDB pointer overflow")?;
ensure!(target < bytes.len(), "MMDB pointer is out of bounds");
return Ok(Header { kind: 1, size: 0, payload, next: payload + count, pointer: Some(target) });
}
let mut cursor = offset + 1;
let kind = if primary == 0 {
let extended = *bytes.get(cursor).context("MMDB extended type is truncated")?;
cursor += 1;
extended.checked_add(7).context("MMDB type overflow")?
} else {
primary
};
let size = match size_bits {
0..=28 => size_bits,
29 => 29 + *bytes.get(cursor).context("MMDB size is truncated")? as usize,
30 => 285 + read_be(bytes, cursor, 2)? as usize,
31 => 65_821 + read_be(bytes, cursor, 3)? as usize,
_ => unreachable!(),
};
cursor += match size_bits { 29 => 1, 30 => 2, 31 => 3, _ => 0 };
let next = match kind {
7 | 11 | 14 => cursor,
_ => cursor.checked_add(size).context("MMDB field size overflow")?,
};
ensure!(next <= bytes.len(), "MMDB field is out of bounds");
Ok(Header { kind, size, payload: cursor, next, pointer: None })
}
fn read_be(bytes: &[u8], offset: usize, count: usize) -> Result<u64> {
ensure!(count <= 8, "MMDB integer is too wide");
let end = offset.checked_add(count).context("MMDB read overflow")?;
let data = bytes.get(offset..end).context("MMDB value is truncated")?;
Ok(data.iter().fold(0u64, |value, byte| (value << 8) | u64::from(*byte)))
}
#[cfg(test)]
include!("mmdb_tests.rs");
src/geoip/mmdb_tests.rs
+92
View File
@@ -0,0 +1,92 @@
#[cfg(test)]
mod tests {
use super::*;
fn string(value: &str) -> Vec<u8> {
assert!(value.len() <= 28);
let mut out = vec![0x40 | value.len() as u8];
out.extend_from_slice(value.as_bytes());
out
}
fn uint16(value: u8) -> [u8; 2] {
[0xa1, value]
}
fn uint32(value: u8) -> [u8; 2] {
[0xc1, value]
}
fn fixture_with_record(record: Vec<u8>) -> Vec<u8> {
let mut bytes = vec![0, 0, 17, 0, 0, 17];
bytes.extend_from_slice(&[0; 16]);
bytes.extend(record);
bytes.extend_from_slice(METADATA_MARKER);
bytes.push(0xe5);
bytes.extend(string("node_count"));
bytes.extend(uint32(1));
bytes.extend(string("record_size"));
bytes.extend(uint16(24));
bytes.extend(string("ip_version"));
bytes.extend(uint16(4));
bytes.extend(string("binary_format_major_version"));
bytes.extend(uint16(2));
bytes.extend(string("database_type"));
bytes.extend(string("GeoLite2-Country"));
bytes
}
fn fixture(country: &str) -> Vec<u8> {
let mut record = vec![0xe1];
record.extend(string("country"));
record.push(0xe1);
record.extend(string("iso_code"));
record.extend(string(country));
fixture_with_record(record)
}
fn sapics_fixture(country: &str) -> Vec<u8> {
let mut record = vec![0xe1];
record.extend(string("country_code"));
record.extend(string(country));
fixture_with_record(record)
}
#[test]
fn reads_standard_country_iso_code_from_mmdb() {
let directory = tempfile::tempdir().unwrap();
let path = directory.path().join("country.mmdb");
std::fs::write(&path, fixture("PL")).unwrap();
let database = Database::load(&path, 1024 * 1024).unwrap();
assert_eq!(database.database_type(), "GeoLite2-Country");
assert_eq!(database.country("192.0.2.1".parse().unwrap()).as_deref(), Some("PL"));
assert!(database.country("2001:db8::1".parse().unwrap()).is_none());
}
#[test]
fn reads_sapics_country_code_from_mmdb() {
let directory = tempfile::tempdir().unwrap();
let path = directory.path().join("country.mmdb");
std::fs::write(&path, sapics_fixture("DE")).unwrap();
let database = Database::load(&path, 1024 * 1024).unwrap();
assert_eq!(database.country("198.51.100.8".parse().unwrap()).as_deref(), Some("DE"));
}
#[test]
fn reads_root_country_code_string_from_custom_mmdb() {
let directory = tempfile::tempdir().unwrap();
let path = directory.path().join("country.mmdb");
std::fs::write(&path, fixture_with_record(string("FR"))).unwrap();
let database = Database::load(&path, 1024 * 1024).unwrap();
assert_eq!(database.country("192.0.2.44".parse().unwrap()).as_deref(), Some("FR"));
}
#[test]
fn rejects_invalid_country_code_without_breaking_lookup() {
let directory = tempfile::tempdir().unwrap();
let path = directory.path().join("country.mmdb");
std::fs::write(&path, fixture("POL")).unwrap();
let database = Database::load(&path, 1024 * 1024).unwrap();
assert!(database.country("203.0.113.4".parse().unwrap()).is_none());
}
}
src/geoip/tests.rs
+14 -12
View File
@@ -1,19 +1,21 @@
mod tests {
use super::*;
#[test]
fn range_lookup_handles_gaps() {
let db = Database {
ipv4: vec![V4Range {
start: u32::from(Ipv4Addr::new(192, 0, 2, 0)),
end: u32::from(Ipv4Addr::new(192, 0, 2, 255)),
country: *b"PL",
}],
ipv6: vec![],
};
assert_eq!(
db.country("192.0.2.7".parse().unwrap()).as_deref(),
Some("PL")
);
let directory = tempfile::tempdir().unwrap();
std::fs::write(
directory.path().join("user-country-ipv4.csv"),
"192.0.2.0,192.0.2.255,PL\n",
)
.unwrap();
std::fs::write(
directory.path().join("user-country-ipv6.csv"),
"2001:db8::,2001:db8::ffff,PL\n",
)
.unwrap();
let db = Database::load_format(directory.path(), DatabaseFormat::CsvPair).unwrap();
assert_eq!(db.country("192.0.2.7".parse().unwrap()).as_deref(), Some("PL"));
assert!(db.country("198.51.100.1".parse().unwrap()).is_none());
}
}
src/geoip/update.rs
+442
View File
@@ -0,0 +1,442 @@
use anyhow::{ensure, Context, Result};
use fs2::FileExt;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::{
fs::{self, File, OpenOptions},
io::Read,
path::{Path, PathBuf},
sync::{Arc, Weak},
time::Duration,
};
use tokio_util::sync::CancellationToken;
use url::Url;
pub const DEFAULT_IPV4_URL: &str =
"https://github.com/sapics/ip-location-db/releases/download/latest/user-country-ipv4.csv";
pub const DEFAULT_IPV6_URL: &str =
"https://github.com/sapics/ip-location-db/releases/download/latest/user-country-ipv6.csv";
pub const DEFAULT_MMDB_URL: &str =
"https://github.com/sapics/ip-location-db/releases/download/latest/user-country.mmdb";
pub const DEFAULT_MAX_BYTES: u64 = 256 * 1024 * 1024;
pub const DEFAULT_INTERVAL_SECONDS: u64 = 24 * 60 * 60;
const LEGACY_RELEASE_URL: &str =
"https://github.com/sapics/ip-location-db/releases/download/latest/";
const LEGACY_CHECKSUM_URL: &str =
"https://github.com/sapics/ip-location-db/releases/download/checksum/";
const MAX_CONFIGURED_BYTES: u64 = DEFAULT_MAX_BYTES;
const MAX_CHECKSUM_BYTES: usize = 4096;
const RETRY_INTERVAL: Duration = Duration::from_secs(15 * 60);
const CONFIG_POLL_INTERVAL: Duration = Duration::from_secs(60);
const LOCK_FILE: &str = ".geoip-update.lock";
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(default, deny_unknown_fields)]
pub struct UpdateSettings {
pub enabled: bool,
pub format: super::DatabaseFormat,
pub ipv4_url: String,
pub ipv6_url: String,
pub mmdb_url: String,
pub ipv4_checksum_url: Option<String>,
pub ipv6_checksum_url: Option<String>,
pub mmdb_checksum_url: Option<String>,
pub max_bytes: u64,
pub interval_seconds: u64,
#[serde(skip_serializing_if = "Option::is_none")]
pub release_url: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub checksum_url: Option<String>,
}
impl Default for UpdateSettings {
fn default() -> Self {
Self {
enabled: true,
format: super::DatabaseFormat::CsvPair,
ipv4_url: DEFAULT_IPV4_URL.into(),
ipv6_url: DEFAULT_IPV6_URL.into(),
mmdb_url: DEFAULT_MMDB_URL.into(),
ipv4_checksum_url: None,
ipv6_checksum_url: None,
mmdb_checksum_url: None,
max_bytes: DEFAULT_MAX_BYTES,
interval_seconds: DEFAULT_INTERVAL_SECONDS,
release_url: None,
checksum_url: None,
}
}
}
impl UpdateSettings {
pub fn validate(&mut self) -> Result<()> {
self.migrate_legacy()?;
self.ipv4_url = self.ipv4_url.trim().to_owned();
self.ipv6_url = self.ipv6_url.trim().to_owned();
self.mmdb_url = self.mmdb_url.trim().to_owned();
match self.format {
super::DatabaseFormat::CsvPair => {
self.ipv4_url = validate_direct_url("GeoIP IPv4 URL", &self.ipv4_url)?;
self.ipv6_url = validate_direct_url("GeoIP IPv6 URL", &self.ipv6_url)?;
if !self.mmdb_url.is_empty() {
self.mmdb_url = validate_direct_url("GeoIP MMDB URL", &self.mmdb_url)?;
}
}
super::DatabaseFormat::Mmdb => {
if self.enabled || !self.mmdb_url.is_empty() {
self.mmdb_url = validate_direct_url("GeoIP MMDB URL", &self.mmdb_url)?;
}
if !self.ipv4_url.is_empty() {
self.ipv4_url = validate_direct_url("GeoIP IPv4 URL", &self.ipv4_url)?;
}
if !self.ipv6_url.is_empty() {
self.ipv6_url = validate_direct_url("GeoIP IPv6 URL", &self.ipv6_url)?;
}
}
}
self.ipv4_checksum_url = validate_optional_url("GeoIP IPv4 checksum URL", self.ipv4_checksum_url.take())?;
self.ipv6_checksum_url = validate_optional_url("GeoIP IPv6 checksum URL", self.ipv6_checksum_url.take())?;
self.mmdb_checksum_url = validate_optional_url("GeoIP MMDB checksum URL", self.mmdb_checksum_url.take())?;
ensure!(
(1024 * 1024..=MAX_CONFIGURED_BYTES).contains(&self.max_bytes),
"GeoIP maximum file size must be between 1 MiB and 256 MiB"
);
ensure!(
(60..=31 * 24 * 60 * 60).contains(&self.interval_seconds),
"GeoIP update interval must be between 60 seconds and 31 days"
);
Ok(())
}
fn migrate_legacy(&mut self) -> Result<()> {
if let Some(release) = self.release_url.take() {
let base = validate_legacy_base("GeoIP legacy release URL", &release)?;
if base != LEGACY_RELEASE_URL
&& self.ipv4_url == DEFAULT_IPV4_URL
&& self.ipv6_url == DEFAULT_IPV6_URL
{
self.ipv4_url = format!("{base}user-country-ipv4.csv");
self.ipv6_url = format!("{base}user-country-ipv6.csv");
}
}
if let Some(checksum) = self.checksum_url.take() {
let base = validate_legacy_base("GeoIP legacy checksum URL", &checksum)?;
if base != LEGACY_CHECKSUM_URL
&& self.ipv4_checksum_url.is_none()
&& self.ipv6_checksum_url.is_none()
{
self.ipv4_checksum_url = Some(format!("{base}user-country-ipv4.csv.sha256"));
self.ipv6_checksum_url = Some(format!("{base}user-country-ipv6.csv.sha256"));
}
}
Ok(())
}
}
#[derive(Clone)]
struct Asset {
name: &'static str,
url: String,
checksum_url: Option<String>,
}
fn assets(settings: &UpdateSettings) -> Vec<Asset> {
match settings.format {
super::DatabaseFormat::CsvPair => vec![
Asset {
name: "user-country-ipv4.csv",
url: settings.ipv4_url.clone(),
checksum_url: settings.ipv4_checksum_url.clone(),
},
Asset {
name: "user-country-ipv6.csv",
url: settings.ipv6_url.clone(),
checksum_url: settings.ipv6_checksum_url.clone(),
},
],
super::DatabaseFormat::Mmdb => vec![Asset {
name: "country.mmdb",
url: settings.mmdb_url.clone(),
checksum_url: settings.mmdb_checksum_url.clone(),
}],
}
}
fn validate_direct_url(label: &str, value: &str) -> Result<String> {
let parsed = Url::parse(value).with_context(|| format!("invalid {label}"))?;
ensure!(parsed.scheme() == "https", "{label} must use https://");
ensure!(parsed.host_str().is_some(), "{label} must contain a host");
ensure!(parsed.username().is_empty() && parsed.password().is_none(), "{label} must not contain URL credentials");
ensure!(parsed.fragment().is_none(), "{label} must not contain a fragment");
Ok(parsed.to_string())
}
fn validate_optional_url(label: &str, value: Option<String>) -> Result<Option<String>> {
let value = value.unwrap_or_default();
let value = value.trim();
if value.is_empty() {
Ok(None)
} else {
validate_direct_url(label, value).map(Some)
}
}
fn validate_legacy_base(label: &str, value: &str) -> Result<String> {
let mut parsed = Url::parse(value).with_context(|| format!("invalid {label}"))?;
ensure!(parsed.scheme() == "https" && parsed.host_str().is_some(), "{label} must be an HTTPS URL");
ensure!(parsed.username().is_empty() && parsed.password().is_none(), "{label} must not contain credentials");
ensure!(parsed.query().is_none() && parsed.fragment().is_none(), "{label} must not contain a query or fragment");
if !parsed.path().ends_with('/') {
parsed.set_path(&format!("{}/", parsed.path()));
}
Ok(parsed.to_string())
}
struct UpdateLock(File);
impl Drop for UpdateLock {
fn drop(&mut self) {
let _ = FileExt::unlock(&self.0);
}
}
fn try_update_lock(data_dir: &Path) -> Result<Option<UpdateLock>> {
let path = data_dir.join(LOCK_FILE);
let mut options = OpenOptions::new();
options.create(true).read(true).write(true);
#[cfg(unix)]
{
use std::os::unix::fs::OpenOptionsExt;
options.mode(0o600).custom_flags(libc::O_NOFOLLOW);
}
let file = options.open(&path).with_context(|| format!("open GeoIP update lock {}", path.display()))?;
match file.try_lock_exclusive() {
Ok(()) => Ok(Some(UpdateLock(file))),
Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => Ok(None),
Err(error) => Err(error).context("lock GeoIP updater"),
}
}
pub async fn poll_loop(app: Weak<crate::runtime::App>, shutdown: CancellationToken) {
let mut observed: Option<UpdateSettings> = None;
let mut next_attempt = tokio::time::Instant::now();
loop {
if shutdown.is_cancelled() { return; }
let Some(app) = app.upgrade() else { return; };
let settings = app.config.load().envelope.config.geoip_updates.clone();
if observed.as_ref() != Some(&settings) {
next_attempt = tokio::time::Instant::now();
observed = Some(settings.clone());
}
let now = tokio::time::Instant::now();
if settings.enabled && now >= next_attempt {
let interval = match update_once(&app, &settings).await {
Ok(UpdateOutcome::Updated(names)) => {
tracing::info!(files = %names.join(", "), "GeoIP database updated and hot-reloaded");
Duration::from_secs(settings.interval_seconds)
}
Ok(UpdateOutcome::Current) => Duration::from_secs(settings.interval_seconds),
Ok(UpdateOutcome::Busy) => RETRY_INTERVAL.min(Duration::from_secs(settings.interval_seconds)),
Err(error) => {
tracing::warn!(error = %format!("{error:#}"), "GeoIP update failed; proxy traffic is unaffected");
RETRY_INTERVAL.min(Duration::from_secs(settings.interval_seconds))
}
};
next_attempt = tokio::time::Instant::now() + interval;
}
let wait = if settings.enabled {
next_attempt.saturating_duration_since(tokio::time::Instant::now()).min(CONFIG_POLL_INTERVAL)
} else {
CONFIG_POLL_INTERVAL
};
drop(app);
tokio::select! {
_ = shutdown.cancelled() => return,
_ = tokio::time::sleep(wait) => {}
}
}
}
enum UpdateOutcome { Updated(Vec<String>), Current, Busy }
#[derive(Debug, Clone, Serialize)]
pub struct UpdateResult {
pub outcome: &'static str,
pub updated_files: Vec<String>,
pub database: super::View,
}
pub async fn update_now(app: &Arc<crate::runtime::App>) -> Result<UpdateResult> {
let settings = app.config.load().envelope.config.geoip_updates.clone();
let outcome = update_once(app, &settings).await?;
let (outcome, updated_files) = match outcome {
UpdateOutcome::Updated(files) => ("updated", files),
UpdateOutcome::Current => ("current", Vec::new()),
UpdateOutcome::Busy => ("busy", Vec::new()),
};
Ok(UpdateResult { outcome, updated_files, database: app.geoip.load().view() })
}
async fn update_once(app: &Arc<crate::runtime::App>, settings: &UpdateSettings) -> Result<UpdateOutcome> {
let data_dir = app.policy.data_dir.as_deref().context("GeoIP updater requires a data directory")?;
let Some(_lock) = try_update_lock(data_dir)? else { return Ok(UpdateOutcome::Busy); };
let directory = data_dir.join("geoip");
ensure_safe_directory(&directory)?;
if settings.format == super::DatabaseFormat::Mmdb && settings.mmdb_url.trim().is_empty() {
reload(app, directory, settings.format).await?;
return Ok(UpdateOutcome::Current);
}
let assets = assets(settings);
let mut expected = Vec::with_capacity(assets.len());
for asset in &assets {
expected.push(match asset.checksum_url.as_deref() {
Some(url) => Some(publisher_checksum(url).await?),
None => None,
});
}
let mut current = Vec::with_capacity(assets.len());
for asset in &assets {
current.push(file_digest(&directory.join(asset.name), settings.max_bytes)?);
}
if assets.iter().enumerate().all(|(index, _)| expected[index].as_ref().is_some_and(|digest| current[index].as_deref() == Some(digest))) {
reload_if_needed(app, directory, settings.format).await?;
return Ok(UpdateOutcome::Current);
}
let stage = data_dir.join(format!(".geoip-stage-{:016x}", rand::random::<u64>()));
private_directory(&stage)?;
let result = async {
for (index, asset) in assets.iter().enumerate() {
let target = stage.join(asset.name);
if expected[index].as_ref().is_some_and(|digest| current[index].as_deref() == Some(digest)) {
copy_private(&directory.join(asset.name), &target, settings.max_bytes)?;
} else {
super::update_http::download(&asset.url, &target, settings.max_bytes, expected[index].as_deref()).await?;
}
}
let validation_stage = stage.clone();
let format = settings.format;
tokio::task::spawn_blocking(move || super::validate_directory_format(&validation_stage, format))
.await.context("GeoIP validation worker failed")??;
let mut changed = Vec::new();
for (index, asset) in assets.iter().enumerate() {
let staged = file_digest(&stage.join(asset.name), settings.max_bytes)?.context("staged GeoIP file disappeared")?;
if current[index].as_deref() != Some(staged.as_str()) {
publish_private(&stage.join(asset.name), &directory.join(asset.name), settings.max_bytes)?;
changed.push(asset.name.to_string());
}
}
reload(app, directory.clone(), settings.format).await?;
Ok::<Vec<String>, anyhow::Error>(changed)
}.await;
let _ = fs::remove_dir_all(&stage);
let changed = result?;
if changed.is_empty() { Ok(UpdateOutcome::Current) } else { Ok(UpdateOutcome::Updated(changed)) }
}
async fn reload_if_needed(app: &Arc<crate::runtime::App>, directory: PathBuf, format: super::DatabaseFormat) -> Result<()> {
let expected = match format { super::DatabaseFormat::CsvPair => "csv_pair", super::DatabaseFormat::Mmdb => "mmdb" };
let view = app.geoip.load().view();
if !view.loaded || view.format != expected {
reload(app, directory, format).await?;
}
Ok(())
}
async fn reload(app: &Arc<crate::runtime::App>, directory: PathBuf, format: super::DatabaseFormat) -> Result<()> {
let database = tokio::task::spawn_blocking(move || super::Database::load_format(&directory, format))
.await.context("GeoIP reload worker failed")??;
publish_database(&app.geoip, database);
Ok(())
}
fn publish_database(slot: &arc_swap::ArcSwap<super::Database>, database: super::Database) {
slot.store(Arc::new(database));
}
async fn publisher_checksum(url: &str) -> Result<String> {
let payload = super::update_http::fetch_bytes(url, MAX_CHECKSUM_BYTES).await?;
let text = std::str::from_utf8(&payload).context("GeoIP checksum is not UTF-8")?;
text.split(|character: char| !character.is_ascii_hexdigit())
.find(|token| token.len() == 64)
.map(str::to_ascii_lowercase)
.context("GeoIP publisher checksum is invalid")
}
fn ensure_safe_directory(directory: &Path) -> Result<()> {
match fs::symlink_metadata(directory) {
Ok(metadata) => ensure!(metadata.is_dir() && !metadata.file_type().is_symlink(), "GeoIP directory must be a real directory"),
Err(error) if error.kind() == std::io::ErrorKind::NotFound => private_directory(directory)?,
Err(error) => return Err(error.into()),
}
Ok(())
}
fn private_directory(path: &Path) -> Result<()> {
fs::create_dir_all(path)?;
let metadata = fs::symlink_metadata(path)?;
ensure!(metadata.is_dir() && !metadata.file_type().is_symlink(), "GeoIP directory must not be a symlink");
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
fs::set_permissions(path, fs::Permissions::from_mode(0o700))?;
}
Ok(())
}
fn file_digest(path: &Path, max_bytes: u64) -> Result<Option<String>> {
let metadata = match fs::symlink_metadata(path) {
Ok(metadata) => metadata,
Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None),
Err(error) => return Err(error.into()),
};
if !metadata.is_file() || metadata.file_type().is_symlink() || metadata.len() > max_bytes { return Ok(None); }
let mut options = OpenOptions::new(); options.read(true);
#[cfg(unix)] { use std::os::unix::fs::OpenOptionsExt; options.custom_flags(libc::O_NOFOLLOW); }
let mut file = options.open(path)?;
let mut digest = Sha256::new();
let mut buffer = [0u8; 1024 * 1024];
loop {
let read = file.read(&mut buffer)?;
if read == 0 { break; }
digest.update(&buffer[..read]);
}
Ok(Some(format!("{:x}", digest.finalize())))
}
fn copy_private(source: &Path, target: &Path, max_bytes: u64) -> Result<()> {
let metadata = fs::symlink_metadata(source)?;
ensure!(metadata.is_file() && !metadata.file_type().is_symlink(), "GeoIP source must be a regular file");
ensure!(metadata.len() <= max_bytes, "GeoIP source exceeds configured maximum size");
let mut source_options = OpenOptions::new(); source_options.read(true);
let mut target_options = OpenOptions::new(); target_options.write(true).create_new(true);
#[cfg(unix)] {
use std::os::unix::fs::OpenOptionsExt;
source_options.custom_flags(libc::O_NOFOLLOW);
target_options.mode(0o600).custom_flags(libc::O_NOFOLLOW);
}
let source = source_options.open(source)?;
let mut input = source.take(max_bytes.saturating_add(1));
let mut output = target_options.open(target)?;
let copied = std::io::copy(&mut input, &mut output)?;
ensure!(copied <= max_bytes, "GeoIP source exceeds configured maximum size");
output.sync_all()?;
Ok(())
}
fn publish_private(source: &Path, target: &Path, max_bytes: u64) -> Result<()> {
let parent = target.parent().context("GeoIP target has no parent")?;
ensure_safe_directory(parent)?;
let temporary = parent.join(format!(".{}.tmp-{:016x}", target.file_name().unwrap_or_default().to_string_lossy(), rand::random::<u64>()));
if let Err(error) = copy_private(source, &temporary, max_bytes) { let _ = fs::remove_file(&temporary); return Err(error); }
if let Err(error) = fs::rename(&temporary, target) { let _ = fs::remove_file(&temporary); return Err(error.into()); }
#[cfg(unix)] {
use std::os::unix::fs::OpenOptionsExt;
let directory = OpenOptions::new().read(true).custom_flags(libc::O_DIRECTORY).open(parent)?;
directory.sync_all()?;
}
Ok(())
}
#[cfg(test)]
include!("update_tests.rs");
src/geoip/update_http.rs
+200
View File
@@ -0,0 +1,200 @@
use anyhow::{ensure, Context, Result};
use axum::body::Body;
use http::{header, Request, Uri};
use http_body_util::{BodyExt, Limited};
use hyper_rustls::HttpsConnectorBuilder;
use hyper_util::{
client::legacy::{connect::HttpConnector, Client},
rt::TokioExecutor,
};
use rustls::ClientConfig;
use sha2::{Digest, Sha256};
use std::{
fs::OpenOptions,
path::Path,
sync::Arc,
time::Duration,
};
use tokio::io::AsyncWriteExt;
use url::Url;
const REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
const MAX_REDIRECTS: usize = 5;
pub(super) async fn fetch_bytes(url: &str, max_bytes: usize) -> Result<Vec<u8>> {
let client = client()?;
let mut current = initial_url(url)?;
for redirects in 0..=MAX_REDIRECTS {
let response = request(&client, &current).await?;
if response.status().is_redirection() {
current = redirect(&current, &response, redirects)?;
continue;
}
ensure_success_and_length(&response, max_bytes as u64)?;
let body = tokio::time::timeout(
REQUEST_TIMEOUT,
Limited::new(response.into_body(), max_bytes).collect(),
)
.await
.context("GeoIP response body timed out")?
.map_err(|error| anyhow::anyhow!("GeoIP download is too large or incomplete: {error}"))?
.to_bytes();
return Ok(body.to_vec());
}
unreachable!("GeoIP redirect loop always returns or fails")
}
pub(super) async fn download(
url: &str,
target: &Path,
max_bytes: u64,
expected_digest: Option<&str>,
) -> Result<String> {
let client = client()?;
let mut current = initial_url(url)?;
for redirects in 0..=MAX_REDIRECTS {
let response = request(&client, &current).await?;
if response.status().is_redirection() {
current = redirect(&current, &response, redirects)?;
continue;
}
ensure_success_and_length(&response, max_bytes)?;
return stream_to_file(response.into_body(), target, max_bytes, expected_digest).await;
}
unreachable!("GeoIP redirect loop always returns or fails")
}
fn client() -> Result<Client<hyper_rustls::HttpsConnector<HttpConnector>, Body>> {
let roots = crate::update_check::load_root_store();
let tls = ClientConfig::builder_with_provider(Arc::new(
rustls::crypto::ring::default_provider(),
))
.with_safe_default_protocol_versions()?
.with_root_certificates(roots)
.with_no_client_auth();
let mut http = HttpConnector::new();
http.enforce_http(false);
http.set_connect_timeout(Some(REQUEST_TIMEOUT));
http.set_nodelay(true);
let connector = HttpsConnectorBuilder::new()
.with_tls_config(tls)
.https_only()
.enable_http1()
.wrap_connector(http);
Ok(Client::builder(TokioExecutor::new()).build(connector))
}
fn initial_url(url: &str) -> Result<Url> {
let parsed = Url::parse(url).context("invalid GeoIP URL")?;
ensure!(parsed.scheme() == "https", "GeoIP downloads require HTTPS");
Ok(parsed)
}
async fn request(
client: &Client<hyper_rustls::HttpsConnector<HttpConnector>, Body>,
current: &Url,
) -> Result<hyper::Response<hyper::body::Incoming>> {
let uri: Uri = current.as_str().parse().context("invalid GeoIP URI")?;
let request = Request::builder()
.method(http::Method::GET)
.uri(uri)
.header(header::ACCEPT, "application/octet-stream, text/plain;q=0.9")
.header(
header::USER_AGENT,
concat!("RustProxy/", env!("CARGO_PKG_VERSION"), " GeoIP-update"),
)
.body(Body::empty())?;
tokio::time::timeout(REQUEST_TIMEOUT, client.request(request))
.await
.context("GeoIP request timed out")?
.with_context(|| format!("GeoIP request failed for {current}"))
}
fn redirect(
current: &Url,
response: &hyper::Response<hyper::body::Incoming>,
redirects: usize,
) -> Result<Url> {
ensure!(redirects < MAX_REDIRECTS, "GeoIP redirect limit exceeded");
let location = response
.headers()
.get(header::LOCATION)
.context("GeoIP redirect is missing Location")?
.to_str()
.context("GeoIP redirect Location is not UTF-8")?;
let target = crate::update_check::redirect_target(current, location)?;
ensure!(target.scheme() == "https", "GeoIP redirect must remain HTTPS");
Ok(target)
}
fn ensure_success_and_length(
response: &hyper::Response<hyper::body::Incoming>,
max_bytes: u64,
) -> Result<()> {
ensure!(
response.status().is_success(),
"GeoIP request returned HTTP {}",
response.status()
);
if let Some(length) = response.headers().get(header::CONTENT_LENGTH) {
let length = length
.to_str()?
.parse::<u64>()
.context("invalid GeoIP Content-Length")?;
ensure!(
length <= max_bytes,
"GeoIP download is larger than the configured safety limit"
);
}
Ok(())
}
async fn stream_to_file(
mut body: hyper::body::Incoming,
target: &Path,
max_bytes: u64,
expected_digest: Option<&str>,
) -> Result<String> {
let mut options = OpenOptions::new();
options.write(true).create_new(true);
#[cfg(unix)]
{
use std::os::unix::fs::OpenOptionsExt;
options.mode(0o600).custom_flags(libc::O_NOFOLLOW);
}
let file = options
.open(target)
.with_context(|| format!("create GeoIP staging file {}", target.display()))?;
let mut output = tokio::fs::File::from_std(file);
let mut digest = Sha256::new();
let mut written = 0u64;
while let Some(frame) = tokio::time::timeout(REQUEST_TIMEOUT, body.frame())
.await
.context("GeoIP response body timed out")?
{
let frame = frame.context("read GeoIP response body")?;
if let Some(data) = frame.data_ref() {
written = written
.checked_add(data.len() as u64)
.context("GeoIP download size overflow")?;
ensure!(
written <= max_bytes,
"GeoIP download is larger than the configured safety limit"
);
digest.update(data);
output.write_all(data).await?;
}
}
output.flush().await?;
output.sync_all().await?;
let actual = format!("{:x}", digest.finalize());
if let Some(expected) = expected_digest {
ensure!(
actual == expected,
"GeoIP SHA-256 mismatch for {}",
target.file_name().unwrap_or_default().to_string_lossy()
);
}
Ok(actual)
}
src/geoip/update_tests.rs
+94
View File
@@ -0,0 +1,94 @@
mod tests {
use super::*;
#[test]
fn defaults_use_direct_csv_sources_without_mandatory_checksums() {
let settings = UpdateSettings::default();
assert!(settings.enabled);
assert_eq!(settings.format, crate::geoip::DatabaseFormat::CsvPair);
assert_eq!(settings.ipv4_url, DEFAULT_IPV4_URL);
assert_eq!(settings.ipv6_url, DEFAULT_IPV6_URL);
assert!(settings.ipv4_checksum_url.is_none());
assert!(settings.ipv6_checksum_url.is_none());
assert_eq!(settings.mmdb_url, DEFAULT_MMDB_URL);
assert!(settings.mmdb_checksum_url.is_none());
assert_eq!(settings.max_bytes, 256 * 1024 * 1024);
assert_eq!(settings.interval_seconds, 24 * 60 * 60);
}
#[test]
fn settings_require_https_and_safe_bounds() {
let mut settings = UpdateSettings::default();
settings.ipv4_url = "http://example.invalid/db.csv".into();
assert!(settings.validate().is_err());
let mut settings = UpdateSettings::default();
settings.interval_seconds = 59;
assert!(settings.validate().is_err());
let mut settings = UpdateSettings::default();
settings.max_bytes = DEFAULT_MAX_BYTES + 1;
assert!(settings.validate().is_err());
}
#[test]
fn local_mmdb_can_disable_downloader_and_leave_url_empty() {
let mut settings = UpdateSettings {
enabled: false,
format: crate::geoip::DatabaseFormat::Mmdb,
mmdb_url: String::new(),
..UpdateSettings::default()
};
settings.validate().unwrap();
assert!(settings.mmdb_url.is_empty());
}
#[test]
fn enabled_mmdb_requires_a_direct_url() {
let mut settings = UpdateSettings {
enabled: true,
format: crate::geoip::DatabaseFormat::Mmdb,
mmdb_url: String::new(),
..UpdateSettings::default()
};
assert!(settings.validate().is_err());
}
#[test]
fn legacy_default_checksum_configuration_migrates_to_checksum_optional_defaults() {
let mut settings = UpdateSettings {
release_url: Some(LEGACY_RELEASE_URL.into()),
checksum_url: Some(LEGACY_CHECKSUM_URL.into()),
..UpdateSettings::default()
};
settings.validate().unwrap();
assert_eq!(settings.ipv4_url, DEFAULT_IPV4_URL);
assert_eq!(settings.ipv6_url, DEFAULT_IPV6_URL);
assert!(settings.ipv4_checksum_url.is_none());
assert!(settings.ipv6_checksum_url.is_none());
assert!(settings.release_url.is_none());
assert!(settings.checksum_url.is_none());
}
#[test]
fn publishing_a_downloaded_database_hot_swaps_future_geoip_lookups() {
let directory = tempfile::tempdir().unwrap();
let ipv4 = directory.path().join("user-country-ipv4.csv");
let ipv6 = directory.path().join("user-country-ipv6.csv");
std::fs::write(&ipv4, "192.0.2.0,192.0.2.255,PL\n").unwrap();
std::fs::write(&ipv6, "2001:db8::,2001:db8::ffff,PL\n").unwrap();
let slot = arc_swap::ArcSwap::from_pointee(
crate::geoip::Database::load_format(directory.path(), crate::geoip::DatabaseFormat::CsvPair).unwrap(),
);
let previous = slot.load_full();
assert_eq!(slot.load().country("192.0.2.7".parse().unwrap()).as_deref(), Some("PL"));
std::fs::write(&ipv4, "192.0.2.0,192.0.2.255,DE\n").unwrap();
std::fs::write(&ipv6, "2001:db8::,2001:db8::ffff,DE\n").unwrap();
publish_database(
&slot,
crate::geoip::Database::load_format(directory.path(), crate::geoip::DatabaseFormat::CsvPair).unwrap(),
);
assert_eq!(slot.load().country("192.0.2.7".parse().unwrap()).as_deref(), Some("DE"));
assert_eq!(previous.country("192.0.2.7".parse().unwrap()).as_deref(), Some("PL"));
}
}