Add persistent key storage to the Batch Poster #957
Conversation
…er and the Caff Node.
…tion to get tests passing
zacshowa
requested review from
ImJeremyHe,
Sneh1999,
jbearer,
jjeangal,
philippecamacho and
varun-doshi
as code owners
Summary of ChangesHello @zacshowa, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request enables the Batch Poster to utilize persistent key storage, mirroring the functionality previously implemented for the TEE Caff node. It involves significant modifications to the Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adapts the persistent key storage mechanism for the batch poster, enabling it to use keys from sources like AWS KMS via the EspressoKeyManager. The changes primarily involve updating the EspressoKeyManager to read an enclave private key, modifying configurations to support this, and updating tests accordingly. The overall approach is sound, but I've identified a potential logic issue in the key acquisition process within EspressoKeyManager that could lead to a runtime panic. My review includes a suggestion to make this logic more robust. The rest of the changes appear correct and align with the PR's objective.
| if keyPairAttestationsPath != "" && chainID != 0 { | ||
| // Read enclave private key | ||
| privKey, err = espresso_tee_utils.ReadEnclavePrivateKey(keyPairAttestationsPath, chainID) | ||
| if err != nil { | ||
| panic(err) | ||
| log.Crit("error reading enclave private key for Espresso Key Manager", "path", keyPairAttestationsPath, "err", err) | ||
| } | ||
| } else { | ||
|
|
||
| } else if servicePersistentPrivateKey != nil { | ||
| privKey = servicePersistentPrivateKey | ||
| } else { | ||
| panic("either keyPairAttestationsPath and chainID must be provided, or servicePersistentPrivateKey must be non-nil") | ||
| } |
There was a problem hiding this comment.
The current logic for obtaining a private key has a potential issue. If keyPairAttestationsPath is provided but espresso_tee_utils.ReadEnclavePrivateKey returns a nil key without an error, the code doesn't fall back to using servicePersistentPrivateKey. This would result in the EspressoKeyManager being created with a nil private key, which will cause a panic later. The logic should be restructured to ensure a fallback mechanism is in place.
The proposed change ensures that if reading from the enclave path doesn't yield a key, it will then attempt to use servicePersistentPrivateKey, and only panic if no key can be obtained from any source.
if keyPairAttestationsPath != "" && chainID != 0 {
// Read enclave private key
var err error
privKey, err = espresso_tee_utils.ReadEnclavePrivateKey(keyPairAttestationsPath, chainID)
if err != nil {
log.Crit("error reading enclave private key for Espresso Key Manager", "path", keyPairAttestationsPath, "err", err)
}
}
if privKey == nil {
if servicePersistentPrivateKey != nil {
privKey = servicePersistentPrivateKey
} else {
panic("could not obtain a private key from any source")
}
}There was a problem hiding this comment.
ReadEnclavePrivateKey only returns (privateKey, nil), or (nil, someErr)
ImJeremyHe
left a comment
ImJeremyHe
left a comment
There was a problem hiding this comment.
I think a simple unit test would help clarify and verify the conditions under which the batcher should be up or down.
| servicePersistentPrivateKey *ecdsa.PrivateKey, | ||
| zkAttestationServiceURL string, | ||
| keyPairAttestationsPath string, | ||
| chainID uint64, |
There was a problem hiding this comment.
I’m a bit confused here and may be missing some context. Could you help clarify the role of servicePersistentPrivateKey and the purpose of keyPairAttestationPath?
Based on the logic below, they appear to overlap in responsibility. Would it be reasonable to abstract a shared struct for private key access?
There was a problem hiding this comment.
I would agree, they absolutely have overlapping responsibilities, currently the batch poster is supposed to utilize the attestation path and read the key in from the KMS when calling the key managers constructor. However, the caff node will be using the KMS in nitro.go and eventually passing the key to the key manager. My intention is for the key manager constructor to be the sole place we interact with the AWS KMS in the future. The only reason the key parameter exists still is because I wanted to limit the scope of this PR to the batcher. I believe with the groundwork laid in this PR, the caff node can be made to use the key manager to handle its key relatively easily.
I think there might be some additional things to consider there though. Like we need to ensure that if for some reason anyone decided to run a caff node and batch poster on the same machine, that the caff node wouldn't try/be able to access the batch posters key.
Closes Asana Ticket: Add persistent key storage to batch poster.
This PR:
Adapts the persistant key storage mechanism made for the TEE Caff node to be used by the batch poster via the
EspressoKeyManagerThis PR does not:
Add any new tests. This work is not really unit testable as we don't want our CI to be creating things in an AWS KMS. However, I have fixed, the behavior of some tests that were seemingly broken.
Change the register on startup approach that the batcher takes to registering it's persistent key. This will not affect the batchers behavior, as this code section already has checks (via the key manager) to determine if it has already registered, so nothing weird will happen on any startup after the initial registration.
Key places to review:
Private key generation in the key manager tests vs e2e tests. I ended up with 2 different solutions, the one done in the key manager tests was written by claude code, and produces the exact same account as the one I produced for the batch poster, but only because hooking into the batch posters config was the only reasonable way I could see to give it a private key during tests. It's a little "hacky" in the batchers constructor, but it seemed like the easiest way to accomplish something specific to tests.
All of the changes in
espresso/key-manager/key_manager.goConstruction of the key manager in
arbnode/batch_poster.goandarbnode/espresso_caff_node.goChanges to the E2E test to provide the batch poster with a dummy persistent key.
Changes to
espresso/key-manager/key_manager_test.goChanges to the EspressoTEEVerifierMock contract in espresso-tee-contracts-legacy. I believe this may fix some changes @lukeiannucci was having recently.
How to test this PR:
Run the E2E test.
Things tested:
E2E test is passing locally.
I will also be testing this against some internal devnet, or against a testnet.