use std::future::Future;
use std::path::Path;
use std::sync::Arc;
use std::time::Duration;
use microvms_edges::identity::HandshakeState;
use crate::adapters::SystemAdapters;
use crate::agents::{AgentSpec, AgentVm, require_specs};
use crate::clock::{Clock, SystemClock};
use crate::control::transport::{SignedTransport, Transport};
use crate::control::{BuildContext, ControlPlane};
use crate::cost::CalendarDate;
use crate::entropy::OsEntropy;
use crate::error::Error;
use crate::identity::{LaunchIdentity, TunnelIdentity};
use crate::names::{NameRecord, NameStore};
use crate::region::Region;
use crate::sandbox::{ControlPlaneMinter, Sandbox};
use crate::session::{ProxyAuth, Session, SessionBuilder, TokenMinter};
pub trait CalendarDateExt {
fn today_utc() -> CalendarDate;
}
impl CalendarDateExt for CalendarDate {
fn today_utc() -> CalendarDate {
SystemClock::new()
.today_utc()
.unwrap_or(CalendarDate::from_ymd(1970, 1, 1))
}
}
pub trait NameRecordExt: Sized {
fn new(
name: impl Into<String>,
microvm_id: impl Into<String>,
endpoint: impl Into<String>,
agent_token: impl Into<String>,
region: impl Into<String>,
) -> Result<Self, Error>;
}
impl NameRecordExt for NameRecord {
fn new(
name: impl Into<String>,
microvm_id: impl Into<String>,
endpoint: impl Into<String>,
agent_token: impl Into<String>,
region: impl Into<String>,
) -> Result<Self, Error> {
let at = SystemClock::new().unix_now().as_secs();
NameRecord::new_at(name, microvm_id, endpoint, agent_token, region, at)
}
}
pub trait LaunchIdentityExt: Sized {
fn generate() -> Result<Self, Error>;
}
impl LaunchIdentityExt for LaunchIdentity {
fn generate() -> Result<Self, Error> {
crate::entropy::launch_identity(&OsEntropy)
}
}
pub trait TunnelIdentityExt {
fn initiator(&self) -> Result<HandshakeState, Error>;
}
impl TunnelIdentityExt for TunnelIdentity {
fn initiator(&self) -> Result<HandshakeState, Error> {
microvms_edges::identity::initiator(self)
}
}
pub trait ControlPlaneExt: Sized {
fn new(region: Region) -> impl Future<Output = Result<Self, Error>> + Send;
fn with_transport(transport: Arc<dyn Transport>, region: Region, clock: Arc<dyn Clock>)
-> Self;
}
impl ControlPlaneExt for ControlPlane {
async fn new(region: Region) -> Result<Self, Error> {
let transport = SignedTransport::new(region.clone()).await?;
Ok(Self::with_transport(
Arc::new(transport),
region,
Arc::new(SystemClock::new()),
))
}
fn with_transport(
transport: Arc<dyn Transport>,
region: Region,
clock: Arc<dyn Clock>,
) -> Self {
ControlPlane::from_ports(
transport,
region,
clock,
Arc::new(OsEntropy),
Arc::new(SystemAdapters),
)
}
}
async fn plane_for(region: Region, port: Option<u16>) -> Result<ControlPlane, Error> {
let control = ControlPlane::new(region).await?;
match port {
Some(port) => control.with_port(port),
None => Ok(control),
}
}
pub trait SandboxExt: Sized {
fn new(region: Region) -> impl Future<Output = Result<Self, Error>> + Send;
fn adopt_in(
region: Region,
microvm_id: impl Into<String>,
endpoint: impl Into<String>,
agent_token: impl Into<String>,
port: Option<u16>,
) -> impl Future<Output = Result<Self, Error>> + Send;
fn from_name(
store: &dyn NameStore,
name: &str,
region: Option<Region>,
port: Option<u16>,
) -> impl Future<Output = Result<Self, Error>> + Send;
}
impl SandboxExt for Sandbox {
async fn new(region: Region) -> Result<Self, Error> {
Ok(Self::with_control_plane(ControlPlane::new(region).await?))
}
fn adopt_in(
region: Region,
microvm_id: impl Into<String>,
endpoint: impl Into<String>,
agent_token: impl Into<String>,
port: Option<u16>,
) -> impl Future<Output = Result<Self, Error>> + Send {
let (microvm_id, endpoint, agent_token) =
(microvm_id.into(), endpoint.into(), agent_token.into());
async move {
Self::adopt(
plane_for(region, port).await?,
microvm_id,
endpoint,
agent_token,
)
.await
}
}
async fn from_name(
store: &dyn NameStore,
name: &str,
region: Option<Region>,
port: Option<u16>,
) -> Result<Self, Error> {
let record = crate::names::resolve(store, name, region.as_ref())?;
let control = plane_for(record.region(), port).await?;
Self::adopt_record(control, record).await
}
}
pub trait AgentVmExt: Sized {
fn from_name(
store: &dyn NameStore,
specs: Vec<AgentSpec>,
name: &str,
region: Option<Region>,
port: Option<u16>,
) -> impl Future<Output = Result<Self, Error>> + Send;
fn adopt_in(
region: Region,
specs: Vec<AgentSpec>,
microvm_id: impl Into<String>,
endpoint: impl Into<String>,
agent_token: impl Into<String>,
port: Option<u16>,
) -> impl Future<Output = Result<Self, Error>> + Send;
}
impl AgentVmExt for AgentVm {
async fn from_name(
store: &dyn NameStore,
specs: Vec<AgentSpec>,
name: &str,
region: Option<Region>,
port: Option<u16>,
) -> Result<Self, Error> {
require_specs(&specs)?;
let sandbox = Sandbox::from_name(store, name, region, port).await?;
AgentVm::new(sandbox, specs)
}
fn adopt_in(
region: Region,
specs: Vec<AgentSpec>,
microvm_id: impl Into<String>,
endpoint: impl Into<String>,
agent_token: impl Into<String>,
port: Option<u16>,
) -> impl Future<Output = Result<Self, Error>> + Send {
let (microvm_id, endpoint, agent_token) =
(microvm_id.into(), endpoint.into(), agent_token.into());
async move {
require_specs(&specs)?;
let sandbox =
Sandbox::adopt_in(region, microvm_id, endpoint, agent_token, port).await?;
AgentVm::new(sandbox, specs)
}
}
}
pub trait SessionExt: Sized {
fn connect(
endpoint: impl Into<String>,
agent_token: impl Into<String>,
minter: Arc<dyn TokenMinter>,
) -> impl Future<Output = Result<Self, Error>> + Send;
fn attach(
region: Region,
microvm_id: impl Into<String>,
endpoint: impl Into<String>,
agent_token: impl Into<String>,
port: Option<u16>,
request_timeout: Option<Duration>,
) -> impl Future<Output = Result<Self, Error>> + Send;
fn direct(endpoint: impl Into<String>, agent_token: impl Into<String>) -> Result<Self, Error>;
}
impl SessionExt for Session {
fn connect(
endpoint: impl Into<String>,
agent_token: impl Into<String>,
minter: Arc<dyn TokenMinter>,
) -> impl Future<Output = Result<Self, Error>> + Send {
let session = Session::builder(endpoint, agent_token)
.with_minter(minter)
.build();
std::future::ready(session)
}
fn attach(
region: Region,
microvm_id: impl Into<String>,
endpoint: impl Into<String>,
agent_token: impl Into<String>,
port: Option<u16>,
request_timeout: Option<Duration>,
) -> impl Future<Output = Result<Self, Error>> + Send {
let (microvm_id, endpoint, agent_token) =
(microvm_id.into(), endpoint.into(), agent_token.into());
async move {
if request_timeout.is_some_and(|timeout| timeout.is_zero()) {
return Err(Error::invalid_arg("request timeout must be positive"));
}
let control = plane_for(region, port).await?;
let port = control.port();
let minter = Arc::new(ControlPlaneMinter::new(Arc::new(control), microvm_id));
let mut builder = Session::builder(endpoint, agent_token)
.with_minter(minter)
.with_port(port);
if let Some(timeout) = request_timeout {
builder = builder.with_timeout(timeout);
}
builder.build()
}
}
fn direct(endpoint: impl Into<String>, agent_token: impl Into<String>) -> Result<Self, Error> {
Session::builder(endpoint, agent_token).build()
}
}
pub trait SessionBuilderExt {
fn build(self) -> Result<Session, Error>;
}
impl SessionBuilderExt for SessionBuilder {
fn build(self) -> Result<Session, Error> {
self.build_with(&SystemAdapters)
}
}
pub trait ProxyAuthExt {
fn new(minter: Arc<dyn TokenMinter>, port: u16) -> Self;
}
impl ProxyAuthExt for ProxyAuth {
fn new(minter: Arc<dyn TokenMinter>, port: u16) -> Self {
ProxyAuth::with_clock(minter, port, Arc::new(SystemClock::new()))
}
}
pub trait BuildContextExt: Sized {
fn from_dir(dir: impl AsRef<Path>) -> Result<Self, Error>;
}
impl BuildContextExt for BuildContext {
fn from_dir(dir: impl AsRef<Path>) -> Result<Self, Error> {
crate::control::context::from_dir(dir)
}
}
pub trait LedgerExt {
fn new(region: &str, root: &std::path::Path) -> Self;
fn open_all_for_vm(root: &std::path::Path, microvm_id: &str) -> Vec<Self>
where
Self: Sized;
}
impl LedgerExt for microvms_app::ledger::Ledger {
fn new(region: &str, root: &std::path::Path) -> Self {
crate::ledger::new(region, root)
}
fn open_all_for_vm(root: &std::path::Path, microvm_id: &str) -> Vec<Self> {
crate::ledger::open_all_for_vm(root, microvm_id)
}
}
pub trait HistoryExt {
fn for_vm(root: &std::path::Path, vm_id: &str) -> Self;
}
impl HistoryExt for microvms_app::history::History {
fn for_vm(root: &std::path::Path, vm_id: &str) -> Self {
crate::history::for_vm(root, vm_id)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::error::ErrorKind;
#[test]
fn with_transport_draws_from_the_os_pool() {
let plane = || {
ControlPlane::with_transport(
Arc::new(microvms_app::testing::FakeControlPlane::new()),
Region::UsEast1,
Arc::new(microvms_app::testing::TestClock::new()),
)
};
let draw = |plane: &ControlPlane| {
let mut bytes = [0_u8; 16];
plane
.entropy()
.fill(&mut bytes)
.expect("the pool is available");
bytes
};
assert_ne!(draw(&plane()), draw(&plane()));
}
#[tokio::test]
async fn from_name_refuses_a_foreign_or_missing_name_before_any_plane() {
struct One(NameRecord);
impl NameStore for One {
fn get(&self, name: &str) -> Result<Option<NameRecord>, Error> {
Ok((name == self.0.name).then(|| self.0.clone()))
}
fn put(&self, _: &NameRecord) -> Result<(), Error> {
unreachable!()
}
fn delete(&self, _: &str) -> Result<bool, Error> {
unreachable!()
}
fn list(&self) -> Result<Vec<NameRecord>, Error> {
Ok(vec![self.0.clone()])
}
}
const TOKEN: &str = "from-name-canary-token";
let record = NameRecord::new(
"ci",
"mvm-abc123",
"https://mvm-abc123.microvm.us-west-2.amazonaws.com",
TOKEN,
"us-west-2",
)
.expect("record");
let store = One(record);
let foreign = Sandbox::from_name(&store, "ci", Some(Region::UsEast1), None)
.await
.expect_err("registered in another region");
assert_eq!(foreign.kind(), ErrorKind::InvalidArg);
let missing = Sandbox::from_name(&store, "nope", None, None)
.await
.expect_err("no such name");
assert_eq!(missing.kind(), ErrorKind::Precondition);
let specs = vec![AgentSpec::new(crate::agents::Agent::Codex)];
let agent = AgentVm::from_name(&store, specs, "ci", Some(Region::UsEast1), None)
.await
.expect_err("registered in another region");
assert_eq!(agent.kind(), ErrorKind::InvalidArg);
assert!(!format!("{foreign} {missing} {agent}").contains(TOKEN));
}
#[test]
fn the_production_build_fills_the_backend_and_the_clock() {
let session = Session::direct("http://127.0.0.1:9", "token").expect("a direct session");
assert!(session.proxy_auth().is_none());
}
#[test]
fn proxy_auth_new_keeps_the_defaults() {
let auth = ProxyAuth::new(
Arc::new(microvms_app::testing::CountingMinter::default()),
crate::session::DEFAULT_AGENT_PORT,
);
assert_eq!(auth.port(), crate::session::DEFAULT_AGENT_PORT);
assert_eq!(auth.mint_count(), 0);
assert!(!auth.is_cached());
let described = format!("{auth:?}");
assert!(
described.contains(&format!(
"refresh_after: {:?}",
crate::session::DEFAULT_REFRESH_AFTER
)),
"{described}"
);
}
}