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

0.34.13_cleaup_code

Browse Source
This commit is contained in:
Mateusz Gruszczyński committed 2026-10-26 14:14:41 +01:00
1 parent ec46a96a29
commit 5c2e50ec06
87 files changed
+4890 -1338
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/geoip/mmdb.rs
+174 -40
View File
@@ -36,7 +36,10 @@ struct Budget {
impl Budget {
fn charge(&mut self) -> Result<()> {
self.values = self.values.saturating_add(1);
ensure!(self.values <= MAX_VALUES, "MMDB value traversal limit exceeded");
ensure!(
self.values <= MAX_VALUES,
"MMDB value traversal limit exceeded"
);
Ok(())
}
}
@@ -45,20 +48,56 @@ 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()))?;
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 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!(
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)
@@ -67,9 +106,17 @@ impl Database {
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");
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,
@@ -100,7 +147,8 @@ impl Database {
}
}
let mut budget = Budget::default();
if let Ok(code) = read_string(&self.bytes, offset, self.data_section_start, 0, &mut budget) {
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());
}
@@ -135,17 +183,27 @@ impl Database {
if node == self.node_count {
return Ok(None);
}
ensure!(node >= self.node_count.saturating_add(16), "invalid MMDB data pointer");
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");
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");
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,
@@ -170,7 +228,9 @@ impl Database {
}
fn metadata_marker(bytes: &[u8]) -> Option<usize> {
let start = bytes.len().saturating_sub(MAX_METADATA_BYTES + METADATA_MARKER.len());
let start = bytes
.len()
.saturating_sub(MAX_METADATA_BYTES + METADATA_MARKER.len());
bytes[start..]
.windows(METADATA_MARKER.len())
.rposition(|window| window == METADATA_MARKER)
@@ -179,17 +239,26 @@ fn metadata_marker(bytes: &[u8]) -> Option<usize> {
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")?;
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")?;
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>> {
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)?;
@@ -209,7 +278,9 @@ fn find_path(bytes: &[u8], offset: usize, base: usize, path: &[&str], depth: usi
if path.len() == 1 {
return Ok(Some(value_offset));
}
if let Some(found) = find_path(bytes, value_offset, base, &path[1..], depth + 1, budget)? {
if let Some(found) =
find_path(bytes, value_offset, base, &path[1..], depth + 1, budget)?
{
return Ok(Some(found));
}
}
@@ -218,7 +289,13 @@ fn find_path(bytes: &[u8], offset: usize, base: usize, path: &[&str], depth: usi
Ok(None)
}
fn read_string(bytes: &[u8], offset: usize, base: usize, depth: usize, budget: &mut Budget) -> Result<String> {
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)?;
@@ -226,31 +303,61 @@ fn read_string(bytes: &[u8], offset: usize, base: usize, depth: usize, budget: &
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())
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> {
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!(
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))
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> {
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)?;
@@ -290,18 +397,30 @@ fn header(bytes: &[u8], offset: usize, base: usize) -> Result<Header> {
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")?,
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) });
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")?;
let extended = *bytes
.get(cursor)
.context("MMDB extended type is truncated")?;
cursor += 1;
extended.checked_add(7).context("MMDB type overflow")?
} else {
@@ -314,20 +433,35 @@ fn header(bytes: &[u8], offset: usize, base: usize) -> Result<Header> {
31 => 65_821 + read_be(bytes, cursor, 3)? as usize,
_ => unreachable!(),
};
cursor += match size_bits { 29 => 1, 30 => 2, 31 => 3, _ => 0 };
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")?,
_ => 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 })
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)))
Ok(data
.iter()
.fold(0u64, |value, byte| (value << 8) | u64::from(*byte)))
}
#[cfg(test)]
src/geoip/mmdb_tests.rs
+12 -3
View File
@@ -59,7 +59,10 @@ mod tests {
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_eq!(
database.country("192.0.2.1".parse().unwrap()).as_deref(),
Some("PL")
);
assert!(database.country("2001:db8::1".parse().unwrap()).is_none());
}
@@ -69,7 +72,10 @@ mod tests {
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"));
assert_eq!(
database.country("198.51.100.8".parse().unwrap()).as_deref(),
Some("DE")
);
}
#[test]
@@ -78,7 +84,10 @@ mod tests {
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"));
assert_eq!(
database.country("192.0.2.44".parse().unwrap()).as_deref(),
Some("FR")
);
}
#[test]
src/geoip/tests.rs
+28 -7
View File
@@ -6,14 +6,32 @@ mod tests {
assert_eq!(std::mem::size_of::<V6Range>(), 34);
assert_eq!(std::mem::size_of::<V4Range>(), 12);
let ranges = [
V6Range { start: 0u128.to_be_bytes(), end: 127u128.to_be_bytes(), country: *b"PL" },
V6Range { start: 256u128.to_be_bytes(), end: 65535u128.to_be_bytes(), country: *b"DE" },
V6Range { start: (1u128 << 120).to_be_bytes(), end: u128::MAX.to_be_bytes(), country: *b"US" },
V6Range {
start: 0u128.to_be_bytes(),
end: 127u128.to_be_bytes(),
country: *b"PL",
},
V6Range {
start: 256u128.to_be_bytes(),
end: 65535u128.to_be_bytes(),
country: *b"DE",
},
V6Range {
start: (1u128 << 120).to_be_bytes(),
end: u128::MAX.to_be_bytes(),
country: *b"US",
},
];
for (address, expected) in [
(0, Some(*b"PL")), (127, Some(*b"PL")), (128, None),
(255, None), (256, Some(*b"DE")), (65535, Some(*b"DE")),
(65536, None), (1u128 << 120, Some(*b"US")), (u128::MAX, Some(*b"US")),
(0, Some(*b"PL")),
(127, Some(*b"PL")),
(128, None),
(255, None),
(256, Some(*b"DE")),
(65535, Some(*b"DE")),
(65536, None),
(1u128 << 120, Some(*b"US")),
(u128::MAX, Some(*b"US")),
] {
assert_eq!(find_v6(&ranges, address.to_be_bytes()), expected);
}
@@ -33,7 +51,10 @@ mod tests {
)
.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_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
+178 -50
View File
@@ -94,9 +94,12 @@ impl UpdateSettings {
}
}
}
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())?;
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"
@@ -166,8 +169,14 @@ 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");
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())
}
@@ -183,9 +192,18 @@ fn validate_optional_url(label: &str, value: Option<String>) -> Result<Option<St
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");
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()));
}
@@ -208,7 +226,9 @@ fn try_update_lock(data_dir: &Path) -> Result<Option<UpdateLock>> {
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()))?;
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),
@@ -220,8 +240,12 @@ pub async fn poll_loop(app: Weak<crate::runtime::App>, shutdown: CancellationTok
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; };
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();
@@ -235,7 +259,9 @@ pub async fn poll_loop(app: Weak<crate::runtime::App>, shutdown: CancellationTok
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)),
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))
@@ -244,7 +270,9 @@ pub async fn poll_loop(app: Weak<crate::runtime::App>, shutdown: CancellationTok
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)
next_attempt
.saturating_duration_since(tokio::time::Instant::now())
.min(CONFIG_POLL_INTERVAL)
} else {
CONFIG_POLL_INTERVAL
};
@@ -256,7 +284,11 @@ pub async fn poll_loop(app: Weak<crate::runtime::App>, shutdown: CancellationTok
}
}
enum UpdateOutcome { Updated(Vec<String>), Current, Busy }
enum UpdateOutcome {
Updated(Vec<String>),
Current,
Busy,
}
#[derive(Debug, Clone, Serialize)]
pub struct UpdateResult {
@@ -273,12 +305,25 @@ pub async fn update_now(app: &Arc<crate::runtime::App>) -> Result<UpdateResult>
UpdateOutcome::Current => ("current", Vec::new()),
UpdateOutcome::Busy => ("busy", Vec::new()),
};
Ok(UpdateResult { outcome, updated_files, database: app.geoip.load().view() })
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); };
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() {
@@ -295,9 +340,16 @@ async fn update_once(app: &Arc<crate::runtime::App>, settings: &UpdateSettings)
}
let mut current = Vec::with_capacity(assets.len());
for asset in &assets {
current.push(file_digest(&directory.join(asset.name), settings.max_bytes)?);
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))) {
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);
}
@@ -307,35 +359,64 @@ async fn update_once(app: &Arc<crate::runtime::App>, settings: &UpdateSettings)
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)) {
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?;
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")??;
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")?;
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)?;
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;
}
.await;
let _ = fs::remove_dir_all(&stage);
let changed = result?;
if changed.is_empty() { Ok(UpdateOutcome::Current) } else { Ok(UpdateOutcome::Updated(changed)) }
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" };
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?;
@@ -343,9 +424,15 @@ async fn reload_if_needed(app: &Arc<crate::runtime::App>, directory: PathBuf, fo
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")??;
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(())
}
@@ -365,7 +452,10 @@ async fn publisher_checksum(url: &str) -> Result<String> {
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"),
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()),
}
@@ -375,7 +465,10 @@ fn ensure_safe_directory(directory: &Path) -> Result<()> {
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");
ensure!(
metadata.is_dir() && !metadata.file_type().is_symlink(),
"GeoIP directory must not be a symlink"
);
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
@@ -390,15 +483,24 @@ fn file_digest(path: &Path, max_bytes: u64) -> Result<Option<String>> {
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); }
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; }
if read == 0 {
break;
}
digest.update(&buffer[..read]);
}
Ok(Some(format!("{:x}", digest.finalize())))
@@ -406,11 +508,20 @@ fn file_digest(path: &Path, max_bytes: u64) -> Result<Option<String>> {
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)] {
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);
@@ -419,7 +530,10 @@ fn copy_private(source: &Path, target: &Path, max_bytes: u64) -> Result<()> {
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");
ensure!(
copied <= max_bytes,
"GeoIP source exceeds configured maximum size"
);
output.sync_all()?;
Ok(())
}
@@ -427,12 +541,26 @@ fn copy_private(source: &Path, target: &Path, max_bytes: u64) -> Result<()> {
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)] {
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)?;
let directory = OpenOptions::new()
.read(true)
.custom_flags(libc::O_DIRECTORY)
.open(parent)?;
directory.sync_all()?;
}
Ok(())
src/geoip/update_http.rs
+10 -13
View File
@@ -9,12 +9,7 @@ use hyper_util::{
};
use rustls::ClientConfig;
use sha2::{Digest, Sha256};
use std::{
fs::OpenOptions,
path::Path,
sync::Arc,
time::Duration,
};
use std::{fs::OpenOptions, path::Path, sync::Arc, time::Duration};
use tokio::io::AsyncWriteExt;
use url::Url;
@@ -66,12 +61,11 @@ pub(super) async fn download(
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 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));
@@ -123,7 +117,10 @@ fn redirect(
.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");
ensure!(
target.scheme() == "https",
"GeoIP redirect must remain HTTPS"
);
Ok(target)
}
src/geoip/update_tests.rs
+22 -5
View File
@@ -76,19 +76,36 @@ mod tests {
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(),
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"));
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(),
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"));
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")
);
}
}