initial public release

src/proxy-manager/.gitignore

@@ -0,0 +1 @@
+target

src/proxy-manager/Cargo.toml

@@ -0,0 +1,21 @@
+[package]
+name = "onefuzz-proxy-manager"
+version = "0.0.1"
+authors = ["fuzzing@microsoft.com"]
+edition = "2018"
+publish = false
+license = "MIT"
+
+[dependencies]
+anyhow = "1.0"
+clap = "2.33"
+env_logger = "0.7"
+futures = "0.3"
+log = "0.4"
+reqwest = { version = "0.10", features = ["json", "stream"] }
+serde = { version = "1.0", features = ["derive"] }
+serde_json = "1.0"
+storage-queue = { path = "../agent/storage-queue" }
+thiserror = "1.0"
+tokio = { version = "0.2", features = ["macros", "rt-threaded", "fs", "process"] }
+url = { version = "2.1", features = ["serde"] }

src/proxy-manager/build.rs

@@ -0,0 +1,36 @@
+use std::error::Error;
+use std::fs::File;
+use std::io::prelude::*;
+use std::process::Command;
+
+fn run_cmd(args: &[&str]) -> Result<String, Box<dyn Error>> {
+    let cmd = Command::new(args[0]).args(&args[1..]).output()?;
+    if cmd.status.success() {
+        Ok(String::from_utf8_lossy(&cmd.stdout).to_string())
+    } else {
+        Err(From::from("failed"))
+    }
+}
+
+fn read_file(filename: &str) -> Result<String, Box<dyn Error>> {
+    let mut file = File::open(filename)?;
+    let mut contents = String::new();
+    file.read_to_string(&mut contents)?;
+
+    Ok(contents)
+}
+
+fn main() -> Result<(), Box<dyn Error>> {
+    let sha = run_cmd(&["git", "rev-parse", "HEAD"])?;
+    let with_changes = if run_cmd(&["git", "diff", "--quiet"]).is_err() {
+        "-local_changes"
+    } else {
+        ""
+    };
+    println!("cargo:rustc-env=GIT_VERSION={}{}", sha, with_changes);
+
+    let version = read_file("../../CURRENT_VERSION")?;
+    println!("cargo:rustc-env=ONEFUZZ_VERSION={}", version);
+
+    Ok(())
+}

src/proxy-manager/data/licenses.json

@@ -0,0 +1 @@
+[]

src/proxy-manager/src/config.rs

@@ -0,0 +1,145 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+use crate::proxy;
+use anyhow::Result;
+use serde::{Deserialize, Serialize};
+use std::{fs::File, io::BufReader, path::PathBuf};
+use storage_queue::QueueClient;
+use thiserror::Error;
+use url::Url;
+
+#[derive(Error, Debug)]
+pub enum ProxyError {
+    #[error("missing argument {0}")]
+    MissingArg(String),
+
+    #[error("missing etag header")]
+    EtagError,
+
+    #[error("unable to open config file")]
+    FileError { source: std::io::Error },
+
+    #[error("unable to parse config file")]
+    ParseError { source: serde_json::error::Error },
+
+    #[error(transparent)]
+    Other {
+        #[from]
+        source: anyhow::Error,
+    },
+}
+
+#[derive(Debug, Deserialize, Serialize, PartialEq, Clone)]
+pub struct Forward {
+    pub src_port: u16,
+    pub dst_ip: String,
+    pub dst_port: u16,
+}
+
+#[derive(Debug, Deserialize, Serialize, PartialEq, Clone)]
+pub struct ConfigData {
+    pub region: String,
+    pub url: Url,
+    pub notification: Url,
+    pub forwards: Vec<Forward>,
+}
+
+#[derive(Debug, Deserialize, Serialize, PartialEq)]
+pub struct NotifyResponse<'a> {
+    pub region: &'a str,
+    pub forwards: Vec<Forward>,
+}
+
+#[derive(Debug, Deserialize, Serialize)]
+pub struct Config {
+    config_path: PathBuf,
+    data: ConfigData,
+    etag: Option<String>,
+}
+
+impl Config {
+    pub fn from_file(path: String) -> Result<Self> {
+        let config_path = PathBuf::from(&path);
+
+        let f = File::open(&config_path).map_err(|source| ProxyError::FileError { source })?;
+        let r = BufReader::new(f);
+        let data: ConfigData =
+            serde_json::from_reader(r).map_err(|source| ProxyError::ParseError { source })?;
+
+        Ok(Self {
+            config_path,
+            data,
+            etag: None,
+        })
+    }
+
+    async fn save(&self) -> Result<()> {
+        let encoded = serde_json::to_string(&self.data)?;
+        tokio::fs::write(&self.config_path, encoded).await?;
+        Ok(())
+    }
+
+    async fn fetch(&mut self) -> Result<bool> {
+        let mut request = reqwest::Client::new().get(self.data.url.clone());
+        if let Some(etag) = &self.etag {
+            request = request.header(reqwest::header::IF_NONE_MATCH, etag);
+        }
+        let response = request.send().await?;
+        let status = response.status();
+
+        if status == reqwest::StatusCode::NOT_MODIFIED {
+            return Ok(false);
+        }
+
+        if !status.is_success() {
+            if status.is_server_error() {
+                bail!("server error");
+            } else {
+                bail!("request failed: {:?}", status);
+            }
+        }
+
+        let etag = response
+            .headers()
+            .get(reqwest::header::ETAG)
+            .ok_or_else(|| ProxyError::EtagError)?
+            .to_str()?
+            .to_owned();
+        let data: ConfigData = response.json().await?;
+        self.etag = Some(etag);
+        if data != self.data {
+            self.data = data;
+            Ok(true)
+        } else {
+            Ok(false)
+        }
+    }
+
+    pub async fn notify(&self) -> Result<()> {
+        let client = QueueClient::new(self.data.notification.clone());
+
+        client
+            .enqueue(NotifyResponse {
+                region: &self.data.region,
+                forwards: self.data.forwards.clone(),
+            })
+            .await?;
+        Ok(())
+    }
+
+    pub async fn update(&mut self) -> Result<bool> {
+        if self.fetch().await? {
+            info!("config updated");
+            self.save().await?;
+        }
+
+        let notified = if proxy::update(&self.data).await? {
+            self.notify().await?;
+            true
+        } else {
+            false
+        };
+        Ok(notified)
+    }
+}

src/proxy-manager/src/main.rs

@@ -0,0 +1,73 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+#[macro_use]
+extern crate log;
+#[macro_use]
+extern crate anyhow;
+#[macro_use]
+extern crate clap;
+
+mod config;
+mod proxy;
+use anyhow::Result;
+use clap::{App, Arg, SubCommand};
+use config::{Config, ProxyError::MissingArg};
+use std::io::Write;
+
+const MINIMUM_NOTIFY_INTERVAL: tokio::time::Duration = std::time::Duration::from_secs(120);
+const POLL_INTERVAL: tokio::time::Duration = tokio::time::Duration::from_secs(5);
+
+async fn run(mut proxy_config: Config) -> Result<()> {
+    let mut last_updated = std::time::Instant::now();
+    loop {
+        info!("checking updates");
+        if proxy_config.update().await? {
+            last_updated = std::time::Instant::now();
+        } else if last_updated + MINIMUM_NOTIFY_INTERVAL < std::time::Instant::now() {
+            proxy_config.notify().await?;
+            last_updated = std::time::Instant::now();
+        }
+
+        tokio::time::delay_for(POLL_INTERVAL).await;
+    }
+}
+
+fn main() -> Result<()> {
+    env_logger::init();
+
+    let license_cmd = SubCommand::with_name("licenses").about("display third-party licenses");
+
+    let version = format!(
+        "{} onefuzz:{} git:{}",
+        crate_version!(),
+        env!("ONEFUZZ_VERSION"),
+        env!("GIT_VERSION")
+    );
+
+    let app = App::new("onefuzz-proxy")
+        .version(version.as_str())
+        .arg(
+            Arg::with_name("config")
+                .long("config")
+                .short("c")
+                .takes_value(true),
+        )
+        .subcommand(license_cmd);
+    let matches = app.get_matches();
+
+    if matches.subcommand_matches("licenses").is_some() {
+        std::io::stdout().write_all(include_bytes!("../data/licenses.json"))?;
+        return Ok(());
+    }
+
+    let config_path = matches
+        .value_of("config")
+        .ok_or_else(|| MissingArg("--config".to_string()))?
+        .parse()?;
+    let proxy = Config::from_file(config_path)?;
+
+    info!("parsed initial config");
+    let mut rt = tokio::runtime::Runtime::new()?;
+    rt.block_on(run(proxy))
+}

src/proxy-manager/src/proxy.rs

@@ -0,0 +1,133 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+use crate::config::ConfigData;
+use anyhow::Result;
+use futures::stream::StreamExt;
+use std::{collections::HashMap, path::Path};
+use tokio::process::Command;
+
+const SYSTEMD_CONFIG_DIR: &str = "/etc/systemd/system";
+const PROXY_PREFIX: &str = "onefuzz-proxy";
+
+fn build(data: &ConfigData) -> Result<HashMap<String, String>> {
+    let mut results = HashMap::new();
+
+    for entry in &data.forwards {
+        let socket_filename = format!("{}-{}.socket", PROXY_PREFIX, entry.src_port);
+        let service_filename = format!("{}-{}.service", PROXY_PREFIX, entry.src_port);
+        let socket = format!(
+            r##"
+[Socket]
+ListenStream=0.0.0.0:{}
+BindIPv6Only=both
+[Install]
+WantedBy=sockets.target
+"##,
+            entry.src_port
+        );
+        let service = format!(
+            r##"
+[Unit]
+Requires=onefuzz-proxy-{}.socket
+After=onefuzz-proxy-{}.socket
+[Service]
+ExecStart=/lib/systemd/systemd-socket-proxyd {}:{}
+"##,
+            entry.src_port, entry.src_port, entry.dst_ip, entry.dst_port
+        );
+        results.insert(socket_filename, socket);
+        results.insert(service_filename, service);
+    }
+
+    Ok(results)
+}
+
+async fn stop_service(service: &str) -> Result<()> {
+    Command::new("systemctl")
+        .arg("stop")
+        .arg(service)
+        .spawn()?
+        .wait_with_output()
+        .await?;
+    Ok(())
+}
+
+async fn start_service(service: &str) -> Result<()> {
+    Command::new("systemctl")
+        .arg("start")
+        .arg(service)
+        .spawn()?
+        .wait_with_output()
+        .await?;
+    Ok(())
+}
+
+async fn restart_systemd() -> Result<()> {
+    Command::new("systemctl")
+        .arg("daemon-reload")
+        .spawn()?
+        .wait_with_output()
+        .await?;
+    Ok(())
+}
+
+pub async fn update(data: &ConfigData) -> Result<bool> {
+    let configs = build(data)?;
+    let mut changed = false;
+
+    let mut config_dir = tokio::fs::read_dir(SYSTEMD_CONFIG_DIR).await?;
+    while let Some(entry) = config_dir.next().await {
+        let entry = match entry {
+            Ok(entry) => entry,
+            Err(err) => {
+                error!("error listing files {}", err);
+                continue;
+            }
+        };
+        let path = entry.path();
+        if !path.is_file() {
+            continue;
+        }
+
+        let file_name = path.file_name().unwrap().to_string_lossy().to_string();
+        if !file_name.starts_with(&PROXY_PREFIX) {
+            continue;
+        }
+
+        if configs.contains_key(&file_name) {
+            let raw = tokio::fs::read(&path).await?;
+            let contents = String::from_utf8_lossy(&raw).to_string();
+            if configs[&file_name] != contents {
+                info!("updating config: {}", file_name);
+
+                tokio::fs::remove_file(&path).await?;
+                stop_service(&file_name).await?;
+                restart_systemd().await?;
+
+                tokio::fs::write(&path, configs[&file_name].clone()).await?;
+                start_service(&file_name).await?;
+                changed = true;
+            }
+        } else {
+            info!("stopping proxy {}", file_name);
+
+            tokio::fs::remove_file(&path).await?;
+            stop_service(&file_name).await?;
+            restart_systemd().await?;
+            changed = true;
+        }
+    }
+
+    for (file_name, content) in &configs {
+        let path = Path::new(SYSTEMD_CONFIG_DIR).join(file_name);
+        if !path.is_file() {
+            info!("adding service {}", file_name);
+            tokio::fs::write(&path, content).await?;
+            start_service(&file_name).await?;
+            changed = true;
+        }
+    }
+
+    Ok(changed)
+}