use std::{future::Future, path::PathBuf};
use agent_client_protocol::{
schema::v1::{
AuthenticateResponse, CompleteElicitationNotification, CreateElicitationRequest,
ElicitationAction, ElicitationId, ElicitationRequestScope, ElicitationUrlMode, RequestId,
},
Error as AcpError,
};
use rho_providers::auth::{
browser::BrowserAvailability,
login_dispatch::{InteractiveLoginCompletion, ProviderAuthentication},
login_prompt::LoginPrompt,
};
use tokio::sync::watch;
use super::{auth::OAuthSignIn, AcpClientPort};
use crate::{
app::{bootstrap, config_repository::ConfigRepository},
credential_store::{self, AppCredentialStore},
};
pub(super) async fn authenticate(
auth: &str,
sign_in: OAuthSignIn,
config: Option<PathBuf>,
request_id: RequestId,
port: &dyn AcpClientPort,
) -> Result<AuthenticateResponse, AcpError> {
if !super::auth::runs_in_process(auth, sign_in) {
return Err(AcpError::invalid_params().data(format!(
"'{auth}' is not an authenticate method this connection advertised"
)));
}
open_credential_store(config)
.map_err(|error| AcpError::internal_error().data(format!("{error:#}")))?;
let mode = ProviderAuthentication::preferred_mode(auth, BrowserAvailability::from_process());
let login = ProviderAuthentication::start_interactive_login_with_availability(
auth,
mode,
BrowserAvailability::Headless,
)
.await
.map_err(|error| sign_in_error(auth, &error))?;
let label = login.provider_label;
let completion = match login.completion {
InteractiveLoginCompletion::Confirm(completion) => completion,
InteractiveLoginCompletion::Unconfirmed { .. } => {
return Err(
AcpError::internal_error().data(format!("{label} sign-in cannot be confirmed"))
);
}
};
let signed_in = async move {
completion
.await
.map_err(|error| sign_in_error(label, &error))?
.save(&AppCredentialStore)
.map_err(|error| {
AcpError::internal_error()
.data(format!("could not save {label} credentials: {error}"))
})
};
let shown = ShownLogin {
auth,
provider_label: label,
prompt: &login.prompt,
};
show_until_signed_in(port, request_id, shown, signed_in).await?;
Ok(AuthenticateResponse::new())
}
pub(super) struct LoginFlight {
running: tokio::sync::Mutex<()>,
latest: watch::Sender<u64>,
}
impl Default for LoginFlight {
fn default() -> Self {
Self {
running: tokio::sync::Mutex::new(()),
latest: watch::channel(0).0,
}
}
}
impl LoginFlight {
pub(super) async fn run<T>(
&self,
login: impl Future<Output = Result<T, AcpError>>,
) -> Result<T, AcpError> {
let mut generation = 0;
self.latest.send_modify(|latest| {
*latest += 1;
generation = *latest;
});
let mut newer = self.latest.subscribe();
let mut superseded = std::pin::pin!(async move {
let _ = newer.wait_for(|latest| *latest != generation).await;
});
let running = tokio::select! {
running = self.running.lock() => running,
() = &mut superseded => return Err(AcpError::request_cancelled()),
};
let result = tokio::select! {
result = login => result,
() = &mut superseded => Err(AcpError::request_cancelled()),
};
drop(running);
result
}
}
fn open_credential_store(config: Option<PathBuf>) -> anyhow::Result<()> {
let repository = ConfigRepository::new(config);
let config = repository.load()?;
if credential_store::needs_explicit_choice(&config) {
anyhow::bail!(
"credential store is unset; choose one with `rho credential-store set os|file`"
);
}
credential_store::activate_for_login(&config, &bootstrap::absolute_config_path(&repository)?)
}
#[derive(Clone, Copy)]
struct ShownLogin<'a> {
auth: &'a str,
provider_label: &'a str,
prompt: &'a LoginPrompt,
}
async fn show_until_signed_in(
port: &dyn AcpClientPort,
request_id: RequestId,
login: ShownLogin<'_>,
signed_in: impl Future<Output = Result<(), AcpError>>,
) -> Result<(), AcpError> {
let elicitation_id = ElicitationId::new(uuid::Uuid::new_v4().to_string());
let mut shown = port.create_elicitation(CreateElicitationRequest::new(
ElicitationUrlMode::new(
ElicitationRequestScope::new(request_id),
elicitation_id.clone(),
login.prompt.copyable_url(),
),
message(login),
));
let mut signed_in = std::pin::pin!(signed_in);
tokio::select! {
result = &mut signed_in => result?,
answer = &mut shown => match answer.map_err(|error| unshown(login, &error))?.action {
ElicitationAction::Accept(_) => signed_in.await?,
ElicitationAction::Decline | ElicitationAction::Cancel => {
return Err(AcpError::request_cancelled());
}
_ => return Err(AcpError::request_cancelled()),
},
}
if let Err(error) = port
.complete_elicitation(CompleteElicitationNotification::new(elicitation_id))
.await
{
tracing::warn!(%error, "could not mark the sign-in elicitation complete");
}
drop(shown);
Ok(())
}
fn message(login: ShownLogin<'_>) -> String {
let label = login.provider_label;
match &login.prompt.user_code {
Some(code) => format!("Sign in to {label} with code {code}."),
None => format!("Sign in to {label} in your browser."),
}
}
fn unshown(login: ShownLogin<'_>, error: &AcpError) -> AcpError {
AcpError::internal_error().data(format!(
"could not show the {} sign-in link: {error}; run `rho login {}` in a terminal instead",
login.provider_label, login.auth
))
}
fn sign_in_error(subject: &str, error: &impl std::fmt::Display) -> AcpError {
AcpError::internal_error().data(format!("could not sign in to {subject}: {error}"))
}
#[cfg(test)]
#[path = "auth_login_tests.rs"]
mod tests;