use std::time::{Duration, Instant, SystemTime};
use crate::account::AccountCredential;
use crate::enroll::protocol::{
AuthorizeRequest, AuthorizeResponse, ClientHint, ErrorResponse, PollOutcome, PollState,
TokenResponse,
};
use crate::{EnrollError, TransportError};
#[derive(Debug, Clone)]
pub struct EnrollRequest {
pub name: String,
pub platform: String,
pub scopes: Vec<String>,
}
#[derive(Debug, Clone, PartialEq)]
pub struct Code {
pub user_code: String,
pub verification_uri: String,
pub expires_at: SystemTime,
}
#[derive(Debug)]
pub enum Progress {
Waiting { remaining: Duration },
Granted(AccountCredential),
Lapsed,
Denied,
}
pub struct Enrollment {
http: reqwest::Client,
base_url: String,
request: AuthorizeRequest,
round: Round,
}
struct Round {
device_code: String,
code: Code,
state: PollState,
deadline: Instant,
next_poll: Instant,
}
impl Round {
fn begin(authorization: AuthorizeResponse) -> Round {
let lifetime = Duration::from_secs(authorization.expires_in.max(0) as u64);
let now = Instant::now();
Round {
device_code: authorization.device_code,
code: Code {
user_code: authorization.user_code,
verification_uri: authorization.verification_uri,
expires_at: SystemTime::now() + lifetime,
},
state: PollState::new(authorization.interval),
deadline: now + lifetime,
next_poll: now,
}
}
}
pub async fn enroll(server_url: &str, request: EnrollRequest) -> Result<Enrollment, EnrollError> {
crate::tls::install_chosen_provider();
let http = reqwest::Client::builder()
.timeout(Duration::from_secs(30))
.user_agent(concat!("zyris/", env!("CARGO_PKG_VERSION")))
.build()
.map_err(unreachable_because)?;
let base_url = http_base(server_url);
let request = AuthorizeRequest {
name: request.name,
platform: request.platform,
scopes: request.scopes,
client_hint: local_client_hint(),
};
let round = Round::begin(authorize(&http, &base_url, &request).await?);
Ok(Enrollment { http, base_url, request, round })
}
impl Enrollment {
pub fn code(&self) -> &Code {
&self.round.code
}
pub async fn poll(&mut self) -> Result<Progress, EnrollError> {
let wait = self.round.next_poll.saturating_duration_since(Instant::now());
if !wait.is_zero() {
tokio::time::sleep(wait).await;
}
if Instant::now() >= self.round.deadline {
return Ok(Progress::Lapsed);
}
let response = post(
&self.http,
&self.base_url,
"/zyris/v1/device/token",
&serde_json::json!({
"device_code": self.round.device_code,
"grant_type": "urn:ietf:params:oauth:grant-type:device_code",
}),
)
.await?;
self.round.next_poll = Instant::now() + self.round.state.interval();
if response.status().is_success() {
let token: TokenResponse = response.json().await.map_err(unreachable_because)?;
return Ok(Progress::Granted(credential_from(&token)));
}
let error = parse_error(response).await;
match self.round.state.on_error(&error) {
PollOutcome::KeepWaiting(interval) => {
self.round.next_poll = Instant::now() + interval;
Ok(Progress::Waiting {
remaining: self.round.deadline.saturating_duration_since(Instant::now()),
})
}
PollOutcome::Expired => Ok(Progress::Lapsed),
PollOutcome::Denied => Ok(Progress::Denied),
PollOutcome::Fatal(message) => {
Err(EnrollError::Unreachable(TransportError::Io(message)))
}
}
}
pub async fn renew(&mut self) -> Result<(), EnrollError> {
let authorization = authorize(&self.http, &self.base_url, &self.request).await?;
self.round = Round::begin(authorization);
Ok(())
}
pub async fn wait(&mut self) -> Result<AccountCredential, EnrollError> {
loop {
match self.poll().await? {
Progress::Waiting { .. } => continue,
Progress::Granted(credential) => return Ok(credential),
Progress::Lapsed => return Err(EnrollError::Lapsed),
Progress::Denied => return Err(EnrollError::Denied),
}
}
}
}
async fn authorize(
http: &reqwest::Client,
base_url: &str,
request: &AuthorizeRequest,
) -> Result<AuthorizeResponse, EnrollError> {
let response = post(http, base_url, "/zyris/v1/device/authorize", request).await?;
if response.status().is_success() {
return response.json().await.map_err(unreachable_because);
}
let status = response.status().as_u16();
let body = response.text().await.unwrap_or_default();
if let Some(scope) = unknown_scope(&body) {
return Err(EnrollError::ScopeUnknown { scope });
}
Err(EnrollError::Unreachable(TransportError::Io(describe(&error_of(status, &body)))))
}
fn unknown_scope(body: &str) -> Option<String> {
let value: serde_json::Value = serde_json::from_str(body).ok()?;
if value.get("error")?.as_str()? != "unknown_scope" {
return None;
}
Some(value.get("scope")?.as_str()?.to_string())
}
pub(crate) fn credential_from(token: &TokenResponse) -> AccountCredential {
AccountCredential::new(
token.access_token.clone(),
token.refresh_token.clone(),
token.node_id.clone(),
token.node_name.clone(),
token.owner_email.clone(),
now_unix() + token.expires_in,
)
}
pub(crate) async fn post<T: serde::Serialize + ?Sized>(
http: &reqwest::Client,
base_url: &str,
path: &str,
body: &T,
) -> Result<reqwest::Response, EnrollError> {
http.post(format!("{base_url}{path}")).json(body).send().await.map_err(unreachable_because)
}
pub(crate) fn unreachable_because(error: reqwest::Error) -> EnrollError {
EnrollError::Unreachable(TransportError::Io(error.to_string()))
}
pub(crate) async fn parse_error(response: reqwest::Response) -> ErrorResponse {
let status = response.status().as_u16();
let body = response.text().await.unwrap_or_default();
error_of(status, &body)
}
fn error_of(status: u16, body: &str) -> ErrorResponse {
let mut error = serde_json::from_str::<ErrorResponse>(body).unwrap_or(ErrorResponse {
error: format!("http_{status}"),
error_description: None,
interval: None,
status: None,
});
error.status = Some(status);
error
}
fn describe(error: &ErrorResponse) -> String {
match &error.error_description {
Some(description) => format!("{}: {description}", error.error),
None => error.error.clone(),
}
}
pub(crate) fn http_base(server_url: &str) -> String {
let without_scheme = server_url
.strip_prefix("wss://")
.map(|rest| format!("https://{rest}"))
.or_else(|| server_url.strip_prefix("ws://").map(|rest| format!("http://{rest}")))
.unwrap_or_else(|| server_url.to_string());
match without_scheme.find("/zyris/") {
Some(index) => without_scheme[..index].to_string(),
None => without_scheme.trim_end_matches('/').to_string(),
}
}
fn local_client_hint() -> ClientHint {
ClientHint {
hostname: hostname(),
os: Some(std::env::consts::OS.to_string()),
agent: Some(concat!("zyris/", env!("CARGO_PKG_VERSION")).to_string()),
}
}
#[cfg(feature = "hostname")]
fn hostname() -> Option<String> {
crate::machine_name()
}
#[cfg(not(feature = "hostname"))]
fn hostname() -> Option<String> {
std::env::var("HOSTNAME")
.ok()
.filter(|v| !v.trim().is_empty())
.or_else(|| std::fs::read_to_string("/etc/hostname").ok().map(|v| v.trim().to_string()))
.filter(|v| !v.is_empty())
}
pub(crate) fn now_unix() -> i64 {
SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_secs() as i64)
.unwrap_or(0)
}
#[cfg(test)]
pub(crate) fn scripted_responder(
scripts: &'static [(&'static str, &'static str, &'static str)],
) -> String {
let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("no loopback port");
let address = listener.local_addr().expect("no local address");
std::thread::spawn(move || {
use std::io::{Read, Write};
let mut used: std::collections::HashMap<String, usize> = std::collections::HashMap::new();
for stream in listener.incoming() {
let Ok(mut stream) = stream else { return };
let mut buf = [0u8; 4096];
let Ok(n) = stream.read(&mut buf) else { continue };
let head = String::from_utf8_lossy(&buf[..n]).to_string();
let path = head.split_whitespace().nth(1).unwrap_or("").to_string();
let matching: Vec<&(&str, &str, &str)> =
scripts.iter().filter(|(suffix, _, _)| path.contains(suffix)).collect();
let (status, body) = if matching.is_empty() {
("404 Not Found", r#"{"error":"nothing scripted for this path"}"#)
} else {
let seen = used.entry(path).or_insert(0);
let (_, status, body) = *matching[(*seen).min(matching.len() - 1)];
*seen += 1;
(status, body)
};
let response = format!(
"HTTP/1.1 {status}\r\ncontent-type: application/json\r\n\
content-length: {}\r\nconnection: close\r\n\r\n{body}",
body.len()
);
let _ = stream.write_all(response.as_bytes());
let _ = stream.flush();
}
});
format!("ws://{address}/zyris/v1/ws")
}
#[cfg(test)]
mod tests {
use super::*;
const FIRST_CODE: &str = r#"{"device_code":"zdc_secret","user_code":"WXQR-7KBD","verification_uri":"https://attacca.example/settings/zyris/device","expires_in":600,"interval":1}"#;
const SECOND_CODE: &str = r#"{"device_code":"zdc_second","user_code":"HTPL-2FMR","verification_uri":"https://attacca.example/settings/zyris/device","expires_in":600,"interval":1}"#;
const GRANTED: &str = r#"{"access_token":"zna_new","refresh_token":"znr_new","expires_in":3600,"scope":"agents:read","node_id":"node-7","node_name":"hello node","owner_email":"allen@example.com"}"#;
const PENDING: &str = r#"{"error":"authorization_pending"}"#;
fn request() -> EnrollRequest {
EnrollRequest {
name: "hello node".into(),
platform: "linux".into(),
scopes: vec!["agents:read".into()],
}
}
#[tokio::test]
async fn the_code_comes_back_as_a_value_and_this_module_writes_to_no_console() {
let url = scripted_responder(&[("/device/authorize", "200 OK", FIRST_CODE)]);
let enrollment = enroll(&url, request()).await.expect("the server answered");
assert_eq!(enrollment.code().user_code, "WXQR-7KBD");
assert_eq!(
enrollment.code().verification_uri,
"https://attacca.example/settings/zyris/device"
);
assert!(
enrollment.code().expires_at > SystemTime::now(),
"a code that is already expired when it is handed over is not a code"
);
let marker = concat!("// zyris:", "nothing below this ships");
let mut parts = include_str!("device.rs").split(marker);
let shipped = parts.next().expect("splitting yields a first part even with no match");
assert!(
parts.next().is_some(),
"the marker bounding the shipping code is gone, so this scan read the whole file as \
test-only and asserted over nothing"
);
assert!(
parts.next().is_none(),
"the boundary is named twice, so the scan stops at whichever comes first"
);
assert!(
shipped.contains("pub async fn enroll"),
"the scan ended before this module's own entry point; the boundary has drifted upwards"
);
for console in [concat!("print", "ln!"), concat!("eprint", "ln!"), concat!("print", "!(")] {
assert!(!shipped.contains(console), "{console} writes where the caller cannot see it");
}
}
#[tokio::test]
async fn a_grant_nobody_has_answered_yet_says_how_long_the_code_has_left() {
let url = scripted_responder(&[
("/device/authorize", "200 OK", FIRST_CODE),
("/device/token", "400 Bad Request", PENDING),
]);
let mut enrollment = enroll(&url, request()).await.expect("the server answered");
match enrollment.poll().await.expect("pending is not a failure") {
Progress::Waiting { remaining } => assert!(
remaining > Duration::from_secs(500) && remaining <= Duration::from_secs(600),
"the countdown must be the code's, got {remaining:?}"
),
other => panic!("expected to still be waiting, got {other:?}"),
}
}
#[tokio::test]
async fn waiting_out_a_pending_grant_ends_with_the_account_credential() {
let url = scripted_responder(&[
("/device/authorize", "200 OK", FIRST_CODE),
("/device/token", "400 Bad Request", PENDING),
("/device/token", "200 OK", GRANTED),
]);
let mut enrollment = enroll(&url, request()).await.expect("the server answered");
let credential = enrollment.wait().await.expect("the person approved it");
assert_eq!(credential.access_token, "zna_new");
assert_eq!(credential.refresh_token, "znr_new");
assert_eq!(credential.node_id, "node-7");
assert_eq!(credential.owner_email, "allen@example.com");
assert!(
credential.bearer(now_unix(), 30).is_some(),
"a credential handed over already spent is not one"
);
}
#[tokio::test]
async fn a_lapsed_code_is_reported_and_no_fresh_one_is_fetched_behind_the_caller() {
let url = scripted_responder(&[
("/device/authorize", "200 OK", FIRST_CODE),
("/device/authorize", "200 OK", SECOND_CODE),
("/device/token", "400 Bad Request", r#"{"error":"expired_token"}"#),
]);
let mut enrollment = enroll(&url, request()).await.expect("the server answered");
assert!(
matches!(enrollment.poll().await.expect("lapsing is not a failure"), Progress::Lapsed),
"the caller has to be told the code stopped working"
);
assert_eq!(
enrollment.code().user_code,
"WXQR-7KBD",
"renewing is the caller's call; a second code must not appear on its own"
);
}
#[tokio::test]
async fn renewing_asks_for_a_new_code_rather_than_showing_the_lapsed_one_again() {
let url = scripted_responder(&[
("/device/authorize", "200 OK", FIRST_CODE),
("/device/authorize", "200 OK", SECOND_CODE),
]);
let mut enrollment = enroll(&url, request()).await.expect("the server answered");
assert_eq!(enrollment.code().user_code, "WXQR-7KBD");
enrollment.renew().await.expect("the server issued another");
assert_eq!(enrollment.code().user_code, "HTPL-2FMR");
assert!(enrollment.code().expires_at > SystemTime::now(), "and it is good for a while");
}
#[tokio::test]
async fn a_declined_request_is_reported_as_declined() {
let url = scripted_responder(&[
("/device/authorize", "200 OK", FIRST_CODE),
(
"/device/token",
"400 Bad Request",
r#"{"error":"access_denied","error_description":"the user declined"}"#,
),
]);
let mut enrollment = enroll(&url, request()).await.expect("the server answered");
assert!(matches!(
enrollment.poll().await.expect("a refusal is an answer"),
Progress::Denied
));
}
#[test]
fn http_base_is_derived_from_the_websocket_url() {
assert_eq!(http_base("wss://attacca.example/zyris/v1/ws"), "https://attacca.example");
assert_eq!(http_base("ws://127.0.0.1:8080/zyris/v1/ws"), "http://127.0.0.1:8080");
assert_eq!(http_base("https://attacca.example/"), "https://attacca.example");
assert_eq!(
http_base("wss://attacca.example:8443/zyris/v1/ws"),
"https://attacca.example:8443"
);
assert_eq!(
http_base(crate::DEFAULT_SERVER_URL),
"https://attacca.cc/api",
"the enrollment base must keep the path prefix the websocket URL carries"
);
}
#[test]
fn the_client_hint_describes_this_machine() {
let hint = local_client_hint();
assert_eq!(hint.os.as_deref(), Some(std::env::consts::OS));
assert!(hint.agent.as_deref().unwrap().starts_with("zyris/"));
}
#[test]
fn an_unparseable_error_body_still_names_the_status() {
assert_eq!(error_of(503, "<html>bad gateway</html>").error, "http_503");
assert_eq!(error_of(503, "<html>bad gateway</html>").status, Some(503));
}
async fn refusal_from(url: &str) -> EnrollError {
match enroll(url, request()).await {
Err(refusal) => refusal,
Ok(enrollment) => {
panic!("expected a refusal, got the code {}", enrollment.code().user_code)
}
}
}
#[tokio::test]
async fn a_scope_the_server_has_never_heard_of_comes_back_by_name() {
let url = scripted_responder(&[(
"/device/authorize",
"422 Unprocessable Entity",
r#"{"error":"unknown_scope","scope":"nodes:write"}"#,
)]);
let refusal = refusal_from(&url).await;
match refusal {
EnrollError::ScopeUnknown { scope } => assert_eq!(scope, "nodes:write"),
other => panic!("a caller cannot drop a scope it was not handed the name of: {other}"),
}
}
#[tokio::test]
async fn a_refusal_that_names_no_scope_is_not_reported_as_an_unknown_one() {
let url = scripted_responder(&[(
"/device/authorize",
"422 Unprocessable Entity",
r#"{"error":"invalid_request"}"#,
)]);
let refusal = refusal_from(&url).await;
assert!(
matches!(refusal, EnrollError::Unreachable(_)),
"only a refusal that names a scope may arrive as one, got {refusal}"
);
assert_eq!(unknown_scope(r#"{"error":"invalid_request"}"#), None);
assert_eq!(unknown_scope("<html>422</html>"), None);
}
}