Skip to content
This repository has been archived by the owner on Nov 1, 2023. It is now read-only.

Enable unmanaged registrations and configuration by environment variables #318

Merged
merged 6 commits into from
Nov 18, 2020
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions src/agent/onefuzz-agent/src/tasks/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -130,8 +130,9 @@ impl Config {
telemetry::set_property(EventData::MachineId(get_machine_id().await?));
telemetry::set_property(EventData::Version(env!("ONEFUZZ_VERSION").to_string()));
telemetry::set_property(EventData::InstanceId(self.common().instance_id));
if let Ok(scaleset) = get_scaleset_name().await {
telemetry::set_property(EventData::ScalesetId(scaleset));
let scaleset = get_scaleset_name().await?;
if let Some(scaleset_name) = &scaleset {
telemetry::set_property(EventData::ScalesetId(scaleset_name.to_string()));
}

info!("agent ready, dispatching task");
Expand Down
32 changes: 8 additions & 24 deletions src/agent/onefuzz-supervisor/src/auth.rs
Original file line number Diff line number Diff line change
Expand Up @@ -106,13 +106,18 @@ impl ClientCredentials {
let mut url = Url::parse("https://login.microsoftonline.com")?;
url.path_segments_mut()
.expect("Authority URL is cannot-be-a-base")
.push(&self.tenant)
.push("oauth2/v2.0/token");
.extend(&[&self.tenant, "oauth2", "v2.0", "token"]);

let response = reqwest::Client::new()
.post(url)
.header("Content-Length", "0")
.form(&self.form_data())
.form(&[
("client_id", self.client_id.to_hyphenated().to_string()),
("client_secret", self.client_secret.expose_ref().to_string()),
("grant_type", "client_credentials".into()),
("tenant", self.tenant.clone()),
("scope", format!("{}.default", self.resource)),
])
.send_retry_default()
.await?
.error_for_status_with_body()
Expand All @@ -122,27 +127,6 @@ impl ClientCredentials {

Ok(body.into())
}

fn form_data(&self) -> FormData {
let scope = format!("{}/.default", self.resource);

FormData {
client_id: self.client_id,
client_secret: self.client_secret.clone(),
grant_type: "client_credentials".into(),
scope,
tenant: self.tenant.clone(),
}
}
}

#[derive(Clone, Deserialize, Eq, PartialEq, Serialize)]
struct FormData {
client_id: Uuid,
client_secret: Secret<String>,
grant_type: String,
scope: String,
tenant: String,
}

// See: https://docs.microsoft.com/en-us/azure/active-directory/develop
Expand Down
58 changes: 55 additions & 3 deletions src/agent/onefuzz-supervisor/src/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -80,11 +80,56 @@ impl StaticConfig {
Ok(config)
}

pub async fn load(config_path: impl AsRef<Path>) -> Result<Self> {
let data = tokio::fs::read(config_path).await?;
pub fn from_file(config_path: impl AsRef<Path>) -> Result<Self> {
let data = std::fs::read(config_path)?;
Self::new(&data)
}

pub fn from_env() -> Result<Self> {
let instance_id = Uuid::parse_str(&std::env::var("ONEFUZZ_INSTANCE_ID")?)?;
let client_id = Uuid::parse_str(&std::env::var("ONEFUZZ_CLIENT_ID")?)?;
let client_secret = std::env::var("ONEFUZZ_CLIENT_SECRET")?.into();
let tenant = std::env::var("ONEFUZZ_TENANT")?;
let onefuzz_url = Url::parse(&std::env::var("ONEFUZZ_URL")?)?;
let pool_name = std::env::var("ONEFUZZ_POOL")?;

let heartbeat_queue = if let Ok(key) = std::env::var("ONEFUZZ_HEARTBEAT") {
Some(Url::parse(&key)?)
} else {
None
};

let instrumentation_key = if let Ok(key) = std::env::var("ONEFUZZ_INSTRUMENTATION_KEY") {
Some(Uuid::parse_str(&key)?)
} else {
None
};

let telemetry_key = if let Ok(key) = std::env::var("ONEFUZZ_TELEMETRY_KEY") {
Some(Uuid::parse_str(&key)?)
} else {
None
};

let credentials = ClientCredentials::new(
client_id,
client_secret,
onefuzz_url.clone().to_string(),
tenant,
)
.into();

Ok(Self {
instance_id,
credentials,
pool_name,
onefuzz_url,
instrumentation_key,
telemetry_key,
heartbeat_queue,
})
}

fn register_url(&self) -> Url {
let mut url = self.onefuzz_url.clone();
url.set_path("/api/agents/registration");
Expand Down Expand Up @@ -152,7 +197,14 @@ impl Registration {

if managed {
let scaleset = onefuzz::machine_id::get_scaleset_name().await?;
url.query_pairs_mut().append_pair("scaleset_id", &scaleset);
match scaleset {
Some(scaleset) => {
url.query_pairs_mut().append_pair("scaleset_id", &scaleset);
}
None => {
anyhow::bail!("managed instance without scaleset name");
}
}
}
// The registration can fail because this call is made before the virtual machine scaleset is done provisioning
// The authentication layer of the service will reject this request when that happens
Expand Down
41 changes: 27 additions & 14 deletions src/agent/onefuzz-supervisor/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -42,12 +42,13 @@ enum Opt {
Run(RunOpt),
Debug(debug::DebugOpt),
Licenses,
Version,
}

#[derive(StructOpt, Debug)]
struct RunOpt {
#[structopt(short, long = "--config", parse(from_os_str))]
config_path: PathBuf,
config_path: Option<PathBuf>,
}

fn main() -> Result<()> {
Expand All @@ -59,25 +60,29 @@ fn main() -> Result<()> {
Opt::Run(opt) => run(opt)?,
Opt::Debug(opt) => debug::debug(opt)?,
Opt::Licenses => licenses()?,
Opt::Version => version()?,
};

Ok(())
}

fn version() -> Result<()> {
println!(
bmc-msft marked this conversation as resolved.
Show resolved Hide resolved
"{} onefuzz:{} git:{}",
crate_version!(),
env!("ONEFUZZ_VERSION"),
env!("GIT_VERSION")
);
Ok(())
}

fn licenses() -> Result<()> {
use std::io::{self, Write};
io::stdout().write_all(include_bytes!("../../data/licenses.json"))?;
Ok(())
}

fn run(opt: RunOpt) -> Result<()> {
info!(
"{} onefuzz:{} git:{}",
crate_version!(),
env!("ONEFUZZ_VERSION"),
env!("GIT_VERSION")
);

if done::is_agent_done()? {
verbose!(
"agent is done, remove lock ({}) to continue",
Expand Down Expand Up @@ -110,9 +115,10 @@ fn run(opt: RunOpt) -> Result<()> {
fn load_config(opt: RunOpt) -> Result<StaticConfig> {
info!("loading supervisor agent config");

let data = std::fs::read(&opt.config_path)?;
let config = StaticConfig::new(&data)?;
verbose!("loaded static config from: {}", opt.config_path.display());
let config = match &opt.config_path {
Some(config_path) => StaticConfig::from_file(config_path)?,
None => StaticConfig::from_env()?,
};

init_telemetry(&config);

Expand All @@ -123,13 +129,20 @@ async fn run_agent(config: StaticConfig) -> Result<()> {
telemetry::set_property(EventData::InstanceId(config.instance_id));
telemetry::set_property(EventData::MachineId(get_machine_id().await?));
telemetry::set_property(EventData::Version(env!("ONEFUZZ_VERSION").to_string()));
if let Ok(scaleset) = get_scaleset_name().await {
telemetry::set_property(EventData::ScalesetId(scaleset));
let scaleset = get_scaleset_name().await?;
if let Some(scaleset_name) = &scaleset {
telemetry::set_property(EventData::ScalesetId(scaleset_name.to_string()));
}

let registration = match config::Registration::load_existing(config.clone()).await {
Ok(registration) => registration,
Err(_) => config::Registration::create_managed(config.clone()).await?,
Err(_) => {
if scaleset.is_some() {
config::Registration::create_managed(config.clone()).await?
} else {
config::Registration::create_unmanaged(config.clone()).await?
}
}
};
verbose!("current registration: {:?}", registration);

Expand Down
33 changes: 17 additions & 16 deletions src/agent/onefuzz/src/machine_id.rs
Original file line number Diff line number Diff line change
Expand Up @@ -60,24 +60,25 @@ pub async fn get_machine_name() -> Result<String> {
Ok(body)
}

pub async fn get_scaleset_name() -> Result<String> {
pub async fn get_scaleset_name() -> Result<Option<String>> {
let path = onefuzz_etc()?.join("scaleset_name");
let body = match fs::read_to_string(&path).await {
Ok(body) => body,
Err(_) => {
let resp = reqwest::Client::new()
.get(VM_SCALESET_NAME)
.timeout(Duration::from_millis(500))
.header("Metadata", "true")
.send_retry_default()
.await?;
let body = resp.text().await?;
write_file(path, &body).await?;
body
}
};
if let Ok(scaleset_name) = fs::read_to_string(&path).await {
return Ok(Some(scaleset_name));
}

Ok(body)
if let Ok(resp) = reqwest::Client::new()
.get(VM_SCALESET_NAME)
.timeout(Duration::from_millis(500))
.header("Metadata", "true")
.send_retry_default()
.await
{
let body = resp.text().await?;
write_file(path, &body).await?;
Ok(Some(body))
} else {
Ok(None)
}
}

#[cfg(target_os = "linux")]
Expand Down