GRIN/nym
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

first pass echo server lib

Browse Source
This commit is contained in:
mfahampshire
2024-12-19 12:55:57 +01:00
parent 0f35f5c909
commit 2fc359f58f
6 changed files with 217 additions and 162 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.lock
Generated
+3
View File
@@ -2196,8 +2196,11 @@ dependencies = [
"bincode",
"bytecodec",
"bytes",
"clap 4.5.20",
"dashmap",
"dirs",
"nym-bin-common",
"nym-crypto",
"nym-sdk",
"serde",
"tokio",
Cargo.toml
+1 -1
View File
@@ -128,7 +128,7 @@ members = [
"nym-node-status-api/nym-node-status-client",
"nym-outfox",
"nym-validator-rewarder",
"tools/echo-server",
# "tools/echo-server",
"tools/internal/ssl-inject",
# "tools/internal/sdk-version-bump",
"tools/internal/testnet-manager",
tools/echo-server/Cargo.toml
+11
View File
@@ -9,6 +9,11 @@ edition.workspace = true
license.workspace = true
rust-version.workspace = true
[[bin]]
name = "echo-server"
path = "src/echo-server.rs"
[dependencies]
anyhow.workspace = true
dashmap.workspace = true
@@ -24,3 +29,9 @@ bytecodec = { workspace = true }
nym-sdk = { path = "../../sdk/rust/nym-sdk/" }
bytes.workspace = true
dirs.workspace = true
clap.workspace = true
nym-bin-common = { path = "../../common/bin-common", features = [
"basic_tracing",
"output_format",
] }
nym-crypto = { path = "../../common/crypto", features = ["asymmetric"] }
tools/echo-server/src/echo-server.rs
+17
View File
@@ -0,0 +1,17 @@
// Copyright 2024 - Nym Technologies SA <contact@nymtech.net>
// SPDX-License-Identifier: GPL-3.0-only
// use clap::{Args, Parser, Subcommand};
// use nym_bin_common::{bin_info, bin_info_owned};
// use std::sync::OnceLock;
// fn pretty_build_info_static() -> &'static str {
// static PRETTY_BUILD_INFORMATION: OnceLock<String> = OnceLock::new();
// PRETTY_BUILD_INFORMATION.get_or_init(|| bin_info!().pretty_print())
// }
// This will be a Clap CLI binary for standalone use use wrapping the lib
fn main() {
todo!()
}
tools/echo-server/src/lib.rs
+185
View File
@@ -0,0 +1,185 @@
// Copyright 2024 - Nym Technologies SA <contact@nymtech.net>
// SPDX-License-Identifier: GPL-3.0-only
use anyhow::Result;
use bytes::Bytes;
use nym_crypto::asymmetric::ed25519;
use nym_sdk::mixnet::Recipient;
use nym_sdk::tcp_proxy;
use nym_sdk::tcp_proxy::NymProxyServer;
use std::env;
use std::fs;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use tokio::io::AsyncWriteExt;
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::Mutex;
use tokio::task;
use tokio_stream::StreamExt;
use tokio_util::sync::CancellationToken;
use tracing::{error, info, warn};
pub struct Metrics {
total_conn: AtomicU64,
active_conn: AtomicU64,
bytes_recv: AtomicU64,
bytes_sent: AtomicU64,
}
impl Metrics {
fn new() -> Self {
Self {
total_conn: AtomicU64::new(0),
active_conn: AtomicU64::new(0),
bytes_recv: AtomicU64::new(0),
bytes_sent: AtomicU64::new(0),
}
}
}
pub struct NymEchoServer {
client: Arc<Mutex<NymProxyServer>>,
client_address: Recipient,
listen_addr: String,
metrics: Arc<Metrics>,
cancel_token: CancellationToken,
}
impl NymEchoServer {
pub async fn new(
&mut self,
gateway: Option<ed25519::PublicKey>,
config_path: Option<&str>,
env: String,
listen_port: &str,
) -> Result<Self> {
let home_dir = dirs::home_dir().expect("Unable to get home directory");
let default_path = format!("{}/tmp/nym-proxy-server-config", home_dir.display());
let config_path = config_path.unwrap_or(&default_path);
let listen_addr = format!("127.0.0.1:{}", listen_port);
Ok(NymEchoServer {
client: Arc::new(Mutex::new(
tcp_proxy::NymProxyServer::new(
&listen_addr,
&config_path,
Some(env.clone()),
gateway,
)
.await
.unwrap(),
)),
client_address: self.client.lock().await.nym_address().clone(),
listen_addr,
metrics: Arc::new(Metrics::new()),
cancel_token: CancellationToken::new(),
})
}
pub async fn run(&mut self) -> Result<()> {
let cancel_token = self.cancel_token.clone();
let client = Arc::clone(&self.client);
task::spawn(async move {
client.lock().await.run_with_shutdown().await?;
Ok::<(), anyhow::Error>(())
});
let all_metrics = Arc::clone(&self.metrics);
tokio::spawn(async move {
loop {
tokio::time::sleep(tokio::time::Duration::from_secs(5)).await;
info!(
"Metrics: total_connections={}, active_connections={}, bytes_received={}, bytes_sent={}",
all_metrics.total_conn.load(Ordering::Relaxed),
all_metrics.active_conn.load(Ordering::Relaxed),
all_metrics.bytes_recv.load(Ordering::Relaxed),
all_metrics.bytes_sent.load(Ordering::Relaxed),
);
}
});
let listener = TcpListener::bind(self.listen_addr.clone()).await?;
loop {
tokio::select! {
stream = listener.accept() => {
let (stream, _) = stream?;
let connection_metrics = Arc::clone(&self.metrics);
connection_metrics.total_conn.fetch_add(1, Ordering::Relaxed);
connection_metrics.active_conn.fetch_add(1, Ordering::Relaxed);
tokio::spawn(NymEchoServer::handle_incoming(
stream, connection_metrics, cancel_token.clone()
));
}
_ = self.cancel_token.cancelled() => {
break Ok(());
}
}
}
}
async fn handle_incoming(
socket: TcpStream,
metrics: Arc<Metrics>,
cancel_token: CancellationToken,
) {
let (read, mut write) = socket.into_split();
let codec = tokio_util::codec::BytesCodec::new();
let mut framed_read = tokio_util::codec::FramedRead::new(read, codec);
loop {
tokio::select! {
Some(result) = framed_read.next() => {
match result {
Ok(bytes) => {
let len = bytes.len();
metrics.bytes_recv.fetch_add(len as u64, Ordering::Relaxed);
if let Err(e) = write.write_all(&bytes).await {
error!("Failed to write to stream with err: {}", e);
break;
}
metrics.bytes_sent.fetch_add(len as u64, Ordering::Relaxed);
}
Err(e) => {
error!("Failed to read from stream with err: {}", e);
break;
}
}
}
_ = cancel_token.cancelled() => {
warn!("Shutdown signal received, closing connection");
break;
}
}
}
metrics
.active_conn
.fetch_sub(1, std::sync::atomic::Ordering::Relaxed);
info!("Connection closed");
}
// TODO
// - make disconnect for proxy_server
// - make disconnect for this
// pub async fn disconnect(&self) {
// self.cancel_token.cancel();
// let client = Arc::clone(&self.client);
// client.lock().await.disconnect_pool().await;
// while metrics.active_conn.load(Ordering::Relaxed) > 0 {
// info!("Waiting on active connections to close: sleeping 100ms");
// tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
// }
// }
pub fn nym_address(&self) -> Recipient {
self.client_address
}
pub fn listen_addr(&self) -> String {
self.listen_addr.clone()
}
}
tools/echo-server/src/main.rs
-161
View File
@@ -1,161 +0,0 @@
use anyhow::Result;
use bytes::Bytes;
use nym_sdk::tcp_proxy;
use std::env;
use std::fs;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use tokio::io::AsyncWriteExt;
use tokio::net::{TcpListener, TcpStream};
use tokio::signal;
use tokio::sync::broadcast;
use tokio::task;
use tokio_stream::StreamExt;
use tracing::{error, info, warn};
struct Metrics {
total_conn: AtomicU64,
active_conn: AtomicU64,
bytes_recv: AtomicU64,
bytes_sent: AtomicU64,
}
impl Metrics {
fn new() -> Self {
Self {
total_conn: AtomicU64::new(0),
active_conn: AtomicU64::new(0),
bytes_recv: AtomicU64::new(0),
bytes_sent: AtomicU64::new(0),
}
}
}
#[tokio::main]
async fn main() -> Result<()> {
// if you run this with DEBUG you see the msg buffer on the ProxyServer, but its quite chatty
tracing_subscriber::fmt()
.with_max_level(tracing::Level::INFO)
.init();
let server_port = env::args()
.nth(1)
.expect("Server listen port not specified");
let tcp_addr = format!("127.0.0.1:{}", server_port);
// This dir gets cleaned up at the end: NOTE if you switch env between tests without letting the file do the automatic cleanup, make sure to manually remove this directory up before running again, otherwise your client will attempt to use these keys for the new env
let home_dir = dirs::home_dir().expect("Unable to get home directory");
let conf_path = format!("{}/tmp/nym-proxy-server-config", home_dir.display());
let env_path = env::args().nth(2).expect("Env file not specified");
let env = env_path.to_string();
let mut proxy_server = tcp_proxy::NymProxyServer::new(&tcp_addr, &conf_path, Some(env.clone()))
.await
.unwrap();
let proxy_nym_addr = *proxy_server.nym_address();
info!("ProxyServer listening out on {}", proxy_nym_addr);
task::spawn(async move {
proxy_server.run_with_shutdown().await?;
Ok::<(), anyhow::Error>(())
});
let (shutdown_sender, _) = broadcast::channel(1);
let metrics = Arc::new(Metrics::new());
let all_metrics = Arc::clone(&metrics);
tokio::spawn(async move {
loop {
tokio::time::sleep(tokio::time::Duration::from_secs(5)).await;
info!(
"Metrics: total_connections={}, active_connections={}, bytes_received={}, bytes_sent={}",
all_metrics.total_conn.load(Ordering::Relaxed),
all_metrics.active_conn.load(Ordering::Relaxed),
all_metrics.bytes_recv.load(Ordering::Relaxed),
all_metrics.bytes_sent.load(Ordering::Relaxed),
);
}
});
let listener = TcpListener::bind(tcp_addr).await?;
loop {
tokio::select! {
_ = signal::ctrl_c() => {
info!("Shutdown signal received, closing server...");
let _ = shutdown_sender.send(());
// TODO we need something like this for the ProxyServer client
break;
}
Ok((socket, _)) = listener.accept() => {
let connection_metrics = Arc::clone(&metrics);
let shutdown_rx = shutdown_sender.subscribe();
connection_metrics.total_conn.fetch_add(1, Ordering::Relaxed);
connection_metrics.active_conn.fetch_add(1, Ordering::Relaxed);
tokio::spawn(async move {
handle_incoming(socket, connection_metrics, shutdown_rx).await;
});
}
}
}
signal::ctrl_c().await?;
info!("Received CTRL+C");
fs::remove_dir_all(conf_path)?;
while metrics.active_conn.load(Ordering::Relaxed) > 0 {
info!("Waiting on active connections to close: sleeping 100ms");
// TODO some kind of hard kill here for the ProxyServer
tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
}
Ok(())
}
async fn handle_incoming(
socket: TcpStream,
metrics: Arc<Metrics>,
mut shutdown_rx: broadcast::Receiver<()>,
) {
let (read, mut write) = socket.into_split();
let codec = tokio_util::codec::BytesCodec::new();
let mut framed_read = tokio_util::codec::FramedRead::new(read, codec);
loop {
tokio::select! {
Some(result) = framed_read.next() => {
match result {
Ok(bytes) => {
let len = bytes.len();
metrics.bytes_recv.fetch_add(len as u64, Ordering::Relaxed);
if let Err(e) = write.write_all(&bytes).await {
error!("Failed to write to stream with err: {}", e);
break;
}
metrics.bytes_sent.fetch_add(len as u64, Ordering::Relaxed);
}
Err(e) => {
error!("Failed to read from stream with err: {}", e);
break;
}
}
}
_ = shutdown_rx.recv() => {
warn!("Shutdown signal received, closing connection");
break;
}
// TODO need to work out a way that if this timesout and breaks but you dont hang up the conn on the client end you can reconnect..maybe. If we just use this as a ping echo server I dont think this is a problem
// EDIT I'm not actually sure we want this functionality? Measuring active connections might be useful though
_ = tokio::time::sleep(tokio::time::Duration::from_secs(120)) => {
info!("Timeout reached, assuming we wont get more messages on this conn, closing");
let close_message = "Closing conn, reconnect if you want to ping again";
let bytes: Bytes = close_message.into();
write.write_all(&bytes).await.expect("Couldn't write to socket");
break;
}
}
}
metrics
.active_conn
.fetch_sub(1, std::sync::atomic::Ordering::Relaxed);
info!("Connection closed");
}