Skip to content

Commit 79957cc

Browse files
antoniupopantoniupop
committed
fix(coprocessor): improve error handling in edge cases
1 parent 73e9e30 commit 79957cc

File tree

2 files changed

+69
-77
lines changed

2 files changed

+69
-77
lines changed

‎coprocessor/fhevm-engine/fhevm-engine-common/src/gpu_memory.rs‎

Lines changed: 4 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -53,10 +53,6 @@ impl SupportedFheCiphertexts {
5353
}
5454
}
5555

56-
pub fn get_supported_ct_size_on_gpu(ct_type: i16) -> u64 {
57-
trivial_encrypt_be_bytes(ct_type, &[1u8]).get_size_on_gpu()
58-
}
59-
6056
// Reserving GPU memory happens in two stages:
6157
// - we add the amount we need atomically to the GPU's reservation pool
6258
// - we check that the new pool fits on GPU
@@ -1724,11 +1720,10 @@ pub fn get_op_size_on_gpu(
17241720
}
17251721
SupportedFheOperations::FheTrivialEncrypt | SupportedFheOperations::FheCast => {
17261722
match (&input_operands[0], &input_operands[1]) {
1727-
(_, SupportedFheCiphertexts::Scalar(op)) => Ok(trivial_encrypt_be_bytes(
1728-
to_be_u16_bit(op) as i16,
1729-
&[1u8],
1730-
)
1731-
.get_size_on_gpu()),
1723+
(_, SupportedFheCiphertexts::Scalar(op)) => {
1724+
Ok(trivial_encrypt_be_bytes(to_be_u16_bit(op) as i16, &[1u8])?
1725+
.get_size_on_gpu())
1726+
}
17321727
(_, _) => Err(FhevmError::UnsupportedFheTypes {
17331728
fhe_operation: format!("{:?}", fhe_operation),
17341729
input_types: input_operands.iter().map(|i| i.type_name()).collect(),

‎coprocessor/fhevm-engine/fhevm-engine-common/src/tfhe_ops.rs‎

Lines changed: 65 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -75,22 +75,25 @@ pub fn deserialize_fhe_ciphertext(
7575
}
7676

7777
/// Function assumes encryption key already set
78-
pub fn trivial_encrypt_be_bytes(output_type: i16, input_bytes: &[u8]) -> SupportedFheCiphertexts {
78+
pub fn trivial_encrypt_be_bytes(
79+
output_type: i16,
80+
input_bytes: &[u8],
81+
) -> Result<SupportedFheCiphertexts, FhevmError> {
7982
let last_byte = if !input_bytes.is_empty() {
8083
input_bytes[input_bytes.len() - 1]
8184
} else {
8285
0
8386
};
8487
match output_type {
85-
0 => SupportedFheCiphertexts::FheBool(
88+
0 => Ok(SupportedFheCiphertexts::FheBool(
8689
FheBool::try_encrypt_trivial(last_byte > 0).expect("trivial encrypt bool"),
87-
),
88-
1 => SupportedFheCiphertexts::FheUint4(
90+
)),
91+
1 => Ok(SupportedFheCiphertexts::FheUint4(
8992
FheUint4::try_encrypt_trivial(last_byte).expect("trivial encrypt 4"),
90-
),
91-
2 => SupportedFheCiphertexts::FheUint8(
93+
)),
94+
2 => Ok(SupportedFheCiphertexts::FheUint8(
9295
FheUint8::try_encrypt_trivial(last_byte).expect("trivial encrypt 8"),
93-
),
96+
)),
9497
3 => {
9598
let mut padded: [u8; 2] = [0; 2];
9699
if !input_bytes.is_empty() {
@@ -105,9 +108,9 @@ pub fn trivial_encrypt_be_bytes(output_type: i16, input_bytes: &[u8]) -> Support
105108
.copy_from_slice(&input_bytes[input_bytes.len() - len..]);
106109
}
107110
let res = u16::from_be_bytes(padded);
108-
SupportedFheCiphertexts::FheUint16(
111+
Ok(SupportedFheCiphertexts::FheUint16(
109112
FheUint16::try_encrypt_trivial(res).expect("trivial encrypt 16"),
110-
)
113+
))
111114
}
112115
4 => {
113116
let mut padded: [u8; 4] = [0; 4];
@@ -123,9 +126,9 @@ pub fn trivial_encrypt_be_bytes(output_type: i16, input_bytes: &[u8]) -> Support
123126
.copy_from_slice(&input_bytes[input_bytes.len() - len..]);
124127
}
125128
let res: u32 = u32::from_be_bytes(padded);
126-
SupportedFheCiphertexts::FheUint32(
129+
Ok(SupportedFheCiphertexts::FheUint32(
127130
FheUint32::try_encrypt_trivial(res).expect("trivial encrypt 32"),
128-
)
131+
))
129132
}
130133
5 => {
131134
let mut padded: [u8; 8] = [0; 8];
@@ -141,9 +144,9 @@ pub fn trivial_encrypt_be_bytes(output_type: i16, input_bytes: &[u8]) -> Support
141144
.copy_from_slice(&input_bytes[input_bytes.len() - len..]);
142145
}
143146
let res: u64 = u64::from_be_bytes(padded);
144-
SupportedFheCiphertexts::FheUint64(
147+
Ok(SupportedFheCiphertexts::FheUint64(
145148
FheUint64::try_encrypt_trivial(res).expect("trivial encrypt 64"),
146-
)
149+
))
147150
}
148151
6 => {
149152
let mut padded: [u8; 16] = [0; 16];
@@ -160,7 +163,7 @@ pub fn trivial_encrypt_be_bytes(output_type: i16, input_bytes: &[u8]) -> Support
160163
}
161164
let res: u128 = u128::from_be_bytes(padded);
162165
let output = FheUint128::try_encrypt_trivial(res).expect("trivial encrypt 128");
163-
SupportedFheCiphertexts::FheUint128(output)
166+
Ok(SupportedFheCiphertexts::FheUint128(output))
164167
}
165168
7 => {
166169
let mut padded: [u8; 32] = [0; 32];
@@ -180,7 +183,7 @@ pub fn trivial_encrypt_be_bytes(output_type: i16, input_bytes: &[u8]) -> Support
180183
let output: FheUint160 = FheUint256::try_encrypt_trivial(be)
181184
.expect("trivial encrypt 160")
182185
.cast_into();
183-
SupportedFheCiphertexts::FheUint160(output)
186+
Ok(SupportedFheCiphertexts::FheUint160(output))
184187
}
185188
8 => {
186189
let mut padded: [u8; 32] = [0; 32];
@@ -198,7 +201,7 @@ pub fn trivial_encrypt_be_bytes(output_type: i16, input_bytes: &[u8]) -> Support
198201
be.copy_from_be_byte_slice(&padded);
199202
}
200203
let output = FheUint256::try_encrypt_trivial(be).expect("trivial encrypt 256");
201-
SupportedFheCiphertexts::FheUint256(output)
204+
Ok(SupportedFheCiphertexts::FheUint256(output))
202205
}
203206
9 => {
204207
let mut padded: [u8; 64] = [0; 64];
@@ -216,7 +219,7 @@ pub fn trivial_encrypt_be_bytes(output_type: i16, input_bytes: &[u8]) -> Support
216219
be.copy_from_be_byte_slice(&padded);
217220
}
218221
let output = FheUint512::try_encrypt_trivial(be).expect("trivial encrypt 512");
219-
SupportedFheCiphertexts::FheBytes64(output)
222+
Ok(SupportedFheCiphertexts::FheBytes64(output))
220223
}
221224
10 => {
222225
let mut padded: [u8; 128] = [0; 128];
@@ -234,7 +237,7 @@ pub fn trivial_encrypt_be_bytes(output_type: i16, input_bytes: &[u8]) -> Support
234237
be.copy_from_be_byte_slice(&padded);
235238
}
236239
let output = FheUint1024::try_encrypt_trivial(be).expect("trivial encrypt 1024");
237-
SupportedFheCiphertexts::FheBytes128(output)
240+
Ok(SupportedFheCiphertexts::FheBytes128(output))
238241
}
239242
11 => {
240243
let mut padded: [u8; 256] = [0; 256];
@@ -252,11 +255,9 @@ pub fn trivial_encrypt_be_bytes(output_type: i16, input_bytes: &[u8]) -> Support
252255
be.copy_from_be_byte_slice(&padded);
253256
}
254257
let output = FheUint2048::try_encrypt_trivial(be).expect("trivial encrypt 2048");
255-
SupportedFheCiphertexts::FheBytes256(output)
256-
}
257-
other => {
258-
panic!("Unknown input type for trivial encryption: {other}")
258+
Ok(SupportedFheCiphertexts::FheBytes256(output))
259259
}
260+
other => Err(FhevmError::UnknownFheType(other as i32)),
260261
}
261262
}
262263

@@ -698,9 +699,10 @@ pub fn check_fhe_operand_types(
698699

699700
Ok(())
700701
}
701-
other => {
702-
panic!("Unexpected branch: {:?}", other)
703-
}
702+
other => Err(FhevmError::UnsupportedFheTypes {
703+
fhe_operation: format!("Unexpected op_type branch: {:?}", other),
704+
input_types: vec![],
705+
}),
704706
}
705707
}
706708
}
@@ -3058,7 +3060,10 @@ pub fn perform_fhe_operation_impl(
30583060
let out: tfhe::FheUint1024 = inp.clone().cast_into();
30593061
Ok(SupportedFheCiphertexts::FheBytes128(out))
30603062
}
3061-
other => panic!("unexpected type: {other}"),
3063+
other => Err(FhevmError::UnknownCastType {
3064+
fhe_operation: format!("{:?}", fhe_operation),
3065+
type_to_cast_to: other,
3066+
}),
30623067
}
30633068
}
30643069
}
@@ -3070,7 +3075,7 @@ pub fn perform_fhe_operation_impl(
30703075
SupportedFheOperations::FheTrivialEncrypt => match (&input_operands[0], &input_operands[1])
30713076
{
30723077
(SupportedFheCiphertexts::Scalar(inp), SupportedFheCiphertexts::Scalar(op)) => {
3073-
Ok(trivial_encrypt_be_bytes(to_be_u16_bit(op) as i16, inp))
3078+
trivial_encrypt_be_bytes(to_be_u16_bit(op) as i16, inp)
30743079
}
30753080
_ => Err(FhevmError::UnsupportedFheTypes {
30763081
fhe_operation: format!("{:?}", fhe_operation),
@@ -3092,7 +3097,7 @@ pub fn perform_fhe_operation_impl(
30923097
};
30933098
let rand_seed = to_be_u128_bit(rand_counter);
30943099
let to_type = to_be_u16_bit(to_type) as i16;
3095-
Ok(generate_random_number(to_type as i16, rand_seed, None))
3100+
generate_random_number(to_type as i16, rand_seed, None)
30963101
}
30973102
SupportedFheOperations::FheRandBounded => {
30983103
let SupportedFheCiphertexts::Scalar(rand_counter) = &input_operands[0] else {
@@ -3115,13 +3120,12 @@ pub fn perform_fhe_operation_impl(
31153120
};
31163121
let rand_seed = to_be_u128_bit(rand_counter);
31173122
let to_type = to_be_u16_bit(to_type) as i16;
3118-
Ok(generate_random_number(
3119-
to_type as i16,
3120-
rand_seed,
3121-
Some(upper_bound),
3122-
))
3123+
generate_random_number(to_type as i16, rand_seed, Some(upper_bound))
31233124
}
3124-
SupportedFheOperations::FheGetInputCiphertext => todo!("Implement FheGetInputCiphertext"),
3125+
SupportedFheOperations::FheGetInputCiphertext => Err(FhevmError::UnsupportedFheTypes {
3126+
fhe_operation: format!("{:?}", fhe_operation),
3127+
input_types: input_operands.iter().map(|i| i.type_name()).collect(),
3128+
}),
31253129
}
31263130
}
31273131

@@ -3288,20 +3292,19 @@ pub fn generate_random_number(
32883292
the_type: i16,
32893293
seed: u128,
32903294
upper_bound: Option<&[u8]>,
3291-
) -> SupportedFheCiphertexts {
3295+
) -> Result<SupportedFheCiphertexts, FhevmError> {
32923296
match the_type {
3293-
0 => {
3294-
SupportedFheCiphertexts::FheBool(FheBool::generate_oblivious_pseudo_random(Seed(seed)))
3295-
}
3297+
0 => Ok(SupportedFheCiphertexts::FheBool(
3298+
FheBool::generate_oblivious_pseudo_random(Seed(seed)),
3299+
)),
32963300
1 => {
32973301
let bit_count = 4;
32983302
let random_bits = upper_bound
32993303
.map(be_number_random_bits)
33003304
.unwrap_or(bit_count)
33013305
.min(bit_count) as u64;
3302-
SupportedFheCiphertexts::FheUint4(FheUint4::generate_oblivious_pseudo_random_bounded(
3303-
Seed(seed),
3304-
random_bits,
3306+
Ok(SupportedFheCiphertexts::FheUint4(
3307+
FheUint4::generate_oblivious_pseudo_random_bounded(Seed(seed), random_bits),
33053308
))
33063309
}
33073310
2 => {
@@ -3310,9 +3313,8 @@ pub fn generate_random_number(
33103313
.map(be_number_random_bits)
33113314
.unwrap_or(bit_count)
33123315
.min(bit_count) as u64;
3313-
SupportedFheCiphertexts::FheUint8(FheUint8::generate_oblivious_pseudo_random_bounded(
3314-
Seed(seed),
3315-
random_bits,
3316+
Ok(SupportedFheCiphertexts::FheUint8(
3317+
FheUint8::generate_oblivious_pseudo_random_bounded(Seed(seed), random_bits),
33163318
))
33173319
}
33183320
3 => {
@@ -3321,9 +3323,8 @@ pub fn generate_random_number(
33213323
.map(be_number_random_bits)
33223324
.unwrap_or(bit_count)
33233325
.min(bit_count) as u64;
3324-
SupportedFheCiphertexts::FheUint16(FheUint16::generate_oblivious_pseudo_random_bounded(
3325-
Seed(seed),
3326-
random_bits,
3326+
Ok(SupportedFheCiphertexts::FheUint16(
3327+
FheUint16::generate_oblivious_pseudo_random_bounded(Seed(seed), random_bits),
33273328
))
33283329
}
33293330
4 => {
@@ -3332,9 +3333,8 @@ pub fn generate_random_number(
33323333
.map(be_number_random_bits)
33333334
.unwrap_or(bit_count)
33343335
.min(bit_count) as u64;
3335-
SupportedFheCiphertexts::FheUint32(FheUint32::generate_oblivious_pseudo_random_bounded(
3336-
Seed(seed),
3337-
random_bits,
3336+
Ok(SupportedFheCiphertexts::FheUint32(
3337+
FheUint32::generate_oblivious_pseudo_random_bounded(Seed(seed), random_bits),
33383338
))
33393339
}
33403340
5 => {
@@ -3343,9 +3343,8 @@ pub fn generate_random_number(
33433343
.map(be_number_random_bits)
33443344
.unwrap_or(bit_count)
33453345
.min(bit_count) as u64;
3346-
SupportedFheCiphertexts::FheUint64(FheUint64::generate_oblivious_pseudo_random_bounded(
3347-
Seed(seed),
3348-
random_bits,
3346+
Ok(SupportedFheCiphertexts::FheUint64(
3347+
FheUint64::generate_oblivious_pseudo_random_bounded(Seed(seed), random_bits),
33493348
))
33503349
}
33513350
6 => {
@@ -3354,62 +3353,60 @@ pub fn generate_random_number(
33543353
.map(be_number_random_bits)
33553354
.unwrap_or(bit_count)
33563355
.min(bit_count) as u64;
3357-
SupportedFheCiphertexts::FheUint128(
3356+
Ok(SupportedFheCiphertexts::FheUint128(
33583357
FheUint128::generate_oblivious_pseudo_random_bounded(Seed(seed), random_bits),
3359-
)
3358+
))
33603359
}
33613360
7 => {
33623361
let bit_count = 160;
33633362
let random_bits = upper_bound
33643363
.map(be_number_random_bits)
33653364
.unwrap_or(bit_count)
33663365
.min(bit_count) as u64;
3367-
SupportedFheCiphertexts::FheUint160(
3366+
Ok(SupportedFheCiphertexts::FheUint160(
33683367
FheUint160::generate_oblivious_pseudo_random_bounded(Seed(seed), random_bits),
3369-
)
3368+
))
33703369
}
33713370
8 => {
33723371
let bit_count = 256;
33733372
let random_bits = upper_bound
33743373
.map(be_number_random_bits)
33753374
.unwrap_or(bit_count)
33763375
.min(bit_count) as u64;
3377-
SupportedFheCiphertexts::FheUint256(
3376+
Ok(SupportedFheCiphertexts::FheUint256(
33783377
FheUint256::generate_oblivious_pseudo_random_bounded(Seed(seed), random_bits),
3379-
)
3378+
))
33803379
}
33813380
9 => {
33823381
let bit_count = 512;
33833382
let random_bits = upper_bound
33843383
.map(be_number_random_bits)
33853384
.unwrap_or(bit_count)
33863385
.min(bit_count) as u64;
3387-
SupportedFheCiphertexts::FheBytes64(
3386+
Ok(SupportedFheCiphertexts::FheBytes64(
33883387
FheUint512::generate_oblivious_pseudo_random_bounded(Seed(seed), random_bits),
3389-
)
3388+
))
33903389
}
33913390
10 => {
33923391
let bit_count = 1024;
33933392
let random_bits = upper_bound
33943393
.map(be_number_random_bits)
33953394
.unwrap_or(bit_count)
33963395
.min(bit_count) as u64;
3397-
SupportedFheCiphertexts::FheBytes128(
3396+
Ok(SupportedFheCiphertexts::FheBytes128(
33983397
FheUint1024::generate_oblivious_pseudo_random_bounded(Seed(seed), random_bits),
3399-
)
3398+
))
34003399
}
34013400
11 => {
34023401
let bit_count = 2048;
34033402
let random_bits = upper_bound
34043403
.map(be_number_random_bits)
34053404
.unwrap_or(bit_count)
34063405
.min(bit_count) as u64;
3407-
SupportedFheCiphertexts::FheBytes256(
3406+
Ok(SupportedFheCiphertexts::FheBytes256(
34083407
FheUint2048::generate_oblivious_pseudo_random_bounded(Seed(seed), random_bits),
3409-
)
3410-
}
3411-
other => {
3412-
panic!("unknown type to trim to: {other}")
3408+
))
34133409
}
3410+
other => Err(FhevmError::UnknownFheType(other as i32)),
34143411
}
34153412
}

0 commit comments

Comments
 (0)