Differential D7661 Diff 25958 services/tunnelbroker/src/websockets/mod.rs

Changeset View

Standalone View

services/tunnelbroker/src/websockets/mod.rs

use futures::future;		use futures::future;
use futures_util::{StreamExt, TryStreamExt};		use futures_util::{StreamExt, TryStreamExt};
use std::net::SocketAddr;		use std::net::SocketAddr;
use std::{env, io::Error};		use std::{env, io::Error};
use tokio::net::{TcpListener, TcpStream};		use tokio::net::{TcpListener, TcpStream};
use tokio::sync::mpsc;		use tokio::sync::mpsc;
		use tokio_tungstenite::tungstenite::Message;
use tracing::{debug, error, info};		use tracing::{debug, error, info};
use tunnelbroker_messages::messages::Messages;		use tunnelbroker_messages::messages::Messages;

use crate::ACTIVE_CONNECTIONS;		use crate::ACTIVE_CONNECTIONS;

pub async fn run_server() -> Result<(), Error> {		pub async fn run_server() -> Result<(), Error> {
let addr = env::var("COMM_TUNNELBROKER_WEBSOCKET_ADDR")		let addr = env::var("COMM_TUNNELBROKER_WEBSOCKET_ADDR")
.unwrap_or_else(\|_\| "127.0.0.1:51001".to_string());		.unwrap_or_else(\|_\| "127.0.0.1:51001".to_string());
Show All 21 Lines	Err(e) => {
);		);
return;		return;
}		}
};		};

let (_outgoing, incoming) = ws_stream.split();		let (_outgoing, incoming) = ws_stream.split();

let handle_incoming = incoming.try_for_each(\|msg\| {		let handle_incoming = incoming.try_for_each(\|msg\| {
let message_text = msg.to_text().unwrap();		debug!("Received message from {}", addr);
debug!("Received message from {}: {}", addr, message_text);		match msg {
		Message::Text(text) => {
match handle_message(message_text) {		match handle_message(&text) {
Ok(_) => {		Ok(_) => {
debug!("Successfully handled message: {}", message_text)		debug!("Successfully handled message: {}", text)
}		}
Err(e) => {		Err(e) => {
error!("Failed to process message: {}", e);		error!("Failed to process message: {}", e);
}		}
};		};
		}
		_ => {
		error!("Invalid message was received");
		}
		}

future::ok(())		future::ok(())
});		});

// Create channel for messages to be passed to this connection		// Create channel for messages to be passed to this connection
let (tx, mut rx) = mpsc::unbounded_channel::<Messages>();		let (tx, mut rx) = mpsc::unbounded_channel::<Messages>();
// TODO: Use device's public key, once we support the SessionRequest message		// TODO: Use device's public key, once we support the SessionRequest message
ACTIVE_CONNECTIONS.insert("test".to_string(), tx.clone());		ACTIVE_CONNECTIONS.insert("test".to_string(), tx.clone());

		debug!("Polling for messages from: {}", addr);
tokio::select! {		tokio::select! {
Some(_) = rx.recv() => { debug!("Received message from channel") },		Some(_) = rx.recv() => { debug!("Received message from channel") },
Ok(_) = handle_incoming => { debug!("Received message from websocket" )},		Ok(_) = handle_incoming => { debug!("Received message from websocket") },
else => {		else => {
info!("Connection with {} closed.", addr);		info!("Connection with {} closed.", addr);
ACTIVE_CONNECTIONS.remove("test");		ACTIVE_CONNECTIONS.remove("test");
}		}
}		}
}		}

fn handle_message(message: &str) -> Result<(), serde_json::Error> {		fn handle_message(message: &str) -> Result<(), serde_json::Error> {
serde_json::from_str::<Messages>(message)?;		serde_json::from_str::<Messages>(message)?;

Ok(())		Ok(())
}		}