github/codex
4
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-02 00:18:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

app-server: persist device key bindings in sqlite (#19206)

Browse Source 
## Why

Device-key providers should only own platform key material. The
account/client binding used to authorize a signing payload is app-server
state, and keeping that state in provider-specific metadata makes the
same check harder to audit and harder to share across platform
implementations.

Persisting the binding in the shared state database gives the device-key
crate a platform-neutral source of truth before it asks a provider to
sign. It also lets app-server move potentially blocking key operations
off the main message processor path, which matters once providers may
wait for OS authentication prompts.

## What changed

- Add a `device_key_bindings` state migration plus `StateRuntime`
helpers keyed by `key_id`.
- Add an async `DeviceKeyBindingStore` abstraction to `codex-device-key`
and use it from `DeviceKeyStore::create` and `DeviceKeyStore::sign`.
- Keep provider calls behind async store methods and run the synchronous
provider work through `spawn_blocking`.
- Wire app-server device-key RPC handling to the SQLite-backed binding
store and spawn response/error delivery tasks for device-key requests.
- Run the turn-start tracing test on the existing larger current-thread
test harness after the larger async surface made the default test stack
too small locally.

## Validation

- `cargo test -p codex-device-key`
- `cargo test -p codex-state device_key`
- `cargo test -p codex-state`
- `cargo test -p codex-app-server device_key`
- `cargo test -p codex-app-server
message_processor::tracing_tests::turn_start_jsonrpc_span_parents_core_turn_spans`
- `cargo test -p codex-app-server`
- `just fix -p codex-device-key`
- `just fix -p codex-state`
- `just fix -p codex-app-server`
- `just bazel-lock-update`
- `just bazel-lock-check`
- `git diff --check`
This commit is contained in:
Ruslan Nigmatullin
2026-04-23 21:55:56 -07:00
committed by GitHub
parent e8d8080818
commit 19badb0be2
11 changed files with 622 additions and 258 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/app-server/src/device_key_api.rs
+95 -4
View File
@@ -1,5 +1,6 @@
use crate::error_code::INTERNAL_ERROR_CODE;
use crate::error_code::INVALID_REQUEST_ERROR_CODE;
use async_trait::async_trait;
use base64::Engine;
use base64::engine::general_purpose::STANDARD;
use codex_app_server_protocol::DeviceKeyAlgorithm;
@@ -13,6 +14,7 @@ use codex_app_server_protocol::DeviceKeySignPayload;
use codex_app_server_protocol::DeviceKeySignResponse;
use codex_app_server_protocol::JSONRPCErrorError;
use codex_device_key::DeviceKeyBinding;
use codex_device_key::DeviceKeyBindingStore;
use codex_device_key::DeviceKeyCreateRequest;
use codex_device_key::DeviceKeyError;
use codex_device_key::DeviceKeyGetPublicRequest;
@@ -24,14 +26,29 @@ use codex_device_key::RemoteControlClientConnectionAudience;
use codex_device_key::RemoteControlClientConnectionSignPayload;
use codex_device_key::RemoteControlClientEnrollmentAudience;
use codex_device_key::RemoteControlClientEnrollmentSignPayload;
use codex_state::DeviceKeyBindingRecord;
use codex_state::StateRuntime;
use std::fmt;
use std::path::PathBuf;
use std::sync::Arc;
use tokio::sync::OnceCell;
#[derive(Clone, Default)]
#[derive(Clone)]
pub(crate) struct DeviceKeyApi {
store: DeviceKeyStore,
}
impl DeviceKeyApi {
pub(crate) fn create(
pub(crate) fn new(sqlite_home: PathBuf, default_provider: String) -> Self {
Self {
store: DeviceKeyStore::new(Arc::new(StateDeviceKeyBindingStore::new(
sqlite_home,
default_provider,
))),
}
}
pub(crate) async fn create(
&self,
params: DeviceKeyCreateParams,
) -> Result<DeviceKeyCreateResponse, JSONRPCErrorError> {
@@ -44,11 +61,12 @@ impl DeviceKeyApi {
client_id: params.client_id,
},
})
.await
.map_err(map_device_key_error)?;
Ok(create_response_from_info(info))
}
pub(crate) fn public(
pub(crate) async fn public(
&self,
params: DeviceKeyPublicParams,
) -> Result<DeviceKeyPublicResponse, JSONRPCErrorError> {
@@ -57,11 +75,12 @@ impl DeviceKeyApi {
.get_public(DeviceKeyGetPublicRequest {
key_id: params.key_id,
})
.await
.map_err(map_device_key_error)?;
Ok(public_response_from_info(info))
}
pub(crate) fn sign(
pub(crate) async fn sign(
&self,
params: DeviceKeySignParams,
) -> Result<DeviceKeySignResponse, JSONRPCErrorError> {
@@ -71,6 +90,7 @@ impl DeviceKeyApi {
key_id: params.key_id,
payload: payload_from_params(params.payload),
})
.await
.map_err(map_device_key_error)?;
Ok(DeviceKeySignResponse {
signature_der_base64: STANDARD.encode(signature.signature_der),
@@ -80,6 +100,77 @@ impl DeviceKeyApi {
}
}
struct StateDeviceKeyBindingStore {
sqlite_home: PathBuf,
default_provider: String,
state_db: OnceCell<Arc<StateRuntime>>,
}
impl StateDeviceKeyBindingStore {
fn new(sqlite_home: PathBuf, default_provider: String) -> Self {
Self {
sqlite_home,
default_provider,
state_db: OnceCell::new(),
}
}
async fn state_db(&self) -> Result<Arc<StateRuntime>, DeviceKeyError> {
let sqlite_home = self.sqlite_home.clone();
let default_provider = self.default_provider.clone();
self.state_db
.get_or_try_init(|| async move {
StateRuntime::init(sqlite_home, default_provider)
.await
.map_err(|err| DeviceKeyError::Platform(err.to_string()))
})
.await
.cloned()
}
}
impl fmt::Debug for StateDeviceKeyBindingStore {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("StateDeviceKeyBindingStore")
.field("sqlite_home", &self.sqlite_home)
.field("default_provider", &self.default_provider)
.finish_non_exhaustive()
}
}
#[async_trait]
impl DeviceKeyBindingStore for StateDeviceKeyBindingStore {
async fn get_binding(&self, key_id: &str) -> Result<Option<DeviceKeyBinding>, DeviceKeyError> {
let state_db = self.state_db().await?;
state_db
.get_device_key_binding(key_id)
.await
.map(|record| {
record.map(|record| DeviceKeyBinding {
account_user_id: record.account_user_id,
client_id: record.client_id,
})
})
.map_err(|err| DeviceKeyError::Platform(err.to_string()))
}
async fn put_binding(
&self,
key_id: &str,
binding: &DeviceKeyBinding,
) -> Result<(), DeviceKeyError> {
let state_db = self.state_db().await?;
state_db
.upsert_device_key_binding(&DeviceKeyBindingRecord {
key_id: key_id.to_string(),
account_user_id: binding.account_user_id.clone(),
client_id: binding.client_id.clone(),
})
.await
.map_err(|err| DeviceKeyError::Platform(err.to_string()))
}
}
fn create_response_from_info(info: DeviceKeyInfo) -> DeviceKeyCreateResponse {
DeviceKeyCreateResponse {
key_id: info.key_id,
codex-rs/app-server/src/message_processor.rs
+55 -72
View File
@@ -325,7 +325,8 @@ impl MessageProcessor {
thread_manager.clone(),
analytics_events_client.clone(),
);
let device_key_api = DeviceKeyApi::default();
let device_key_api =
DeviceKeyApi::new(config.sqlite_home.clone(), config.model_provider_id.clone());
let external_agent_config_api =
ExternalAgentConfigApi::new(config.codex_home.to_path_buf());
let fs_api = FsApi::new(
@@ -882,8 +883,7 @@ impl MessageProcessor {
},
params,
device_key_requests_allowed,
)
.await;
);
}
ClientRequest::DeviceKeyPublic { request_id, params } => {
self.handle_device_key_public(
@@ -893,8 +893,7 @@ impl MessageProcessor {
},
params,
device_key_requests_allowed,
)
.await;
);
}
ClientRequest::DeviceKeySign { request_id, params } => {
self.handle_device_key_sign(
@@ -904,8 +903,7 @@ impl MessageProcessor {
},
params,
device_key_requests_allowed,
)
.await;
);
}
ClientRequest::FsReadFile { request_id, params } => {
self.handle_fs_read_file(
@@ -1173,96 +1171,81 @@ impl MessageProcessor {
}
}
async fn handle_device_key_create(
fn handle_device_key_create(
&self,
request_id: ConnectionRequestId,
params: DeviceKeyCreateParams,
device_key_requests_allowed: bool,
) {
if self
.reject_device_key_request_over_remote_transport(
request_id.clone(),
"device/key/create",
device_key_requests_allowed,
)
.await
{
return;
}
match self.device_key_api.create(params) {
Ok(response) => self.outgoing.send_response(request_id, response).await,
Err(error) => self.outgoing.send_error(request_id, error).await,
}
self.spawn_device_key_request(
request_id,
"device/key/create",
device_key_requests_allowed,
move |device_key_api| async move { device_key_api.create(params).await },
);
}
async fn handle_device_key_public(
fn handle_device_key_public(
&self,
request_id: ConnectionRequestId,
params: DeviceKeyPublicParams,
device_key_requests_allowed: bool,
) {
if self
.reject_device_key_request_over_remote_transport(
request_id.clone(),
"device/key/public",
device_key_requests_allowed,
)
.await
{
return;
}
match self.device_key_api.public(params) {
Ok(response) => self.outgoing.send_response(request_id, response).await,
Err(error) => self.outgoing.send_error(request_id, error).await,
}
self.spawn_device_key_request(
request_id,
"device/key/public",
device_key_requests_allowed,
move |device_key_api| async move { device_key_api.public(params).await },
);
}
async fn handle_device_key_sign(
fn handle_device_key_sign(
&self,
request_id: ConnectionRequestId,
params: DeviceKeySignParams,
device_key_requests_allowed: bool,
) {
if self
.reject_device_key_request_over_remote_transport(
request_id.clone(),
"device/key/sign",
device_key_requests_allowed,
)
.await
{
return;
}
match self.device_key_api.sign(params) {
Ok(response) => self.outgoing.send_response(request_id, response).await,
Err(error) => self.outgoing.send_error(request_id, error).await,
}
self.spawn_device_key_request(
request_id,
"device/key/sign",
device_key_requests_allowed,
move |device_key_api| async move { device_key_api.sign(params).await },
);
}
async fn reject_device_key_request_over_remote_transport(
fn spawn_device_key_request<R, F, Fut>(
&self,
request_id: ConnectionRequestId,
method: &str,
method: &'static str,
device_key_requests_allowed: bool,
) -> bool {
if device_key_requests_allowed {
return false;
}
run_request: F,
) where
R: serde::Serialize + Send + 'static,
F: FnOnce(DeviceKeyApi) -> Fut + Send + 'static,
Fut: Future<Output = Result<R, JSONRPCErrorError>> + Send + 'static,
{
let device_key_api = self.device_key_api.clone();
let outgoing = Arc::clone(&self.outgoing);
tokio::spawn(async move {
if !device_key_requests_allowed {
outgoing
.send_error(
request_id,
JSONRPCErrorError {
code: INVALID_REQUEST_ERROR_CODE,
message: format!("{method} is not available over remote transports"),
data: None,
},
)
.await;
return;
}
self.outgoing
.send_error(
request_id,
JSONRPCErrorError {
code: INVALID_REQUEST_ERROR_CODE,
message: format!("{method} is not available over remote transports"),
data: None,
},
)
.await;
true
match run_request(device_key_api).await {
Ok(response) => outgoing.send_response(request_id, response).await,
Err(error) => outgoing.send_error(request_id, error).await,
}
});
}
async fn handle_external_agent_config_detect(