Skip to content

Commit d04ee84

Browse files
eudelins-zamaeudelins-zama
committed
chore(kms-connector): adapt storage urls retrieval
1 parent e31ee8a commit d04ee84

File tree

12 files changed

+88
-117
lines changed

12 files changed

+88
-117
lines changed

‎kms-connector/Cargo.lock‎

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎kms-connector/Cargo.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ gw-listener.path = "crates/gw-listener"
1818
kms-worker.path = "crates/kms-worker"
1919
tx-sender.path = "crates/tx-sender"
2020
connector-utils.path = "crates/utils"
21-
fhevm_gateway_bindings = { git = "https://github.com/zama-ai/fhevm.git", rev = "a51230f0d138586b7b9e49fc1550b554117ad4b1", default-features = false }
21+
fhevm_gateway_bindings = { git = "https://github.com/zama-ai/fhevm.git", rev = "f9671e161fbff94159bd7355c618954df2ec00fb", default-features = false }
2222
kms-grpc = { git = "https://github.com/zama-ai/kms.git", tag = "v0.11.0-26", default-features = true }
2323
tfhe = "=1.3.2"
2424

‎kms-connector/connector-db/migrations/20250902144114_update_sns_ciphertext.sql‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1 +1,2 @@
11
ALTER TYPE sns_ciphertext_material DROP ATTRIBUTE coprocessor_tx_sender_addresses CASCADE;
2+
ALTER TYPE sns_ciphertext_material ADD ATTRIBUTE storage_urls TEXT[] CASCADE;

‎kms-connector/crates/gw-listener/src/core/event_publisher.rs‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,8 @@ impl DbEventPublisher {
6060
let sns_ciphertexts_db = request
6161
.snsCtMaterials
6262
.iter()
63-
.map(SnsCiphertextMaterialDbItem::from)
63+
.zip(request.storageUrls)
64+
.map(|(sns_ct, url)| SnsCiphertextMaterialDbItem::new(sns_ct, url))
6465
.collect::<Vec<SnsCiphertextMaterialDbItem>>();
6566

6667
sqlx::query!(
@@ -80,7 +81,8 @@ impl DbEventPublisher {
8081
let sns_ciphertexts_db = request
8182
.snsCtMaterials
8283
.iter()
83-
.map(SnsCiphertextMaterialDbItem::from)
84+
.zip(request.storageUrls)
85+
.map(|(sns_ct, url)| SnsCiphertextMaterialDbItem::new(sns_ct, url))
8486
.collect::<Vec<SnsCiphertextMaterialDbItem>>();
8587

8688
sqlx::query!(

‎kms-connector/crates/kms-worker/src/core/event_processor/decryption.rs‎

Lines changed: 5 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ use crate::core::{
88
use alloy::{
99
hex,
1010
primitives::{Address, Bytes, U256},
11-
sol_types::{Eip712Domain, SolValue},
11+
sol_types::Eip712Domain,
1212
};
1313
use anyhow::anyhow;
1414
use connector_utils::types::KmsGrpcRequest;
@@ -49,6 +49,7 @@ impl DecryptionProcessor {
4949
&self,
5050
decryption_id: U256,
5151
sns_materials: Vec<SnsCiphertextMaterial>,
52+
storage_urls: Vec<Vec<String>>,
5253
extra_data: Vec<u8>,
5354
user_decrypt_data: Option<UserDecryptionExtraData>,
5455
) -> anyhow::Result<KmsGrpcRequest> {
@@ -62,7 +63,7 @@ impl DecryptionProcessor {
6263
info!("Extracted key_id {key_id} from snsCtMaterials[0]");
6364

6465
let ciphertexts = self
65-
.prepare_ciphertexts(decryption_id, &key_id, sns_materials, &extra_data)
66+
.prepare_ciphertexts(decryption_id, &key_id, sns_materials, storage_urls)
6667
.await?;
6768

6869
// Convert alloy domain to protobuf domain
@@ -114,12 +115,11 @@ impl DecryptionProcessor {
114115
decryption_id: U256,
115116
key_id: &str,
116117
sns_materials: Vec<SnsCiphertextMaterial>,
117-
extra_data: &[u8],
118+
storage_urls: Vec<Vec<String>>,
118119
) -> anyhow::Result<Vec<TypedCiphertext>> {
119-
let s3_url_matrix = decode_s3_url_matrix(extra_data)?;
120120
let sns_ciphertext_materials = self
121121
.s3_service
122-
.retrieve_sns_ciphertext_materials(sns_materials, s3_url_matrix)
122+
.retrieve_sns_ciphertext_materials(sns_materials, storage_urls)
123123
.await;
124124

125125
if sns_ciphertext_materials.is_empty() {
@@ -144,11 +144,6 @@ impl DecryptionProcessor {
144144
}
145145
}
146146

147-
fn decode_s3_url_matrix(extra_data: &[u8]) -> anyhow::Result<Vec<Vec<String>>> {
148-
let (_version, urls): (u32, Vec<Vec<String>>) = SolValue::abi_decode_sequence(extra_data)?;
149-
Ok(urls)
150-
}
151-
152147
pub struct UserDecryptionExtraData {
153148
pub user_address: Address,
154149
pub public_key: Bytes,
@@ -162,40 +157,3 @@ impl UserDecryptionExtraData {
162157
}
163158
}
164159
}
165-
166-
#[cfg(test)]
167-
mod tests {
168-
use super::*;
169-
170-
#[test]
171-
fn test_urls_decoding() {
172-
let url_matrix = vec![
173-
vec![
174-
String::from("http://localhost"),
175-
String::from("https://127.0.0.1:80/"),
176-
],
177-
vec![String::from("http://localhost:81/test")],
178-
];
179-
180-
let extra_data = hex::decode(REMIX_GENERATED_EXTRA_DATA).unwrap();
181-
let decoded_url_matrix = decode_s3_url_matrix(&extra_data).unwrap();
182-
assert_eq!(url_matrix, decoded_url_matrix);
183-
}
184-
185-
// Encoded bytes retrieved with the following function in Remix
186-
// ```
187-
// function encode() public pure returns (bytes memory) {
188-
// string[] memory url_array1 = new string[](2);
189-
// url_array1[0] = "http://localhost";
190-
// url_array1[1] = "https://127.0.0.1:80/";
191-
// string[] memory url_array2 = new string[](1);
192-
// url_array2[0] = "http://localhost:81/test";
193-
// string[][] memory url_matrix = new string[][](2);
194-
// url_matrix[0] = url_array1;
195-
// url_matrix[1] = url_array2;
196-
// uint32 version = 1;
197-
// return abi.encode(version, url_matrix);
198-
// }
199-
// ```
200-
const REMIX_GENERATED_EXTRA_DATA: &str = "0x000000000000000000000000000000000000000000000000000000000000000100000000000000000000000000000000000000000000000000000000000000400000000000000000000000000000000000000000000000000000000000000002000000000000000000000000000000000000000000000000000000000000004000000000000000000000000000000000000000000000000000000000000001200000000000000000000000000000000000000000000000000000000000000002000000000000000000000000000000000000000000000000000000000000004000000000000000000000000000000000000000000000000000000000000000800000000000000000000000000000000000000000000000000000000000000010687474703a2f2f6c6f63616c686f737400000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000001568747470733a2f2f3132372e302e302e313a38302f0000000000000000000000000000000000000000000000000000000000000000000000000000000000000100000000000000000000000000000000000000000000000000000000000000200000000000000000000000000000000000000000000000000000000000000018687474703a2f2f6c6f63616c686f73743a38312f746573740000000000000000";
201-
}

‎kms-connector/crates/kms-worker/src/core/event_processor/processor.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,7 @@ impl DbEventProcessor {
8888
.prepare_decryption_request(
8989
req.decryptionId,
9090
req.snsCtMaterials,
91+
req.storageUrls,
9192
req.extraData.into(),
9293
None,
9394
)
@@ -98,6 +99,7 @@ impl DbEventProcessor {
9899
.prepare_decryption_request(
99100
req.decryptionId,
100101
req.snsCtMaterials,
102+
req.storageUrls,
101103
req.extraData.into(),
102104
Some(UserDecryptionExtraData::new(req.userAddress, req.publicKey)),
103105
)

‎kms-connector/crates/kms-worker/src/core/event_processor/s3.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -40,11 +40,11 @@ impl S3Service {
4040
pub async fn retrieve_sns_ciphertext_materials(
4141
&self,
4242
sns_materials: Vec<SnsCiphertextMaterial>,
43-
s3_urls_metrix: Vec<Vec<String>>,
43+
s3_urls_matrix: Vec<Vec<String>>,
4444
) -> Vec<TypedCiphertext> {
4545
let mut sns_ciphertext_materials = Vec::new();
4646

47-
for (sns_material, mut s3_urls) in sns_materials.into_iter().zip(s3_urls_metrix) {
47+
for (sns_material, mut s3_urls) in sns_materials.into_iter().zip(s3_urls_matrix) {
4848
// 1. For each SNS material, we try to retrieve its ciphertext from multiple possible S3 URLs
4949
// 1.1. We try to fetch the ciphertext for `self.s3_ct_retrieval_retries` times for each S3 URL
5050
// 2. Once we successfully retrieve a ciphertext from any of those URLs, we break out of the S3 URLs loop

‎kms-connector/crates/kms-worker/tests/event_picker/parallel.rs‎

Lines changed: 19 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,11 @@
1-
use std::time::Duration;
2-
31
use alloy::primitives::U256;
42
use connector_utils::{
53
tests::{rand::rand_sns_ct, setup::TestInstanceBuilder},
64
types::{GatewayEvent, db::SnsCiphertextMaterialDbItem},
75
};
86
use fhevm_gateway_bindings::decryption::Decryption::PublicDecryptionRequest;
97
use kms_worker::core::{Config, DbEventPicker, EventPicker};
8+
use std::time::Duration;
109
use tokio::time::timeout;
1110

1211
#[tokio::test]
@@ -18,11 +17,7 @@ async fn test_parallel_event_picker_one_events() -> anyhow::Result<()> {
1817
let mut event_picker1 = DbEventPicker::connect(test_instance.db().clone(), &config).await?;
1918

2019
let id0 = U256::ZERO;
21-
let sns_ct = vec![rand_sns_ct()];
22-
let sns_ciphertexts_db = sns_ct
23-
.iter()
24-
.map(SnsCiphertextMaterialDbItem::from)
25-
.collect::<Vec<SnsCiphertextMaterialDbItem>>();
20+
let sns_ciphertexts_db = vec![rand_sns_ct()];
2621

2722
println!("Inserting only one PublicDecryptionRequest for two event picker...");
2823
sqlx::query!(
@@ -47,7 +42,11 @@ async fn test_parallel_event_picker_one_events() -> anyhow::Result<()> {
4742
events0,
4843
vec![GatewayEvent::PublicDecryption(PublicDecryptionRequest {
4944
decryptionId: id0,
50-
snsCtMaterials: sns_ct.clone(),
45+
storageUrls: sns_ciphertexts_db
46+
.iter()
47+
.map(|s| s.storage_urls.clone())
48+
.collect(),
49+
snsCtMaterials: sns_ciphertexts_db.iter().map(|s| s.into()).collect(),
5150
extraData: vec![].into()
5251
})]
5352
);
@@ -68,11 +67,7 @@ async fn test_parallel_event_picker_two_events() -> anyhow::Result<()> {
6867

6968
let id0 = U256::ZERO;
7069
let id1 = U256::ONE;
71-
let sns_ct = vec![rand_sns_ct()];
72-
let sns_ciphertexts_db = sns_ct
73-
.iter()
74-
.map(SnsCiphertextMaterialDbItem::from)
75-
.collect::<Vec<SnsCiphertextMaterialDbItem>>();
70+
let sns_ciphertexts_db = vec![rand_sns_ct()];
7671

7772
println!("Inserting two PublicDecryptionRequest for two event picker...");
7873
sqlx::query!(
@@ -86,7 +81,7 @@ async fn test_parallel_event_picker_two_events() -> anyhow::Result<()> {
8681
sqlx::query!(
8782
"INSERT INTO public_decryption_requests VALUES ($1, $2, $3) ON CONFLICT DO NOTHING",
8883
id1.as_le_slice(),
89-
sns_ciphertexts_db as Vec<SnsCiphertextMaterialDbItem>,
84+
sns_ciphertexts_db.clone() as Vec<SnsCiphertextMaterialDbItem>,
9085
vec![],
9186
)
9287
.execute(test_instance.db())
@@ -101,15 +96,23 @@ async fn test_parallel_event_picker_two_events() -> anyhow::Result<()> {
10196
events0,
10297
vec![GatewayEvent::PublicDecryption(PublicDecryptionRequest {
10398
decryptionId: id0,
104-
snsCtMaterials: sns_ct.clone(),
99+
storageUrls: sns_ciphertexts_db
100+
.iter()
101+
.map(|s| s.storage_urls.clone())
102+
.collect(),
103+
snsCtMaterials: sns_ciphertexts_db.iter().map(|s| s.into()).collect(),
105104
extraData: vec![].into(),
106105
})]
107106
);
108107
assert_eq!(
109108
events1,
110109
vec![GatewayEvent::PublicDecryption(PublicDecryptionRequest {
111110
decryptionId: id1,
112-
snsCtMaterials: sns_ct,
111+
storageUrls: sns_ciphertexts_db
112+
.iter()
113+
.map(|s| s.storage_urls.clone())
114+
.collect(),
115+
snsCtMaterials: sns_ciphertexts_db.iter().map(|s| s.into()).collect(),
113116
extraData: vec![].into(),
114117
})]
115118
);

‎kms-connector/crates/kms-worker/tests/event_picker/simple.rs‎

Lines changed: 23 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,3 @@
1-
use std::time::Duration;
2-
31
use connector_utils::{
42
tests::{
53
rand::{rand_address, rand_digest, rand_public_key, rand_sns_ct, rand_u256},
@@ -15,6 +13,7 @@ use fhevm_gateway_bindings::{
1513
},
1614
};
1715
use kms_worker::core::{Config, DbEventPicker, EventPicker};
16+
use std::time::Duration;
1817

1918
#[tokio::test]
2019
async fn test_pick_public_decryption() -> anyhow::Result<()> {
@@ -24,17 +23,13 @@ async fn test_pick_public_decryption() -> anyhow::Result<()> {
2423
DbEventPicker::connect(test_instance.db().clone(), &Config::default()).await?;
2524

2625
let decryption_id = rand_u256();
27-
let sns_ct = vec![rand_sns_ct()];
28-
let sns_ciphertexts_db = sns_ct
29-
.iter()
30-
.map(SnsCiphertextMaterialDbItem::from)
31-
.collect::<Vec<SnsCiphertextMaterialDbItem>>();
26+
let sns_ciphertexts_db = vec![rand_sns_ct()];
3227

3328
println!("Triggering Postgres notification with PublicDecryptionRequest insertion...");
3429
sqlx::query!(
3530
"INSERT INTO public_decryption_requests VALUES ($1, $2, $3) ON CONFLICT DO NOTHING",
3631
decryption_id.as_le_slice(),
37-
sns_ciphertexts_db as Vec<SnsCiphertextMaterialDbItem>,
32+
sns_ciphertexts_db.clone() as Vec<SnsCiphertextMaterialDbItem>,
3833
vec![],
3934
)
4035
.execute(test_instance.db())
@@ -48,7 +43,11 @@ async fn test_pick_public_decryption() -> anyhow::Result<()> {
4843
events,
4944
vec![GatewayEvent::PublicDecryption(PublicDecryptionRequest {
5045
decryptionId: decryption_id,
51-
snsCtMaterials: sns_ct,
46+
storageUrls: sns_ciphertexts_db
47+
.iter()
48+
.map(|s| s.storage_urls.clone())
49+
.collect(),
50+
snsCtMaterials: sns_ciphertexts_db.iter().map(|s| s.into()).collect(),
5251
extraData: vec![].into(),
5352
})]
5453
);
@@ -64,19 +63,15 @@ async fn test_pick_user_decryption() -> anyhow::Result<()> {
6463
DbEventPicker::connect(test_instance.db().clone(), &Config::default()).await?;
6564

6665
let decryption_id = rand_u256();
67-
let sns_ct = vec![rand_sns_ct()];
66+
let sns_ciphertexts_db = vec![rand_sns_ct()];
6867
let user_address = rand_address();
6968
let public_key = rand_public_key();
70-
let sns_ciphertexts_db = sns_ct
71-
.iter()
72-
.map(SnsCiphertextMaterialDbItem::from)
73-
.collect::<Vec<SnsCiphertextMaterialDbItem>>();
7469

7570
println!("Triggering Postgres notification with UserDecryptionRequest insertion...");
7671
sqlx::query!(
7772
"INSERT INTO user_decryption_requests VALUES ($1, $2, $3, $4, $5) ON CONFLICT DO NOTHING",
7873
decryption_id.as_le_slice(),
79-
sns_ciphertexts_db as Vec<SnsCiphertextMaterialDbItem>,
74+
sns_ciphertexts_db.clone() as Vec<SnsCiphertextMaterialDbItem>,
8075
user_address.as_slice(),
8176
&public_key,
8277
vec![],
@@ -92,7 +87,11 @@ async fn test_pick_user_decryption() -> anyhow::Result<()> {
9287
events,
9388
vec![GatewayEvent::UserDecryption(UserDecryptionRequest {
9489
decryptionId: decryption_id,
95-
snsCtMaterials: sns_ct,
90+
storageUrls: sns_ciphertexts_db
91+
.iter()
92+
.map(|s| s.storage_urls.clone())
93+
.collect(),
94+
snsCtMaterials: sns_ciphertexts_db.iter().map(|s| s.into()).collect(),
9695
userAddress: user_address,
9796
publicKey: public_key.into(),
9897
extraData: vec![].into(),
@@ -283,16 +282,13 @@ async fn test_polling_backup() -> anyhow::Result<()> {
283282
let test_instance = TestInstanceBuilder::db_setup().await?;
284283

285284
let decryption_id = rand_u256();
286-
let sns_ct = vec![rand_sns_ct()];
287-
let sns_ciphertexts_db = sns_ct
288-
.iter()
289-
.map(SnsCiphertextMaterialDbItem::from)
290-
.collect::<Vec<SnsCiphertextMaterialDbItem>>();
285+
let sns_ciphertexts_db = vec![rand_sns_ct()];
286+
291287
println!("Inserting PublicDecryptionRequest before starting the event picker...");
292288
sqlx::query!(
293289
"INSERT INTO public_decryption_requests VALUES ($1, $2, $3) ON CONFLICT DO NOTHING",
294290
decryption_id.as_le_slice(),
295-
sns_ciphertexts_db as Vec<SnsCiphertextMaterialDbItem>,
291+
sns_ciphertexts_db.clone() as Vec<SnsCiphertextMaterialDbItem>,
296292
vec![],
297293
)
298294
.execute(test_instance.db())
@@ -312,7 +308,11 @@ async fn test_polling_backup() -> anyhow::Result<()> {
312308
events,
313309
vec![GatewayEvent::PublicDecryption(PublicDecryptionRequest {
314310
decryptionId: decryption_id,
315-
snsCtMaterials: sns_ct,
311+
storageUrls: sns_ciphertexts_db
312+
.iter()
313+
.map(|s| s.storage_urls.clone())
314+
.collect(),
315+
snsCtMaterials: sns_ciphertexts_db.iter().map(|s| s.into()).collect(),
316316
extraData: vec![].into(),
317317
})]
318318
);

‎kms-connector/crates/utils/src/tests/rand.rs‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
1+
use crate::types::db::SnsCiphertextMaterialDbItem;
12
use alloy::primitives::{Address, FixedBytes, U256};
2-
use fhevm_gateway_bindings::decryption::Decryption::SnsCiphertextMaterial;
33
use rand::Rng;
44

55
pub fn rand_u256() -> U256 {
@@ -22,10 +22,11 @@ pub fn rand_digest() -> FixedBytes<32> {
2222
rand::rng().random::<[u8; 32]>().into()
2323
}
2424

25-
pub fn rand_sns_ct() -> SnsCiphertextMaterial {
26-
SnsCiphertextMaterial {
27-
keyId: rand_u256(),
28-
ctHandle: rand::rng().random::<[u8; 32]>().into(),
29-
snsCiphertextDigest: rand::rng().random::<[u8; 32]>().into(),
25+
pub fn rand_sns_ct() -> SnsCiphertextMaterialDbItem {
26+
SnsCiphertextMaterialDbItem {
27+
key_id: rand::rng().random::<[u8; 32]>(),
28+
ct_handle: rand::rng().random::<[u8; 32]>(),
29+
sns_ciphertext_digest: rand::rng().random::<[u8; 32]>(),
30+
storage_urls: vec!["http://localhost:9000/ct/128".to_string()],
3031
}
3132
}

0 commit comments

Comments
 (0)