Skip to content

Commit

Permalink
Code Refactor
Browse files Browse the repository at this point in the history
  • Loading branch information
@Hasan6979
Hasan6979 committed Oct 18, 2024
1 parent 8da4bca commit d07eb05
Show file tree
Hide file tree
Showing 5 changed files with 96 additions and 78 deletions.
1 change: 1 addition & 0 deletions Cargo.lock
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion clis/teliod/Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,9 +5,9 @@ edition = "2021"

[dependencies]
clap.workspace = true
regex.workspace = true
reqwest = { version = "0.11.16", default-features = false, features = [
"json",
"blocking",
"rustls-tls",
] }
serde = { workspace = true }
Expand Down
37 changes: 18 additions & 19 deletions clis/teliod/src/core_api.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,8 @@ use serde_json::Value;
use std::sync::Arc;
use telio::crypto::PublicKey;
use thiserror::Error;
use tracing::info;
use tracing::{debug, info};

const API_BASE: &str = "https://api.nordvpn.com/v1";

#[cfg(windows)]
Expand All @@ -19,7 +20,7 @@ const OS_NAME: &str = "linux";
#[derive(Debug, Error)]
pub enum Error {
#[error(transparent)]
Reqwest(#[from] reqwest::Error),
ReqwestError(#[from] reqwest::Error),
#[error(transparent)]
DeserializeError(#[from] serde_json::Error),
#[error("Machine Identifier not found due to Error: {0}")]
Expand All @@ -32,8 +33,8 @@ pub enum Error {
InvalidResponse,
}

#[derive(Debug, Default, Serialize, Deserialize)]
struct MeshDev {
#[derive(Default, Serialize, Deserialize)]
struct MeshConfig {
public_key: PublicKey,
hardware_identifier: String,
os: String,
Expand All @@ -46,9 +47,9 @@ pub async fn load_identifier_from_api(
auth_token: &String,
public_key: PublicKey,
) -> Result<String, Error> {
info!("fetching machine identifier");
debug!("Fetching machine identifier");
let client = Client::new();
let register = client
let response = client
.get(&format!("{}/meshnet/machines", API_BASE))
.header(
header::AUTHORIZATION,
Expand All @@ -58,9 +59,8 @@ pub async fn load_identifier_from_api(
.send()
.await?;

info!("Status {}", register.status());
let status = register.status().clone();
let json_data: Value = serde_json::from_str(&register.text().await?)?;
let status = response.status().clone();
let json_data: Value = serde_json::from_str(&response.text().await?)?;

if let Some(items) = json_data.as_array() {
for item in items {
Expand All @@ -86,7 +86,7 @@ pub async fn load_identifier_from_api(
}

pub async fn update_machine(client_config: &ClientConfig) -> Result<StatusCode, Error> {
info!("Updating machine");
debug!("Updating machine");
let client = Client::new();
Ok(client
.patch(&format!(
Expand All @@ -100,7 +100,7 @@ pub async fn update_machine(client_config: &ClientConfig) -> Result<StatusCode,
)
.header(header::CONTENT_TYPE, "application/json")
.header(header::ACCEPT, "application/json")
.json(&MeshDev {
.json(&MeshConfig {
public_key: client_config.public_key,
hardware_identifier: client_config.hw_identifier.clone(),
os: OS_NAME.to_owned(),
Expand All @@ -114,14 +114,13 @@ pub async fn update_machine(client_config: &ClientConfig) -> Result<StatusCode,
}

pub async fn get_meshmap(client_config: Arc<ClientConfig>) -> Result<String, Error> {
info!("Getting meshmap");
debug!("Getting meshmap");
let client = Client::new();
Ok({
client
.get(&format!(
"{}/meshnet/machines/{}/map",
API_BASE,
client_config.machine_identifier.clone()
API_BASE, client_config.machine_identifier
))
.header(
header::AUTHORIZATION,
Expand All @@ -142,15 +141,15 @@ pub async fn register_machine(
) -> Result<String, Error> {
info!("Registering machine");
let client = Client::new();
let result = client
let response = client
.post(&format!("{}/meshnet/machines", API_BASE))
.header(
header::AUTHORIZATION,
format!("Bearer token:{}", auth_token),
)
.header(header::CONTENT_TYPE, "application/json")
.header(header::ACCEPT, "application/json")
.json(&MeshDev {
.json(&MeshConfig {
public_key: public_key,
hardware_identifier: hw_identifier.clone(),
os: OS_NAME.to_owned(),
Expand All @@ -162,15 +161,15 @@ pub async fn register_machine(
.await?;

// Save the machine identifier received from API
if result.status() == StatusCode::CREATED {
let response: Value = serde_json::from_str(&result.text().await?)?;
if response.status() == StatusCode::CREATED {
let response: Value = serde_json::from_str(&response.text().await?)?;
if let Some(machine_identifier) = response.get("identifier").and_then(|i| i.as_str()) {
info!("Machine Registered!");
return Ok(machine_identifier.to_owned());
} else {
Err(Error::InvalidResponse)
}
} else {
Err(Error::PeerRegisteringError(result.status()))
Err(Error::PeerRegisteringError(response.status()))
}
}
132 changes: 75 additions & 57 deletions clis/teliod/src/main.rs
View file
Original file line number Diff line number Diff line change
@@ -1,44 +1,42 @@
//! Main and implementation of config and commands for Teliod - simple telio daemon for Linux and OpenWRT

use crate::comms::DaemonSocket;
use crate::core_api::{
get_meshmap as get_meshmap_from_server, load_identifier_from_api, register_machine,
update_machine, Error as ApiError,
};
use clap::Parser;
use futures::stream::StreamExt;
use nix::libc::SIGTERM;
use reqwest::StatusCode;
use serde::{de, Deserialize, Deserializer, Serialize, Serializer};
use serde_json::error::Error as SerdeJsonError;
use signal_hook_tokio::Signals;
use std::cmp::min;
use std::{
fs::{self, File},
str::FromStr,
sync::Arc,
};
use thiserror::Error as ThisError;
use tokio::{sync::mpsc, task::JoinError, time::Duration};
use tracing::{debug, error, info, level_filters::LevelFilter, trace, warn};

mod comms;
mod core_api;

use crate::comms::DaemonSocket;
use serde::{de, Deserialize, Deserializer, Serialize, Serializer};
use serde_json::error::Error as SerdeJsonError;
use std::cmp::min;
use telio::{
crypto::{PublicKey, SecretKey},
device::{Device, DeviceConfig, Error as DeviceError},
telio_model::{config::Config as MeshMap, features::Features},
telio_utils::select,
telio_wg::AdapterType,
};
use thiserror::Error as ThisError;
use tokio::{sync::mpsc, task::JoinError, time::Duration};
use tracing::{debug, error, info, level_filters::LevelFilter, trace, warn};

use crate::core_api::{
get_meshmap as get_meshmap_from_server, load_identifier_from_api, register_machine,
update_machine, Error as ApiError,
};
use futures::stream::StreamExt;
mod comms;
mod core_api;

const MAX_RETRIES: u32 = 5;
const BASE_DELAY_MS: u64 = 500;
const MAX_BACKOFF_TIME: u64 = 5000;

#[derive(Serialize, Deserialize, Debug, Clone)]
#[derive(Serialize, Deserialize, Debug)]
struct TeliodDaemonConfig {
#[serde(
deserialize_with = "deserialize_log_level",
Expand All @@ -47,6 +45,7 @@ struct TeliodDaemonConfig {
log_level: LevelFilter,
log_file_path: String,
hw_identifier: Option<String>,
#[serde(deserialize_with = "deserialize_authentication_token")]
authentication_token: String,
private_key: Option<SecretKey>,
machine_identifier: Option<String>,
Expand Down Expand Up @@ -88,6 +87,18 @@ where
serializer.serialize_str(level_str)
}

fn deserialize_authentication_token<'de, D: Deserializer<'de>>(
deserializer: D,
) -> Result<String, D::Error> {
let raw_string: String = de::Deserialize::deserialize(deserializer)?;
let re = regex::Regex::new("[0-9a-f]{64}").map_err(de::Error::custom)?;
if re.is_match(&raw_string) {
Ok(raw_string)
} else {
Err(de::Error::custom("Incorrect authentication token"))
}
}

#[derive(Parser, Debug)]
#[clap()]
#[derive(Serialize, Deserialize)]
Expand All @@ -111,7 +122,9 @@ enum Cmd {

#[derive(Debug)]
pub enum TelioTaskCmd {
// Command to set the downloaded meshmap to telio instance
UpdateMeshmap(MeshMap),
// Triggers downloading meshmap from server
GetMeshmap,
}

Expand All @@ -133,16 +146,14 @@ enum TeliodError {
DaemonIsNotRunning,
#[error("Daemon is running")]
DaemonIsRunning,
#[error("Token not fully loaded")]
TokenError,
#[error(transparent)]
CoreApiError(#[from] ApiError),
#[error(transparent)]
DeviceError(#[from] DeviceError),
#[error("Unable to update machine due to Error: {0}")]
UpdateMachineError(StatusCode),
#[error("Max retries exceeding to update machine")]
UpdateMachineTimeoutError,
UpdateMachineTimeout,
}

#[tokio::main]
Expand Down Expand Up @@ -196,7 +207,10 @@ impl CommandListener {
}
}

async fn init_api(config: &mut TeliodDaemonConfig, config_path: String) -> Result<(), TeliodError> {
async fn init_api(
config: &mut TeliodDaemonConfig,
config_path: String,
) -> Result<ClientConfig, TeliodError> {
let mut client_config = ClientConfig::new(config).await?;

let mut retries = 0;
Expand All @@ -209,21 +223,24 @@ async fn init_api(config: &mut TeliodDaemonConfig, config_path: String) -> Resul
return Err(TeliodError::UpdateMachineError(status));
}
StatusCode::NOT_FOUND => {
client_config.machine_identifier = if let Ok(id) =
load_identifier_from_api(&config.authentication_token, client_config.public_key)
.await
client_config.machine_identifier = match load_identifier_from_api(
&config.authentication_token,
client_config.public_key,
)
.await
{
id
} else {
debug!("Machine not yet registered");
// Retry machine update after registering. To make sure everything was successful
// If registering fails. Close the daemon
register_machine(
&client_config.hw_identifier,
client_config.public_key,
&client_config.auth_token,
)
.await?
Ok(id) => id,
Err(e) => {
info!("Unable to load identifier due to {e}");
// Retry machine update after registering. To make sure everything was successful
// If registering fails. Close the daemon
register_machine(
&client_config.hw_identifier,
client_config.public_key,
&client_config.auth_token,
)
.await?
}
}
}
_ => {
Expand All @@ -236,7 +253,7 @@ async fn init_api(config: &mut TeliodDaemonConfig, config_path: String) -> Resul
retries += 1;
} else {
// If max retries exceeded, exit daemon
return Err(TeliodError::UpdateMachineTimeoutError);
return Err(TeliodError::UpdateMachineTimeout);
}
}
}
Expand All @@ -248,7 +265,7 @@ async fn init_api(config: &mut TeliodDaemonConfig, config_path: String) -> Resul

let mut file = File::create(config_path)?;
serde_json::to_writer(&mut file, &config)?;
Ok(())
Ok(client_config)
}

// From async context Telio needs to be run in separate task
Expand All @@ -265,19 +282,19 @@ fn telio_task(
None,
)?;

let client_arc = std::sync::Arc::new(client_config);
let config_ptr = std::sync::Arc::new(client_config);
// TODO: This is temporary to be removed later on when we have proper integration
// tests with core API. This is to not look for tokens in a test environment
// right now as the values are dummy and program will not run as it expects
// real tokens.
if !client_arc.auth_token.eq("") {
start_telio(&mut telio, client_arc.private_key)?;
get_meshmap(client_arc.clone(), tx_channel.clone());
if !config_ptr.auth_token.eq("") {
start_telio(&mut telio, config_ptr.private_key)?;
get_meshmap(config_ptr.clone(), tx_channel.clone());

while let Some(cmd) = rx_channel.blocking_recv() {
info!("Got command {:?}", cmd);
match cmd {
TelioTaskCmd::GetMeshmap => get_meshmap(client_arc.clone(), tx_channel.clone()),
TelioTaskCmd::GetMeshmap => get_meshmap(config_ptr.clone(), tx_channel.clone()),
TelioTaskCmd::UpdateMeshmap(map) => {
if let Err(e) = telio.set_config(&Some(map)) {
error!("Unable to set meshmap due to {e}");
Expand All @@ -296,7 +313,13 @@ fn get_meshmap(client_config: Arc<ClientConfig>, tx: mpsc::Sender<TelioTaskCmd>)
let result = get_meshmap_from_server(config_clone).await;
match result {
Ok(map) => {
let meshmap: MeshMap = serde_json::from_str(&map).unwrap();
let meshmap: MeshMap = match serde_json::from_str(&map) {
Ok(map) => map,
Err(e) => {
error!("Unable to parse meshmap due to {e}");
return;
}
};
trace!("Meshmap {:#?}", meshmap);
if let Err(e) = tx.send(TelioTaskCmd::UpdateMeshmap(meshmap)).await {
error!("Unable to send meshmap due to {e}");
Expand Down Expand Up @@ -386,21 +409,16 @@ async fn daemon_event_loop(
// telio task
let (tx, rx) = mpsc::channel(10);

// TODO: This is temporary to be removed later on when we have proper integration
// tests with core API. This is to not look for tokens in a test environment
// right now as the values are dummy and program will not run as it expects
// real tokens.
// TODO: This if condition and ::default call is temporary to be removed later
// on when we have proper integration tests with core API.
// This is to not look for tokens in a test environment right now as the values
// are dummy and program will not run as it expects real tokens.
let mut client_config = ClientConfig::default();
if !config.authentication_token.eq("abcd1234") {
init_api(&mut config, config_path).await?;

client_config = ClientConfig {
hw_identifier: config.hw_identifier.ok_or(TeliodError::TokenError)?,
auth_token: config.authentication_token,
machine_identifier: config.machine_identifier.ok_or(TeliodError::TokenError)?,
private_key: config.private_key.ok_or(TeliodError::TokenError)?,
public_key: config.private_key.ok_or(TeliodError::TokenError)?.public(),
};
if !config
.authentication_token
.eq("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa")
{
client_config = init_api(&mut config, config_path).await?;
}
let tx_clone = tx.clone();
let telio_task_handle =
Expand Down
Loading

0 comments on commit d07eb05

Please sign in to comment.