File size: 3,702 Bytes
afa0cbf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | use std::sync::Arc;
use std::sync::RwLock;
use codex_app_server_protocol::ChatgptAuthTokensRefreshParams;
use codex_app_server_protocol::ChatgptAuthTokensRefreshReason;
use codex_app_server_protocol::ChatgptAuthTokensRefreshResponse;
use codex_app_server_protocol::ServerRequestPayload;
use codex_login::CodexAuth;
use codex_login::ExternalAuthFuture;
use codex_login::auth::ExternalAuth;
use codex_login::auth::ExternalAuthRefreshContext;
use codex_login::auth::ExternalAuthRefreshReason;
use tokio::time::Duration;
use tokio::time::timeout;
use crate::outgoing_message::OutgoingMessageSender;
const EXTERNAL_AUTH_REFRESH_TIMEOUT: Duration = Duration::from_secs(10);
pub(crate) struct ExternalAuthBridge {
outgoing: Arc<OutgoingMessageSender>,
auth: RwLock<CodexAuth>,
}
impl ExternalAuthBridge {
pub(crate) fn new(outgoing: Arc<OutgoingMessageSender>, auth: CodexAuth) -> Self {
Self {
outgoing,
auth: RwLock::new(auth),
}
}
async fn refresh(&self, context: ExternalAuthRefreshContext) -> std::io::Result<CodexAuth> {
let reason = match context.reason {
ExternalAuthRefreshReason::Unauthorized => ChatgptAuthTokensRefreshReason::Unauthorized,
};
let params = ChatgptAuthTokensRefreshParams {
reason,
previous_account_id: context.previous_account_id,
};
let (request_id, rx) = self
.outgoing
.send_request(ServerRequestPayload::ChatgptAuthTokensRefresh(params))
.await;
let result = match timeout(EXTERNAL_AUTH_REFRESH_TIMEOUT, rx).await {
Ok(result) => {
let result = result.map_err(|err| {
std::io::Error::other(format!("auth refresh request canceled: {err}"))
})?;
result.map_err(|err| {
// Don't log err.message because it may contain a token.
let code = err.code;
std::io::Error::other(format!("auth refresh request failed: code={code}"))
})?
}
Err(_) => {
let _canceled = self.outgoing.cancel_request(&request_id).await;
return Err(std::io::Error::other(format!(
"auth refresh request timed out after {}s",
EXTERNAL_AUTH_REFRESH_TIMEOUT.as_secs()
)));
}
};
// Don't propagate parser error messages because they may contain a token.
let response: ChatgptAuthTokensRefreshResponse = serde_json::from_value(result)
.map_err(|_| std::io::Error::other("invalid auth refresh response"))?;
let auth = CodexAuth::from_external_chatgpt_tokens(
response.access_token.as_str(),
response.chatgpt_account_id.as_str(),
response.chatgpt_plan_type.as_deref(),
)
.map_err(|err| {
std::io::Error::new(err.kind(), "auth refresh returned invalid credentials")
})?;
*self
.auth
.write()
.map_err(|_| std::io::Error::other("external auth lock is poisoned"))? = auth.clone();
Ok(auth)
}
}
impl ExternalAuth for ExternalAuthBridge {
fn resolve(&self) -> ExternalAuthFuture<'_, CodexAuth> {
Box::pin(async {
self.auth
.read()
.map(|auth| auth.clone())
.map_err(|_| std::io::Error::other("external auth lock is poisoned"))
})
}
fn refresh(&self, context: ExternalAuthRefreshContext) -> ExternalAuthFuture<'_, CodexAuth> {
Box::pin(ExternalAuthBridge::refresh(self, context))
}
}
|