Skip to content

Commit ee5b60f

Browse files
isaacdecodedeudelins-zama
isaacdecoded
and
eudelins-zama
authored
feat(common): add PRSSInit and KeyReshareSameSet flows (#1170)
* feat(gateway-contracts): implement exceptional keygen corrective flow * feat(gateway-contracts): implement exceptional keygen corrective flow * test(gateway-contracts): add test for exceptional keygen corrective flow * chore(gateway-contracts): update bindings and selectors * ci(gateway-contracts): enable upgrade CI test for KMSGeneration contract * test(gateway-contracts): fix mock contracts tests * ci(gateway-contracts): fix upgrade CI test for KMSGeneration contract * refactor(gateway-contracts): remove epoch ID logic from keygen and rename reshare method * refactor(gateway-contracts): properly mock the epoch ID * refactor(gateway-contracts): replace epoch mention by keyReshare * feat(kms-connector): prss_init and retry keygen reshare support (#1151) * feat(kms-connector): support prss init flow * feat(kms-connector): support retry keygen reshare flow * chore(gateway-contracts): implement hardhat tasks for key resharing methods * feat(gateway-contracts): implement exceptional keygen corrective flow * feat(gateway-contracts): implement exceptional keygen corrective flow * test(gateway-contracts): add test for exceptional keygen corrective flow * chore(gateway-contracts): update bindings and selectors * ci(gateway-contracts): enable upgrade CI test for KMSGeneration contract * test(gateway-contracts): fix mock contracts tests * ci(gateway-contracts): fix upgrade CI test for KMSGeneration contract * refactor(gateway-contracts): remove epoch ID logic from keygen and rename reshare method * refactor(gateway-contracts): properly mock the epoch ID * refactor(gateway-contracts): replace epoch mention by keyReshare * chore(gateway-contracts): implement hardhat tasks for key resharing methods * test(gateway-contracts): add key resharing tasks unit tests * chore(kms-connector): improve error handling * chore(test-suite): update connector and gateway * fix(kms-connector): fix prss init encoding --------- Co-authored-by: eudelins-zama <simon.eudeline@zama.ai>
1 parent 85e7388 commit ee5b60f

37 files changed

+2565
-193
lines changed

‎gateway-contracts/contracts/KMSGeneration.sol‎

Lines changed: 47 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,12 @@ import { Strings } from "@openzeppelin/contracts/utils/Strings.sol";
1010
import { UUPSUpgradeableEmptyProxy } from "./shared/UUPSUpgradeableEmptyProxy.sol";
1111
import { GatewayConfigChecks } from "./shared/GatewayConfigChecks.sol";
1212
import { GatewayOwnable } from "./shared/GatewayOwnable.sol";
13-
import { PREP_KEYGEN_COUNTER_BASE, KEY_COUNTER_BASE, CRS_COUNTER_BASE } from "./shared/KMSRequestCounters.sol";
13+
import {
14+
PREP_KEYGEN_COUNTER_BASE,
15+
KEY_COUNTER_BASE,
16+
CRS_COUNTER_BASE,
17+
KEY_RESHARE_COUNTER_BASE
18+
} from "./shared/KMSRequestCounters.sol";
1419

1520
/**
1621
* @title KMSGeneration contract
@@ -163,6 +168,8 @@ contract KMSGeneration is
163168
// ----------------------------------------------------------------------------------------------
164169
/// @notice The parameters type used for the request
165170
mapping(uint256 requestId => ParamsType paramsType) requestParamsType;
171+
/// @notice The number of key resharing, used to generate the keyReshareIds.
172+
uint256 keyReshareCounter;
166173
}
167174

168175
/**
@@ -192,14 +199,19 @@ contract KMSGeneration is
192199
$.prepKeygenCounter = PREP_KEYGEN_COUNTER_BASE;
193200
$.keyCounter = KEY_COUNTER_BASE;
194201
$.crsCounter = CRS_COUNTER_BASE;
202+
$.keyReshareCounter = KEY_RESHARE_COUNTER_BASE;
195203
}
196204

197205
/**
198206
* @notice Re-initializes the contract from V1.
199207
*/
200208
/// @custom:oz-upgrades-unsafe-allow missing-initializer-call
201209
/// @custom:oz-upgrades-validate-as-initializer
202-
function reinitializeV2() public virtual reinitializer(REINITIALIZER_VERSION) {}
210+
function reinitializeV2() public virtual reinitializer(REINITIALIZER_VERSION) {
211+
KMSGenerationStorage storage $ = _getKMSGenerationStorage();
212+
213+
$.keyReshareCounter = KEY_RESHARE_COUNTER_BASE;
214+
}
203215

204216
/**
205217
* @notice See {IKMSGeneration-keygen}.
@@ -417,6 +429,39 @@ contract KMSGeneration is
417429
}
418430
}
419431

432+
/**
433+
* @notice See {IKMSGeneration-prssInit}.
434+
*/
435+
function prssInit() external virtual onlyGatewayOwner {
436+
emit PRSSInit();
437+
}
438+
439+
/**
440+
* @notice See {IKMSGeneration-keyReshareSameSet}.
441+
* @dev ⚠️ This function should only be called under exceptional circumstances.
442+
* It is intended for corrective flows when a previous resharing attempt failed.
443+
* Use with caution since incorrect usage may cause inconsistent key generation states.
444+
*/
445+
function keyReshareSameSet(uint256 keyId) external virtual onlyGatewayOwner {
446+
KMSGenerationStorage storage $ = _getKMSGenerationStorage();
447+
448+
if (!$.isRequestDone[keyId]) {
449+
revert KeyNotGenerated(keyId);
450+
}
451+
452+
// Get the prepKeygenId associated to the keyId and its params type.
453+
uint256 prepKeygenId = $.keygenIdPairs[keyId];
454+
ParamsType paramsType = $.requestParamsType[prepKeygenId];
455+
456+
// Generate a globally unique keyReshareId for the key resharing.
457+
// The counter is initialized at deployment such that keyReshareId's first byte uniquely
458+
// represents a key reshare request, with format: [0000 0110 | counter_1..31]
459+
$.keyReshareCounter++;
460+
uint256 keyReshareId = $.keyReshareCounter;
461+
462+
emit KeyReshareSameSet(prepKeygenId, keyId, keyReshareId, paramsType);
463+
}
464+
420465
/**
421466
* @notice See {IKMSGeneration-getKeyParamsType}.
422467
*/

‎gateway-contracts/contracts/interfaces/IKMSGeneration.sol‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,22 @@ interface IKMSGeneration {
7373
*/
7474
event ActivateCrs(uint256 crsId, string[] kmsNodeStorageUrls, bytes crsDigest);
7575

76+
/**
77+
* @notice Emitted to trigger the initialization of the PRSS (Pseudo-Random Secret Sharing).
78+
* @dev This is a temporary event to initialize PRSS until implementation of a proper key resharing.
79+
*/
80+
event PRSSInit();
81+
82+
/**
83+
* @notice Emitted to trigger the reshare of the specified key ID.
84+
* @dev This is a temporary event to reshare the specified key ID until implementation of a proper key resharing.
85+
* @param prepKeygenId The ID of the preprocessing keygen request.
86+
* @param keyId The ID of the key to reshare.
87+
* @param keyReshareId The ID of the key reshare request.
88+
* @param paramsType The type of FHE parameters to use.
89+
*/
90+
event KeyReshareSameSet(uint256 prepKeygenId, uint256 keyId, uint256 keyReshareId, ParamsType paramsType);
91+
7692
/**
7793
* @notice Error thrown when a KMS node has already signed for a preprocessing keygen response.
7894
* @param prepKeygenId The ID of the preprocessing keygen request.
@@ -142,6 +158,19 @@ interface IKMSGeneration {
142158
*/
143159
function crsgenResponse(uint256 crsId, bytes calldata crsDigest, bytes calldata signature) external;
144160

161+
/**
162+
* @notice Trigger the initialization of the PRSS (Pseudo-Random Secret Sharing).
163+
* @dev This is a temporary method to initialize PRSS until implementation of a proper key resharing.
164+
*/
165+
function prssInit() external;
166+
167+
/**
168+
* @notice Trigger the reshare of the given key ID.
169+
* @dev This is a temporary method to reshare the specified key ID until implementation of a proper key resharing.
170+
* @param keyId The ID of the key to reshare.
171+
*/
172+
function keyReshareSameSet(uint256 keyId) external;
173+
145174
/**
146175
* @notice Get the parameters type used for the key generation.
147176
* @param keyId The ID of the key.

‎gateway-contracts/contracts/mocks/KMSGenerationMock.sol‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,9 +27,14 @@ contract KMSGenerationMock {
2727

2828
event ActivateCrs(uint256 crsId, string[] kmsNodeStorageUrls, bytes crsDigest);
2929

30+
event PRSSInit();
31+
32+
event KeyReshareSameSet(uint256 prepKeygenId, uint256 keyId, uint256 keyReshareId, ParamsType paramsType);
33+
3034
uint256 prepKeygenCounter = 3 << 248;
3135
uint256 keyCounter = 4 << 248;
3236
uint256 crsCounter = 5 << 248;
37+
uint256 keyReshareCounter = 6 << 248;
3338

3439
function keygen(ParamsType paramsType) external {
3540
prepKeygenCounter++;
@@ -64,4 +69,18 @@ contract KMSGenerationMock {
6469

6570
emit ActivateCrs(crsId, kmsNodeStorageUrls, crsDigest);
6671
}
72+
73+
function prssInit() external {
74+
emit PRSSInit();
75+
}
76+
77+
function keyReshareSameSet(uint256 keyId) external {
78+
prepKeygenCounter++;
79+
uint256 prepKeygenId = prepKeygenCounter;
80+
keyReshareCounter++;
81+
uint256 keyReshareId = keyReshareCounter;
82+
ParamsType paramsType;
83+
84+
emit KeyReshareSameSet(prepKeygenId, keyId, keyReshareId, paramsType);
85+
}
6786
}

‎gateway-contracts/contracts/shared/KMSRequestCounters.sol‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,8 @@ enum RequestType {
1212
UserDecrypt, // 2
1313
PrepKeygen, // 3
1414
Keygen, // 4
15-
Crsgen // 5
15+
Crsgen, // 5
16+
KeyReshare // 6
1617
}
1718

1819
// Bit position to left shift for initializing the counters
@@ -35,3 +36,6 @@ uint256 constant KEY_COUNTER_BASE = uint256(RequestType.Keygen) << REQUEST_TYPE_
3536

3637
// CRS generation requestId format in bytes: [0000 0101 | counter_1..31]
3738
uint256 constant CRS_COUNTER_BASE = uint256(RequestType.Crsgen) << REQUEST_TYPE_SHIFT;
39+
40+
// Key resharing requestId format in bytes: [0000 0110 | counter_1..31]
41+
uint256 constant KEY_RESHARE_COUNTER_BASE = uint256(RequestType.KeyReshare) << REQUEST_TYPE_SHIFT;

‎gateway-contracts/hardhat.config.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import "./tasks/generateKmsMaterials";
1919
import "./tasks/getters";
2020
import "./tasks/ownership";
2121
import "./tasks/pauseContracts";
22+
import "./tasks/reshareKeys";
2223
import "./tasks/safeSmartAccounts";
2324
import "./tasks/upgradeContracts";
2425

‎gateway-contracts/rust_bindings/src/decryption.rs‎

Lines changed: 4 additions & 4 deletions
Large diffs are not rendered by default.

0 commit comments

Comments
 (0)