mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
[codex] request desktop attestation from app (#20619)
## Summary TL;DR: teaches `codex-rs` / app-server to request a desktop-provided attestation token and attach it as `x-oai-attestation` on the scoped ChatGPT Codex request paths.  ## Details This PR teaches the Codex app-server runtime how to request and attach an attestation token. It does not generate DeviceCheck tokens directly; instead, it relies on the connected desktop app to advertise that it can generate attestation and then asks that app for a fresh header value when needed. The flow is: 1. The Codex desktop app connects to app-server. 2. During `initialize`, the app can advertise that it supports `requestAttestation`. 3. Before app-server calls selected ChatGPT Codex endpoints, it sends the internal server request `attestation/generate` to the app. 4. app-server receives a pre-encoded header value back. 5. app-server forwards that value as `x-oai-attestation` on the scoped outbound requests. The code in this repo is mostly protocol and runtime plumbing: it adds the app-server request/response shape, introduces an attestation provider in core, wires that provider into Responses / compaction / realtime setup paths, and covers the intended scoping with tests. The signed macOS DeviceCheck generation remains owned by the desktop app PR. ## Related PR - Codex desktop app implementation: https://github.com/openai/openai/pull/878649 ## Validation <details> <summary>Tests run</summary> ```sh cargo test -p codex-app-server-protocol cargo test -p codex-core attestation --lib cargo test -p codex-app-server --lib attestation ``` Also ran: ```sh just fix -p codex-core just fix -p codex-app-server just fix -p codex-app-server-protocol just fmt just write-app-server-schema ``` </details> <details> <summary>E2E DeviceCheck validation</summary> First validated the signed desktop app boundary directly: launched a packaged signed `Codex.app`, sent `attestation/generate`, decoded the returned `v1.` attestation header, and validated the extracted DeviceCheck token with `personal/jm/verify_devicecheck_token.py` using bundle ID `com.openai.codex`. Apple returned `status_code: 200` and `is_ok: true`. Then ran the fuller app + app-server flow. The packaged `Codex.app` launched a current-branch app-server via `CODEX_CLI_PATH`, and a local MITM proxy intercepted outbound `chatgpt.com` traffic. The app-server requested `attestation/generate` from the real Electron app process, and the intercepted `/backend-api/codex/responses` traffic included `x-oai-attestation` on both routes: ```text GET /backend-api/codex/responses Upgrade: websocket x-oai-attestation: present POST /backend-api/codex/responses Upgrade: none x-oai-attestation: present ``` The captured header decoded to a DeviceCheck token that also validated with Apple for `com.openai.codex` (`status_code: 200`, `is_ok: true`, team `2DC432GLL2`). </details> --------- Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
@@ -0,0 +1,217 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::http::HeaderValue;
|
||||
use codex_app_server_protocol::AttestationGenerateParams;
|
||||
use codex_app_server_protocol::AttestationGenerateResponse;
|
||||
use codex_app_server_protocol::ServerRequestPayload;
|
||||
use codex_core::AttestationContext;
|
||||
use codex_core::AttestationProvider;
|
||||
use codex_core::GenerateAttestationFuture;
|
||||
use serde::Serialize;
|
||||
use tokio::time::Duration;
|
||||
use tokio::time::timeout;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::outgoing_message::OutgoingMessageSender;
|
||||
use crate::thread_state::ThreadStateManager;
|
||||
|
||||
const ATTESTATION_GENERATE_TIMEOUT: Duration = Duration::from_millis(100);
|
||||
|
||||
pub(crate) fn app_server_attestation_provider(
|
||||
outgoing: Arc<OutgoingMessageSender>,
|
||||
thread_state_manager: ThreadStateManager,
|
||||
) -> Arc<dyn AttestationProvider> {
|
||||
Arc::new(AppServerAttestationProvider {
|
||||
outgoing,
|
||||
thread_state_manager,
|
||||
})
|
||||
}
|
||||
|
||||
struct AppServerAttestationProvider {
|
||||
outgoing: Arc<OutgoingMessageSender>,
|
||||
thread_state_manager: ThreadStateManager,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for AppServerAttestationProvider {
|
||||
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
formatter
|
||||
.debug_struct("AppServerAttestationProvider")
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl AttestationProvider for AppServerAttestationProvider {
|
||||
fn header_for_request(&self, context: AttestationContext) -> GenerateAttestationFuture<'_> {
|
||||
let outgoing = self.outgoing.clone();
|
||||
let thread_state_manager = self.thread_state_manager.clone();
|
||||
Box::pin(async move {
|
||||
request_attestation_header_value_with_timeout(
|
||||
outgoing,
|
||||
thread_state_manager,
|
||||
context.thread_id,
|
||||
ATTESTATION_GENERATE_TIMEOUT,
|
||||
)
|
||||
.await
|
||||
.and_then(|value| HeaderValue::from_bytes(value.as_bytes()).ok())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
async fn request_attestation_header_value_with_timeout(
|
||||
outgoing: Arc<OutgoingMessageSender>,
|
||||
thread_state_manager: ThreadStateManager,
|
||||
thread_id: codex_protocol::ThreadId,
|
||||
timeout_duration: Duration,
|
||||
) -> Option<String> {
|
||||
let connection_id = thread_state_manager
|
||||
.first_attestation_capable_connection_for_thread(thread_id)
|
||||
.await?;
|
||||
|
||||
let connection_ids = [connection_id];
|
||||
let (request_id, rx) = outgoing
|
||||
.send_request_to_connections(
|
||||
Some(&connection_ids),
|
||||
ServerRequestPayload::AttestationGenerate(AttestationGenerateParams {}),
|
||||
/*thread_id*/ None,
|
||||
)
|
||||
.await;
|
||||
|
||||
let result = match timeout(timeout_duration, rx).await {
|
||||
Ok(Ok(Ok(result))) => result,
|
||||
Ok(Ok(Err(err))) => {
|
||||
warn!(
|
||||
code = err.code,
|
||||
message = %err.message,
|
||||
"attestation generation request failed"
|
||||
);
|
||||
return app_server_attestation_header_value(
|
||||
AppServerAttestationStatus::RequestFailed,
|
||||
/*token*/ None,
|
||||
);
|
||||
}
|
||||
Ok(Err(err)) => {
|
||||
warn!("attestation generation request canceled: {err}");
|
||||
return app_server_attestation_header_value(
|
||||
AppServerAttestationStatus::RequestCanceled,
|
||||
/*token*/ None,
|
||||
);
|
||||
}
|
||||
Err(_) => {
|
||||
let _canceled = outgoing.cancel_request(&request_id).await;
|
||||
warn!(
|
||||
timeout_seconds = timeout_duration.as_secs(),
|
||||
"attestation generation request timed out"
|
||||
);
|
||||
return app_server_attestation_header_value(
|
||||
AppServerAttestationStatus::Timeout,
|
||||
/*token*/ None,
|
||||
);
|
||||
}
|
||||
};
|
||||
|
||||
match serde_json::from_value::<AttestationGenerateResponse>(result) {
|
||||
Ok(response) => app_server_attestation_header_value(
|
||||
AppServerAttestationStatus::Ok,
|
||||
Some(&response.token),
|
||||
),
|
||||
Err(err) => {
|
||||
warn!("failed to deserialize attestation generation response: {err}");
|
||||
app_server_attestation_header_value(
|
||||
AppServerAttestationStatus::MalformedResponse,
|
||||
/*token*/ None,
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
enum AppServerAttestationStatus {
|
||||
Ok,
|
||||
Timeout,
|
||||
RequestFailed,
|
||||
RequestCanceled,
|
||||
MalformedResponse,
|
||||
}
|
||||
|
||||
impl AppServerAttestationStatus {
|
||||
const fn code(self) -> u8 {
|
||||
match self {
|
||||
Self::Ok => 0,
|
||||
Self::Timeout => 1,
|
||||
Self::RequestFailed => 2,
|
||||
Self::RequestCanceled => 3,
|
||||
Self::MalformedResponse => 4,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct AppServerAttestationEnvelope<'a> {
|
||||
v: u8,
|
||||
s: u8,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
t: Option<&'a str>,
|
||||
}
|
||||
|
||||
fn app_server_attestation_header_value(
|
||||
status: AppServerAttestationStatus,
|
||||
token: Option<&str>,
|
||||
) -> Option<String> {
|
||||
serde_json::to_string(&AppServerAttestationEnvelope {
|
||||
v: 1,
|
||||
s: status.code(),
|
||||
t: token,
|
||||
})
|
||||
.map_err(|err| warn!("failed to serialize app-server attestation envelope: {err}"))
|
||||
.ok()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::AppServerAttestationStatus;
|
||||
use super::app_server_attestation_header_value;
|
||||
use pretty_assertions::assert_eq;
|
||||
|
||||
#[test]
|
||||
fn app_server_attestation_header_value_wraps_opaque_client_payloads() {
|
||||
assert_eq!(
|
||||
app_server_attestation_header_value(
|
||||
AppServerAttestationStatus::Ok,
|
||||
Some("v1.opaque-client-payload"),
|
||||
),
|
||||
Some(r#"{"v":1,"s":0,"t":"v1.opaque-client-payload"}"#.to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn app_server_attestation_header_value_reports_app_server_failures() {
|
||||
assert_eq!(
|
||||
app_server_attestation_header_value(
|
||||
AppServerAttestationStatus::Timeout,
|
||||
/*token*/ None,
|
||||
),
|
||||
Some(r#"{"v":1,"s":1}"#.to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
app_server_attestation_header_value(
|
||||
AppServerAttestationStatus::RequestFailed,
|
||||
/*token*/ None,
|
||||
),
|
||||
Some(r#"{"v":1,"s":2}"#.to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
app_server_attestation_header_value(
|
||||
AppServerAttestationStatus::RequestCanceled,
|
||||
/*token*/ None,
|
||||
),
|
||||
Some(r#"{"v":1,"s":3}"#.to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
app_server_attestation_header_value(
|
||||
AppServerAttestationStatus::MalformedResponse,
|
||||
/*token*/ None
|
||||
),
|
||||
Some(r#"{"v":1,"s":4}"#.to_string())
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -74,6 +74,7 @@ use tracing_subscriber::util::SubscriberInitExt;
|
||||
|
||||
mod analytics_utils;
|
||||
mod app_server_tracing;
|
||||
mod attestation;
|
||||
mod bespoke_event_handling;
|
||||
mod command_exec;
|
||||
mod config;
|
||||
@@ -938,7 +939,14 @@ pub async fn run_main_with_transport_options(
|
||||
),
|
||||
)
|
||||
.await;
|
||||
processor.connection_initialized(connection_id).await;
|
||||
processor
|
||||
.connection_initialized(
|
||||
connection_id,
|
||||
connection_state
|
||||
.session
|
||||
.request_attestation(),
|
||||
)
|
||||
.await;
|
||||
connection_state
|
||||
.outbound_initialized
|
||||
.store(true, std::sync::atomic::Ordering::Release);
|
||||
|
||||
@@ -187,6 +187,7 @@ mod tests {
|
||||
thread_store,
|
||||
Some(state_db.clone()),
|
||||
"11111111-1111-4111-8111-111111111111".to_string(),
|
||||
/*attestation_provider*/ None,
|
||||
));
|
||||
thread_manager.start_thread(good_config).await?;
|
||||
thread_manager.start_thread(bad_config).await?;
|
||||
|
||||
@@ -4,6 +4,7 @@ use std::sync::Arc;
|
||||
use std::sync::OnceLock;
|
||||
use std::sync::atomic::AtomicBool;
|
||||
|
||||
use crate::attestation::app_server_attestation_provider;
|
||||
use crate::config_manager::ConfigManager;
|
||||
use crate::connection_rpc_gate::ConnectionRpcGate;
|
||||
use crate::error_code::invalid_request;
|
||||
@@ -34,6 +35,7 @@ use crate::request_processors::WindowsSandboxRequestProcessor;
|
||||
use crate::request_serialization::QueuedInitializedRequest;
|
||||
use crate::request_serialization::RequestSerializationQueueKey;
|
||||
use crate::request_serialization::RequestSerializationQueues;
|
||||
use crate::thread_state::ConnectionCapabilities;
|
||||
use crate::thread_state::ThreadStateManager;
|
||||
use crate::transport::AppServerTransport;
|
||||
use crate::transport::RemoteControlHandle;
|
||||
@@ -82,6 +84,7 @@ use tokio::time::timeout;
|
||||
use tracing::Instrument;
|
||||
|
||||
const EXTERNAL_AUTH_REFRESH_TIMEOUT: Duration = Duration::from_secs(10);
|
||||
|
||||
#[derive(Clone)]
|
||||
struct ExternalAuthRefreshBridge {
|
||||
outgoing: Arc<OutgoingMessageSender>,
|
||||
@@ -186,6 +189,7 @@ pub(crate) struct InitializedConnectionSessionState {
|
||||
pub(crate) opted_out_notification_methods: HashSet<String>,
|
||||
pub(crate) app_server_client_name: String,
|
||||
pub(crate) client_version: String,
|
||||
pub(crate) request_attestation: bool,
|
||||
}
|
||||
|
||||
impl Default for ConnectionSessionState {
|
||||
@@ -231,6 +235,12 @@ impl ConnectionSessionState {
|
||||
.map(|session| session.client_version.as_str())
|
||||
}
|
||||
|
||||
pub(crate) fn request_attestation(&self) -> bool {
|
||||
self.initialized
|
||||
.get()
|
||||
.is_some_and(|session| session.request_attestation)
|
||||
}
|
||||
|
||||
pub(crate) fn initialize(&self, session: InitializedConnectionSessionState) -> Result<(), ()> {
|
||||
self.initialized.set(session).map_err(|_| ())
|
||||
}
|
||||
@@ -280,6 +290,7 @@ impl MessageProcessor {
|
||||
auth_manager.set_external_auth(Arc::new(ExternalAuthRefreshBridge {
|
||||
outgoing: outgoing.clone(),
|
||||
}));
|
||||
let thread_state_manager = ThreadStateManager::new();
|
||||
// The thread store is intentionally process-scoped. Config reloads can
|
||||
// affect per-thread behavior, but they must not move newly started,
|
||||
// resumed, or forked threads to a different persistence backend/root.
|
||||
@@ -293,13 +304,16 @@ impl MessageProcessor {
|
||||
Arc::clone(&thread_store),
|
||||
state_db.clone(),
|
||||
installation_id,
|
||||
Some(app_server_attestation_provider(
|
||||
outgoing.clone(),
|
||||
thread_state_manager.clone(),
|
||||
)),
|
||||
));
|
||||
thread_manager
|
||||
.plugins_manager()
|
||||
.set_analytics_events_client(analytics_events_client.clone());
|
||||
|
||||
let pending_thread_unloads = Arc::new(Mutex::new(HashSet::new()));
|
||||
let thread_state_manager = ThreadStateManager::new();
|
||||
let thread_watch_manager =
|
||||
crate::thread_status::ThreadWatchManager::new_with_outgoing(outgoing.clone());
|
||||
let thread_list_state_permit = Arc::new(Semaphore::new(/*permits*/ 1));
|
||||
@@ -620,9 +634,18 @@ impl MessageProcessor {
|
||||
.await;
|
||||
}
|
||||
|
||||
pub(crate) async fn connection_initialized(&self, connection_id: ConnectionId) {
|
||||
pub(crate) async fn connection_initialized(
|
||||
&self,
|
||||
connection_id: ConnectionId,
|
||||
request_attestation: bool,
|
||||
) {
|
||||
self.thread_processor
|
||||
.connection_initialized(connection_id)
|
||||
.connection_initialized(
|
||||
connection_id,
|
||||
ConnectionCapabilities {
|
||||
request_attestation,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
@@ -718,7 +741,12 @@ impl MessageProcessor {
|
||||
.await?;
|
||||
if connection_initialized {
|
||||
self.thread_processor
|
||||
.connection_initialized(connection_id)
|
||||
.connection_initialized(
|
||||
connection_id,
|
||||
ConnectionCapabilities {
|
||||
request_attestation: session.request_attestation(),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
}
|
||||
return Ok(());
|
||||
|
||||
@@ -267,7 +267,7 @@ impl OutgoingMessageSender {
|
||||
RequestId::Integer(self.next_server_request_id.fetch_add(1, Ordering::Relaxed))
|
||||
}
|
||||
|
||||
async fn send_request_to_connections(
|
||||
pub(crate) async fn send_request_to_connections(
|
||||
&self,
|
||||
connection_ids: Option<&[ConnectionId]>,
|
||||
request: ServerRequestPayload,
|
||||
|
||||
@@ -473,6 +473,7 @@ use crate::error_code::internal_error;
|
||||
use crate::error_code::invalid_request;
|
||||
use crate::filters::compute_source_filters;
|
||||
use crate::filters::source_kind_matches;
|
||||
use crate::thread_state::ConnectionCapabilities;
|
||||
use crate::thread_state::ThreadListenerCommand;
|
||||
use crate::thread_state::ThreadState;
|
||||
use crate::thread_state::ThreadStateManager;
|
||||
|
||||
@@ -65,15 +65,17 @@ impl InitializeRequestProcessor {
|
||||
// experimental API). Proposed direction is instance-global first-write-wins
|
||||
// with initialize-time mismatch rejection.
|
||||
let analytics_initialize_params = params.clone();
|
||||
let (experimental_api_enabled, opt_out_notification_methods) = match params.capabilities {
|
||||
Some(capabilities) => (
|
||||
capabilities.experimental_api,
|
||||
capabilities
|
||||
.opt_out_notification_methods
|
||||
.unwrap_or_default(),
|
||||
),
|
||||
None => (false, Vec::new()),
|
||||
};
|
||||
let (experimental_api_enabled, request_attestation, opt_out_notification_methods) =
|
||||
match params.capabilities {
|
||||
Some(capabilities) => (
|
||||
capabilities.experimental_api,
|
||||
capabilities.request_attestation,
|
||||
capabilities
|
||||
.opt_out_notification_methods
|
||||
.unwrap_or_default(),
|
||||
),
|
||||
None => (false, false, Vec::new()),
|
||||
};
|
||||
let ClientInfo {
|
||||
name,
|
||||
title: _title,
|
||||
@@ -95,6 +97,7 @@ impl InitializeRequestProcessor {
|
||||
opted_out_notification_methods: opt_out_notification_methods.into_iter().collect(),
|
||||
app_server_client_name: name.clone(),
|
||||
client_version: version,
|
||||
request_attestation,
|
||||
})
|
||||
.is_err()
|
||||
{
|
||||
|
||||
@@ -2230,9 +2230,13 @@ impl ThreadRequestProcessor {
|
||||
self.thread_manager.subscribe_thread_created()
|
||||
}
|
||||
|
||||
pub(crate) async fn connection_initialized(&self, connection_id: ConnectionId) {
|
||||
pub(crate) async fn connection_initialized(
|
||||
&self,
|
||||
connection_id: ConnectionId,
|
||||
capabilities: ConnectionCapabilities,
|
||||
) {
|
||||
self.thread_state_manager
|
||||
.connection_initialized(connection_id)
|
||||
.connection_initialized(connection_id, capabilities)
|
||||
.await;
|
||||
}
|
||||
|
||||
|
||||
@@ -1115,7 +1115,9 @@ mod thread_processor_behavior_tests {
|
||||
let connection = ConnectionId(1);
|
||||
let (cancel_tx, cancel_rx) = oneshot::channel();
|
||||
|
||||
manager.connection_initialized(connection).await;
|
||||
manager
|
||||
.connection_initialized(connection, ConnectionCapabilities::default())
|
||||
.await;
|
||||
manager
|
||||
.try_ensure_connection_subscribed(
|
||||
thread_id, connection, /*experimental_raw_events*/ false,
|
||||
@@ -1158,8 +1160,12 @@ mod thread_processor_behavior_tests {
|
||||
let connection_b = ConnectionId(2);
|
||||
let (cancel_tx, mut cancel_rx) = oneshot::channel();
|
||||
|
||||
manager.connection_initialized(connection_a).await;
|
||||
manager.connection_initialized(connection_b).await;
|
||||
manager
|
||||
.connection_initialized(connection_a, ConnectionCapabilities::default())
|
||||
.await;
|
||||
manager
|
||||
.connection_initialized(connection_b, ConnectionCapabilities::default())
|
||||
.await;
|
||||
manager
|
||||
.try_ensure_connection_subscribed(
|
||||
thread_id,
|
||||
@@ -1203,8 +1209,12 @@ mod thread_processor_behavior_tests {
|
||||
let connection_a = ConnectionId(1);
|
||||
let connection_b = ConnectionId(2);
|
||||
|
||||
manager.connection_initialized(connection_a).await;
|
||||
manager.connection_initialized(connection_b).await;
|
||||
manager
|
||||
.connection_initialized(connection_a, ConnectionCapabilities::default())
|
||||
.await;
|
||||
manager
|
||||
.connection_initialized(connection_b, ConnectionCapabilities::default())
|
||||
.await;
|
||||
manager
|
||||
.try_ensure_connection_subscribed(
|
||||
thread_id,
|
||||
@@ -1249,7 +1259,9 @@ mod thread_processor_behavior_tests {
|
||||
let thread_id = ThreadId::from_string("ad7f0408-99b8-4f6e-a46f-bd0eec433370")?;
|
||||
let connection = ConnectionId(1);
|
||||
|
||||
manager.connection_initialized(connection).await;
|
||||
manager
|
||||
.connection_initialized(connection, ConnectionCapabilities::default())
|
||||
.await;
|
||||
let threads_to_unload = manager.remove_connection(connection).await;
|
||||
assert_eq!(threads_to_unload, Vec::<ThreadId>::new());
|
||||
|
||||
@@ -1264,4 +1276,79 @@ mod thread_processor_behavior_tests {
|
||||
assert!(!manager.has_subscribers(thread_id).await);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn first_attestation_capable_connection_for_thread_only_uses_thread_subscribers()
|
||||
-> Result<()> {
|
||||
let manager = ThreadStateManager::new();
|
||||
let thread_id = ThreadId::from_string("dfbd9a95-2f44-470a-8bd8-1cfc04efc243")?;
|
||||
let other_thread_id = ThreadId::from_string("6c9a74e4-5e59-479e-90bf-5c5798bb50aa")?;
|
||||
let unrelated_supported_connection = ConnectionId(1);
|
||||
let earlier_supported_connection = ConnectionId(2);
|
||||
let later_supported_connection = ConnectionId(3);
|
||||
let unsupported_connection = ConnectionId(4);
|
||||
|
||||
manager
|
||||
.connection_initialized(
|
||||
unrelated_supported_connection,
|
||||
ConnectionCapabilities {
|
||||
request_attestation: true,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
manager
|
||||
.connection_initialized(
|
||||
earlier_supported_connection,
|
||||
ConnectionCapabilities {
|
||||
request_attestation: true,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
manager
|
||||
.connection_initialized(
|
||||
later_supported_connection,
|
||||
ConnectionCapabilities {
|
||||
request_attestation: true,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
manager
|
||||
.connection_initialized(unsupported_connection, ConnectionCapabilities::default())
|
||||
.await;
|
||||
|
||||
assert!(
|
||||
manager
|
||||
.try_add_connection_to_thread(other_thread_id, unrelated_supported_connection)
|
||||
.await
|
||||
);
|
||||
assert!(
|
||||
manager
|
||||
.try_add_connection_to_thread(thread_id, later_supported_connection)
|
||||
.await
|
||||
);
|
||||
assert!(
|
||||
manager
|
||||
.try_add_connection_to_thread(thread_id, earlier_supported_connection)
|
||||
.await
|
||||
);
|
||||
assert!(
|
||||
manager
|
||||
.try_add_connection_to_thread(thread_id, unsupported_connection)
|
||||
.await
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
manager
|
||||
.first_attestation_capable_connection_for_thread(thread_id)
|
||||
.await,
|
||||
Some(earlier_supported_connection)
|
||||
);
|
||||
assert_eq!(
|
||||
manager
|
||||
.first_attestation_capable_connection_for_thread(other_thread_id)
|
||||
.await,
|
||||
Some(unrelated_supported_connection)
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -199,11 +199,16 @@ impl ThreadEntry {
|
||||
|
||||
#[derive(Default)]
|
||||
struct ThreadStateManagerInner {
|
||||
live_connections: HashSet<ConnectionId>,
|
||||
live_connections: HashMap<ConnectionId, ConnectionCapabilities>,
|
||||
threads: HashMap<ThreadId, ThreadEntry>,
|
||||
thread_ids_by_connection: HashMap<ConnectionId, HashSet<ThreadId>>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Default)]
|
||||
pub(crate) struct ConnectionCapabilities {
|
||||
pub(crate) request_attestation: bool,
|
||||
}
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
pub(crate) struct ThreadStateManager {
|
||||
state: Arc<Mutex<ThreadStateManagerInner>>,
|
||||
@@ -214,12 +219,36 @@ impl ThreadStateManager {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
pub(crate) async fn connection_initialized(&self, connection_id: ConnectionId) {
|
||||
pub(crate) async fn connection_initialized(
|
||||
&self,
|
||||
connection_id: ConnectionId,
|
||||
capabilities: ConnectionCapabilities,
|
||||
) {
|
||||
self.state
|
||||
.lock()
|
||||
.await
|
||||
.live_connections
|
||||
.insert(connection_id);
|
||||
.insert(connection_id, capabilities);
|
||||
}
|
||||
|
||||
pub(crate) async fn first_attestation_capable_connection_for_thread(
|
||||
&self,
|
||||
thread_id: ThreadId,
|
||||
) -> Option<ConnectionId> {
|
||||
let state = self.state.lock().await;
|
||||
state
|
||||
.threads
|
||||
.get(&thread_id)?
|
||||
.connection_ids
|
||||
.iter()
|
||||
.filter_map(|connection_id| {
|
||||
state
|
||||
.live_connections
|
||||
.get(connection_id)?
|
||||
.request_attestation
|
||||
.then_some(*connection_id)
|
||||
})
|
||||
.min_by_key(|connection_id| connection_id.0)
|
||||
}
|
||||
|
||||
pub(crate) async fn subscribed_connection_ids(&self, thread_id: ThreadId) -> Vec<ConnectionId> {
|
||||
@@ -338,7 +367,7 @@ impl ThreadStateManager {
|
||||
) -> Option<Arc<Mutex<ThreadState>>> {
|
||||
let thread_state = {
|
||||
let mut state = self.state.lock().await;
|
||||
if !state.live_connections.contains(&connection_id) {
|
||||
if !state.live_connections.contains_key(&connection_id) {
|
||||
return None;
|
||||
}
|
||||
state
|
||||
@@ -366,7 +395,7 @@ impl ThreadStateManager {
|
||||
connection_id: ConnectionId,
|
||||
) -> bool {
|
||||
let mut state = self.state.lock().await;
|
||||
if !state.live_connections.contains(&connection_id) {
|
||||
if !state.live_connections.contains_key(&connection_id) {
|
||||
return false;
|
||||
}
|
||||
state
|
||||
|
||||
Reference in New Issue
Block a user