Skip to content

Commit a651ebf

Browse files
committed
fix: consolidated numeric overflow/underflow prevention with saturating operations
This PR consolidates the following numeric safety fixes: - #39: Use saturating casts in git_info to prevent overflow - #40: Use saturating casts for token counts in streaming - #41: Use saturating subtraction to prevent underflow in compaction All changes use saturating arithmetic operations: - Replaced direct casts with saturating_sub and try_into - Prevents panic on numeric overflow/underflow conditions
1 parent c398212 commit a651ebf

File tree

3 files changed

+32
-8
lines changed

3 files changed

+32
-8
lines changed

src/cortex-compact/src/compactor.rs

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -106,7 +106,10 @@ impl Compactor {
106106
}];
107107
new_items.extend(items.into_iter().skip(preserved_start));
108108

109-
let tokens_after = current_tokens - tokens_in_compacted + summary_tokens;
109+
// Use saturating arithmetic to prevent underflow if tokens_in_compacted > current_tokens
110+
let tokens_after = current_tokens
111+
.saturating_sub(tokens_in_compacted)
112+
.saturating_add(summary_tokens);
110113

111114
let result =
112115
CompactionResult::success(summary, current_tokens, tokens_after, items_removed);

src/cortex-engine/src/git_info.rs

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -67,9 +67,17 @@ impl GitInfo {
6767
let status = git_command(&root, &["status", "--porcelain"]).unwrap_or_default();
6868
let is_dirty = !status.is_empty();
6969

70-
// Count changes
71-
let changes = status.lines().filter(|l| !l.starts_with("??")).count() as u32;
72-
let untracked = status.lines().filter(|l| l.starts_with("??")).count() as u32;
70+
// Count changes (saturating to u32::MAX to prevent truncation)
71+
let changes = status
72+
.lines()
73+
.filter(|l| !l.starts_with("??"))
74+
.count()
75+
.min(u32::MAX as usize) as u32;
76+
let untracked = status
77+
.lines()
78+
.filter(|l| l.starts_with("??"))
79+
.count()
80+
.min(u32::MAX as usize) as u32;
7381

7482
// Get tags
7583
let tags = git_command(&root, &["tag", "--points-at", "HEAD"])
@@ -385,7 +393,7 @@ impl GitStash {
385393
let message = parts.get(2).unwrap_or(&"").to_string();
386394

387395
stashes.push(GitStash {
388-
index: i as u32,
396+
index: i.min(u32::MAX as usize) as u32,
389397
message,
390398
branch,
391399
});

src/cortex-engine/src/streaming.rs

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -26,12 +26,25 @@ pub struct StreamTokenUsage {
2626
pub total_tokens: u32,
2727
}
2828

29+
/// Safely convert an i64 token count to u32 with saturation.
30+
/// Negative values clamp to 0, values > u32::MAX clamp to u32::MAX.
31+
#[inline]
32+
fn saturating_i64_to_u32(value: i64) -> u32 {
33+
if value <= 0 {
34+
0
35+
} else if value > u32::MAX as i64 {
36+
u32::MAX
37+
} else {
38+
value as u32
39+
}
40+
}
41+
2942
impl From<crate::client::TokenUsage> for StreamTokenUsage {
3043
fn from(usage: crate::client::TokenUsage) -> Self {
3144
Self {
32-
prompt_tokens: usage.input_tokens as u32,
33-
completion_tokens: usage.output_tokens as u32,
34-
total_tokens: usage.total_tokens as u32,
45+
prompt_tokens: saturating_i64_to_u32(usage.input_tokens),
46+
completion_tokens: saturating_i64_to_u32(usage.output_tokens),
47+
total_tokens: saturating_i64_to_u32(usage.total_tokens),
3548
}
3649
}
3750
}

0 commit comments

Comments
 (0)