Skip to content

Commit 8552534

Browse files
committed
Use u32 for struct and member variables in IWP, saturating to u32::MAX
To prevent panics during addition if `usize` is `u64`, use `u32` member variables internally. TK use `u32::saturating_add` instead of basic addition. However, to use `u32::saturating_add()`, the variables need to be of type `u32`. Therefore, this commit transforms the internal types to `u32`. In so doing, add a `const` function `saturate_to_u32()` which saturates a usize to `u32`. Also, replace `compact_size::encoded_size()` with a new function `encoded_size()` which returns a `u32`, avoiding needless casts.
1 parent b0981fc commit 8552534

File tree

1 file changed

+27
-10
lines changed

1 file changed

+27
-10
lines changed

bitcoin/src/blockdata/transaction.rs

Lines changed: 27 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -932,8 +932,8 @@ pub const fn predict_weight_from_slices(
932932
/// associated constants/methods.
933933
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
934934
pub struct InputWeightPrediction {
935-
script_size: usize,
936-
witness_size: usize,
935+
script_size: u32,
936+
witness_size: u32,
937937
}
938938

939939
impl InputWeightPrediction {
@@ -996,6 +996,23 @@ impl InputWeightPrediction {
996996
/// [`InputWeightPrediction::new`].
997997
pub const P2TR_KEY_NON_DEFAULT_SIGHASH: Self = InputWeightPrediction::from_slice(0, &[65]);
998998

999+
const fn saturate_to_u32(x: usize) -> u32 {
1000+
if x > u32::MAX as usize {
1001+
u32::MAX
1002+
} else {
1003+
x as u32 //cast ok, condition prevents larger than u32::MAX.
1004+
}
1005+
}
1006+
1007+
const fn encoded_size(value: usize) -> u32 {
1008+
match value {
1009+
0..=0xFC => 1,
1010+
0xFD..=0xFFFF => 3,
1011+
0x10000..=0xFFFFFFFF => 5,
1012+
_ => 9,
1013+
}
1014+
}
1015+
9991016
/// Input weight prediction corresponding to spending of P2WPKH output using [signature
10001017
/// grinding].
10011018
///
@@ -1065,16 +1082,16 @@ impl InputWeightPrediction {
10651082
T::Item: Borrow<usize>,
10661083
{
10671084
let (count, total_size) = witness_element_lengths.into_iter().fold(
1068-
(0usize, 0),
1085+
(0usize, 0u32),
10691086
|(count, total_size), elem_len| {
10701087
let elem_len = *elem_len.borrow();
1071-
let elem_size = elem_len + compact_size::encoded_size(elem_len);
1088+
let elem_size = Self::saturate_to_u32(elem_len) + Self::encoded_size(elem_len);
10721089
(count + 1, total_size + elem_size)
10731090
},
10741091
);
1075-
let witness_size =
1076-
if count > 0 { total_size + compact_size::encoded_size(count) } else { 0 };
1077-
let script_size = input_script_len + compact_size::encoded_size(input_script_len);
1092+
let witness_size = if count > 0 { total_size + Self::encoded_size(count) } else { 0 };
1093+
let script_size =
1094+
Self::saturate_to_u32(input_script_len) + Self::encoded_size(input_script_len);
10781095

10791096
InputWeightPrediction { script_size, witness_size }
10801097
}
@@ -1090,17 +1107,17 @@ impl InputWeightPrediction {
10901107
// for loops not supported in const fn
10911108
while i < witness_element_lengths.len() {
10921109
let elem_len = witness_element_lengths[i];
1093-
let elem_size = elem_len + compact_size::encoded_size_const(elem_len as u64);
1110+
let elem_size = Self::saturate_to_u32(elem_len) + Self::encoded_size(elem_len);
10941111
total_size += elem_size;
10951112
i += 1;
10961113
}
10971114
let witness_size = if !witness_element_lengths.is_empty() {
1098-
total_size + compact_size::encoded_size_const(witness_element_lengths.len() as u64)
1115+
total_size + Self::encoded_size(witness_element_lengths.len())
10991116
} else {
11001117
0
11011118
};
11021119
let script_size =
1103-
input_script_len + compact_size::encoded_size_const(input_script_len as u64);
1120+
Self::saturate_to_u32(input_script_len) + Self::encoded_size(input_script_len);
11041121

11051122
InputWeightPrediction { script_size, witness_size }
11061123
}

0 commit comments

Comments
 (0)