Skip to content

refactor: add psk receiver#5552

Merged
jmayclin merged 6 commits intoaws:mainfrom
jmayclin:2025-10-09-psk-receiver
Oct 12, 2025
Merged

refactor: add psk receiver#5552
jmayclin merged 6 commits intoaws:mainfrom
jmayclin:2025-10-09-psk-receiver

Conversation

@jmayclin
Copy link
Copy Markdown
Contributor

Description of changes:

This adds the PskReceiver component of the HMAC based PSK system.

Testing:

Following the approach from the PSK Provider. We extract the update into a single "poll update" method which can be neatly unit tested.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.

@github-actions github-actions bot added the s2n-core team label Oct 10, 2025
@jmayclin jmayclin requested review from jouho and maddeleine October 10, 2025 16:05
// matching PSK identity
let session_name = b"test_session".to_vec();
let psk_identity = PskIdentity::new(&session_name, &epoch_secret).unwrap();
secret_state.find_match(psk_identity.clone()).unwrap();
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do you want to assert anything on the returned PSK? It looks like you're not doing anything with what you returned.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

s2n-tls doesn't expose any getters for PSKs ☹️ . They are totally opaque, so there isn't really anything to assert. Unless I write them to bytes and look at that, but it felt like overkill.

Comment on lines +207 to +214
let current_epoch = epoch_schedule::current_epoch();

let update = secret_state
.fetch_secrets(&kms_client, current_epoch, &failure_notification)
.await;
if update.is_err() {
anyhow::bail!("failed to fetch keys during startup");
}
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You technically don't need these lines since you just do the same thing in the async block. But I suppose you're wanting it for the error message?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ya, we want to ensure that a successful call to initialize means that all of the epoch keys have been successfully fetched.

If we just relied on the async block, then we might say that we had successfully initialized even though we have only fetched a single secret.

}

#[cfg(test)]
mod secret_state_tests {
Copy link
Copy Markdown
Contributor

@maddeleine maddeleine Oct 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What about a test that the psk provider/receiver fetch generates the same thing? I mean, seems like the only difference is that the server fetches one more epoch than the client. I'm mostly interested in testing that the hmac output are the same.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hmm, it's gonna be a bit messy, because the secret state is internal so there isn't any place that we can get both the internal state of the receiver and the internal state of the provider, unless we add test-only getters.

More generally, I feel comfortable relying on the handshake tests and unit tests to assert that the results are individually what we want.

If we want more tests, I think I'd push for abstracting the unix timestamp into a global clock that we can mock, and using the tokio fast-forward utilities to skip time.

jmayclin and others added 2 commits October 10, 2025 16:32
Co-authored-by: maddeleine <59030281+maddeleine@users.noreply.github.com>
Co-authored-by: maddeleine <59030281+maddeleine@users.noreply.github.com>
@jmayclin jmayclin requested a review from maddeleine October 10, 2025 23:44
struct ReceiverSecrets {
smoothing_factor: Duration,
pub trusted_key_arns: Vec<KeyArn>,
pub epoch_secrets: RwLock<HashMap<u64, HashMap<KeyArn, EpochSecret>>>,
Copy link
Copy Markdown
Contributor

@jouho jouho Oct 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why do we need RwLock here - presumably to allow many concurrent handshake callbacks to read it while permitting exclusive writes when we update the secret?

Should we also have tests for concurrent client connections and make sure secrets are updated as expected?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Testing these types of concurrent things is very difficult, so it's relatively rare to have explicit test cases for it. In general, I'm comfortable relying on the API of the read/write lock for safety, but I can look at adding some concurrent usage as part of #5552 (comment)

jmayclin and others added 2 commits October 10, 2025 17:43
Co-authored-by: Jou Ho <43765840+jouho@users.noreply.github.com>
Co-authored-by: Jou Ho <43765840+jouho@users.noreply.github.com>
@jmayclin jmayclin added this pull request to the merge queue Oct 12, 2025
Merged via the queue into aws:main with commit 1eb6fac Oct 12, 2025
50 checks passed
@jmayclin jmayclin deleted the 2025-10-09-psk-receiver branch October 12, 2025 21:30
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants