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)]