use std::net::SocketAddr;
use anyhow::Result;
use cli_shared::{ClientConfig, UserConfig};
use proto::{AuthToken, ProtocolError};
use crate::{credentials, grpc_hosted::HostedGrpcClient};
pub enum HostedAuthMode {
ConfigToken,
CredentialFallback,
}
pub struct HostedSession {
config: ClientConfig,
}
impl HostedSession {
pub fn build(
user_config: &UserConfig,
server_key: Option<String>,
mode: HostedAuthMode,
) -> Result<Self> {
let (token, credential_proof_key) = match mode {
HostedAuthMode::ConfigToken => (user_config.remote_token()?, None),
HostedAuthMode::CredentialFallback => {
let mut token = user_config.remote_token()?;
let mut credential_proof_key = None;
if token.is_none()
&& let Some(ref key) = server_key
&& let Ok(Some(cred)) = credentials::resolve_credential_for_server(key)
{
token = Some(AuthToken::new(cred.token, "credential-store"));
credential_proof_key = cred.private_key_pem;
}
(token, credential_proof_key)
}
};
let mut config = user_config.heddle_client_config(token)?;
if let Some(key) = server_key {
config = config.with_server_key(key);
}
if let Some(pem) = credential_proof_key
&& config.auth_proof_key_pem.is_none()
{
config = config.with_auth_proof_key_pem(pem);
}
Ok(Self { config })
}
pub async fn connect(&self, addr: SocketAddr) -> Result<HostedGrpcClient, ProtocolError> {
let mut client = HostedGrpcClient::connect(addr, &self.config).await?;
client.auto_rotate_if_needed().await;
Ok(client)
}
}
impl HostedGrpcClient {
pub async fn open_session(
addr: SocketAddr,
user_config: &UserConfig,
server_key: Option<String>,
mode: HostedAuthMode,
) -> Result<Self> {
Ok(HostedSession::build(user_config, server_key, mode)?
.connect(addr)
.await?)
}
}
#[cfg(test)]
mod tests {
#[test]
fn session_connect_rotates_credentials_after_connect() {
let source = include_str!("session.rs");
let connect_idx = source
.find("HostedGrpcClient::connect(addr, &self.config)")
.expect("session.rs must connect with the resolved addr");
let after_connect = &source[connect_idx..];
let rotate_offset = after_connect
.find("auto_rotate_if_needed")
.expect("auto_rotate_if_needed must appear in session.rs");
assert!(
rotate_offset < 400,
"auto_rotate_if_needed must follow HostedGrpcClient::connect within the \
same async block (found {rotate_offset} chars later)",
);
}
}