pezkuwichain/pezkuwi-telemetry
1
0
Fork 0
mirror of https://github.com/pezkuwichain/pezkuwi-telemetry.git synced 2026-07-17 01:55:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

fmt

Browse Source
This commit is contained in:
James Wilson
2021-07-01 09:38:26 +01:00
parent 16747dd66c
commit 509542e460
25 changed files with 787 additions and 634 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/shard/src/aggregator.rs
+96 -54
View File
@@ -1,10 +1,15 @@
use common::{internal_messages::{self, ShardNodeId}, node_message, AssignId, node_types::BlockHash};
use std::sync::Arc;
use std::sync::atomic::AtomicU64;
use crate::connection::{create_ws_connection, Message};
use common::{
internal_messages::{self, ShardNodeId},
node_message,
node_types::BlockHash,
AssignId,
};
use futures::{channel::mpsc, future};
use futures::{ Sink, SinkExt, StreamExt };
use std::collections::{ HashSet };
use crate::connection::{ create_ws_connection, Message };
use futures::{Sink, SinkExt, StreamExt};
use std::collections::HashSet;
use std::sync::atomic::AtomicU64;
use std::sync::Arc;
/// A unique Id is assigned per websocket connection (or more accurately,
/// per thing-that-subscribes-to-the-aggregator). That connection might send
@@ -16,18 +21,18 @@ type ConnId = u64;
/// from the telemetry core. This can be private since the only
/// external messages are via subscriptions that take
/// [`FromWebsocket`] instances.
#[derive(Clone,Debug)]
#[derive(Clone, Debug)]
enum ToAggregator {
DisconnectedFromTelemetryCore,
ConnectedToTelemetryCore,
FromWebsocket(ConnId, FromWebsocket),
FromTelemetryCore(internal_messages::FromTelemetryCore)
FromTelemetryCore(internal_messages::FromTelemetryCore),
}
/// An incoming socket connection can provide these messages.
/// Until a node has been Added via [`FromWebsocket::Add`],
/// messages from it will be ignored.
#[derive(Clone,Debug)]
#[derive(Clone, Debug)]
pub enum FromWebsocket {
/// Fire this when the connection is established.
Initialize {
@@ -35,22 +40,22 @@ pub enum FromWebsocket {
/// the websocket connection and force the node to reconnect
/// so that it sends its system info again incase the telemetry
/// core has restarted.
close_connection: mpsc::Sender<()>
close_connection: mpsc::Sender<()>,
},
/// Tell the aggregator about a new node.
Add {
message_id: node_message::NodeMessageId,
ip: Option<std::net::IpAddr>,
node: common::node_types::NodeDetails,
genesis_hash: BlockHash
genesis_hash: BlockHash,
},
/// Update/pass through details about a node.
Update {
message_id: node_message::NodeMessageId,
payload: node_message::Payload
payload: node_message::Payload,
},
/// Make a note when the node disconnects.
Disconnected
Disconnected,
}
pub type FromAggregator = internal_messages::FromShardAggregator;
@@ -67,7 +72,7 @@ struct AggregatorInternal {
/// Send messages to the aggregator from websockets via this. This is
/// stored here so that anybody holding an `Aggregator` handle can
/// make use of it.
tx_to_aggregator: mpsc::Sender<ToAggregator>
tx_to_aggregator: mpsc::Sender<ToAggregator>,
}
impl Aggregator {
@@ -77,21 +82,21 @@ impl Aggregator {
// Map responses from our connection into messages that will be sent to the aggregator:
let tx_from_connection = tx_to_aggregator.clone().with(|msg| {
future::ok::<_,mpsc::SendError>(match msg {
future::ok::<_, mpsc::SendError>(match msg {
Message::Connected => ToAggregator::ConnectedToTelemetryCore,
Message::Disconnected => ToAggregator::DisconnectedFromTelemetryCore,
Message::Data(data) => ToAggregator::FromTelemetryCore(data)
Message::Data(data) => ToAggregator::FromTelemetryCore(data),
})
});
// Establish a resiliant connection to the core (this retries as needed):
let tx_to_telemetry_core = create_ws_connection(
tx_from_connection,
telemetry_uri
).await;
let tx_to_telemetry_core = create_ws_connection(tx_from_connection, telemetry_uri).await;
// Handle any incoming messages in our handler loop:
tokio::spawn(Aggregator::handle_messages(rx_from_external, tx_to_telemetry_core));
tokio::spawn(Aggregator::handle_messages(
rx_from_external,
tx_to_telemetry_core,
));
// Return a handle to our aggregator:
Ok(Aggregator(Arc::new(AggregatorInternal {
@@ -103,8 +108,11 @@ impl Aggregator {
// This is spawned into a separate task and handles any messages coming
// in to the aggregator. If nobody is tolding the tx side of the channel
// any more, this task will gracefully end.
async fn handle_messages(mut rx_from_external: mpsc::Receiver<ToAggregator>, mut tx_to_telemetry_core: mpsc::Sender<FromAggregator>) {
use internal_messages::{ FromShardAggregator, FromTelemetryCore };
async fn handle_messages(
mut rx_from_external: mpsc::Receiver<ToAggregator>,
mut tx_to_telemetry_core: mpsc::Sender<FromAggregator>,
) {
use internal_messages::{FromShardAggregator, FromTelemetryCore};
// Just as an optimisation, we can keep track of whether we're connected to the backend
// or not, and ignore incoming messages while we aren't.
@@ -140,41 +148,64 @@ impl Aggregator {
connected_to_telemetry_core = true;
log::info!("Connected to telemetry core");
},
}
ToAggregator::DisconnectedFromTelemetryCore => {
connected_to_telemetry_core = false;
log::info!("Disconnected from telemetry core");
},
ToAggregator::FromWebsocket(_conn_id, FromWebsocket::Initialize { close_connection }) => {
}
ToAggregator::FromWebsocket(
_conn_id,
FromWebsocket::Initialize { close_connection },
) => {
// We boot all connections on a reconnect-to-core to force new systemconnected
// messages to be sent. We could boot on muting, but need to be careful not to boot
// connections where we mute one set of messages it sends and not others.
close_connections.push(close_connection);
},
ToAggregator::FromWebsocket(conn_id, FromWebsocket::Add { message_id, ip, node, genesis_hash }) => {
}
ToAggregator::FromWebsocket(
conn_id,
FromWebsocket::Add {
message_id,
ip,
node,
genesis_hash,
},
) => {
// Don't bother doing anything else if we're disconnected, since we'll force the
// node to reconnect anyway when the backend does:
if !connected_to_telemetry_core { continue }
if !connected_to_telemetry_core {
continue;
}
// Generate a new "local ID" for messages from this connection:
let local_id = to_local_id.assign_id((conn_id, message_id));
// Send the message to the telemetry core with this local ID:
let _ = tx_to_telemetry_core.send(FromShardAggregator::AddNode {
ip,
node,
genesis_hash,
local_id
}).await;
},
ToAggregator::FromWebsocket(conn_id, FromWebsocket::Update { message_id, payload }) => {
let _ = tx_to_telemetry_core
.send(FromShardAggregator::AddNode {
ip,
node,
genesis_hash,
local_id,
})
.await;
}
ToAggregator::FromWebsocket(
conn_id,
FromWebsocket::Update {
message_id,
payload,
},
) => {
// Ignore incoming messages if we're not connected to the backend:
if !connected_to_telemetry_core { continue }
if !connected_to_telemetry_core {
continue;
}
// Get the local ID, ignoring the message if none match:
let local_id = match to_local_id.get_id(&(conn_id, message_id)) {
Some(id) => id,
None => continue
None => continue,
};
// ignore the message if this node has been muted:
@@ -183,28 +214,35 @@ impl Aggregator {
}
// Send the message to the telemetry core with this local ID:
let _ = tx_to_telemetry_core.send(FromShardAggregator::UpdateNode {
local_id,
payload
}).await;
},
let _ = tx_to_telemetry_core
.send(FromShardAggregator::UpdateNode { local_id, payload })
.await;
}
ToAggregator::FromWebsocket(disconnected_conn_id, FromWebsocket::Disconnected) => {
// Find all of the local IDs corresponding to the disconnected connection ID and
// remove them, telling Telemetry Core about them too. This could be more efficient,
// but the mapping isn't currently cached and it's not a super frequent op.
let local_ids_disconnected: Vec<_> = to_local_id.iter()
let local_ids_disconnected: Vec<_> = to_local_id
.iter()
.filter(|(_, &(conn_id, _))| disconnected_conn_id == conn_id)
.map(|(local_id, _)| local_id)
.collect();
for local_id in local_ids_disconnected {
to_local_id.remove_by_id(local_id);
let _ = tx_to_telemetry_core.send(FromShardAggregator::RemoveNode { local_id }).await;
let _ = tx_to_telemetry_core
.send(FromShardAggregator::RemoveNode { local_id })
.await;
}
},
ToAggregator::FromTelemetryCore(FromTelemetryCore::Mute { local_id, reason: _ }) => {
}
ToAggregator::FromTelemetryCore(FromTelemetryCore::Mute {
local_id,
reason: _,
}) => {
// Ignore incoming messages if we're not connected to the backend:
if !connected_to_telemetry_core { continue }
if !connected_to_telemetry_core {
continue;
}
// Mute the local ID we've been told to:
muted.insert(local_id);
@@ -217,13 +255,17 @@ impl Aggregator {
pub fn subscribe_node(&self) -> impl Sink<FromWebsocket, Error = anyhow::Error> + Unpin {
// Assign a unique aggregator-local ID to each connection that subscribes, and pass
// that along with every message to the aggregator loop:
let conn_id: ConnId = self.0.conn_id.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
let conn_id: ConnId = self
.0
.conn_id
.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
let tx_to_aggregator = self.0.tx_to_aggregator.clone();
// Calling `send` on this Sink requires Unpin. There may be a nicer way than this,
// but pinning by boxing is the easy solution for now:
Box::pin(tx_to_aggregator.with(move |msg| async move {
Ok(ToAggregator::FromWebsocket(conn_id, msg))
}))
Box::pin(
tx_to_aggregator
.with(move |msg| async move { Ok(ToAggregator::FromWebsocket(conn_id, msg)) }),
)
}
}
}
backend/shard/src/connection.rs
+52 -23
View File
@@ -1,25 +1,28 @@
use futures::channel::{ mpsc };
use futures::{ Sink, SinkExt, StreamExt };
use futures::channel::mpsc;
use futures::{Sink, SinkExt, StreamExt};
use tokio::net::TcpStream;
use tokio_util::compat::{ TokioAsyncReadCompatExt };
use tokio_util::compat::TokioAsyncReadCompatExt;
#[derive(Clone,Debug)]
#[derive(Clone, Debug)]
pub enum Message<Out> {
Connected,
Disconnected,
Data(Out)
Data(Out),
}
/// Connect to a websocket server, retrying the connection if we're disconnected.
/// - Sends messages when disconnected, reconnected or data received from the connection.
/// - Returns a channel that allows you to send messages to the connection.
/// - Messages all encoded/decoded from bincode.
pub async fn create_ws_connection<In, Out, S, E>(mut tx_to_external: S, telemetry_uri: http::Uri) -> mpsc::Sender<In>
pub async fn create_ws_connection<In, Out, S, E>(
mut tx_to_external: S,
telemetry_uri: http::Uri,
) -> mpsc::Sender<In>
where
S: Sink<Message<Out>, Error = E> + Unpin + Send + Clone + 'static,
E: std::fmt::Debug + std::fmt::Display + Send + 'static,
In: serde::Serialize + Send + 'static,
Out: serde::de::DeserializeOwned + Send + 'static
Out: serde::de::DeserializeOwned + Send + 'static,
{
// Set up a proxy channel to relay messages to the telemetry core, and return one end of it.
// Once a connection to the backend is established, we pass messages along to it. If the connection
@@ -42,7 +45,8 @@ where
connected = true;
// Inform the handler loop that we've reconnected.
tx_to_external.send(Message::Connected)
tx_to_external
.send(Message::Connected)
.await
.expect("must be able to send reconnect msg");
@@ -51,14 +55,20 @@ where
if let Err(e) = tx_to_connection.send(msg).await {
// Issue forwarding a message to the telemetry core?
// Give up and try to reconnect on the next loop iteration.
log::error!("Error sending message to websocker server (will reconnect): {}", e);
log::error!(
"Error sending message to websocker server (will reconnect): {}",
e
);
break;
}
}
},
}
Err(e) => {
// Issue connecting? Wait and try again on the next loop iteration.
log::error!("Error connecting to websocker server (will reconnect): {}", e);
log::error!(
"Error connecting to websocker server (will reconnect): {}",
e
);
}
};
@@ -79,15 +89,18 @@ where
/// This spawns a connection to a websocket server, serializing/deserialziing
/// from bincode as messages are sent or received.
async fn create_ws_connection_no_retry<In, Out, S, E>(mut tx_to_external: S, telemetry_uri: http::Uri) -> anyhow::Result<mpsc::Sender<In>>
async fn create_ws_connection_no_retry<In, Out, S, E>(
mut tx_to_external: S,
telemetry_uri: http::Uri,
) -> anyhow::Result<mpsc::Sender<In>>
where
S: Sink<Message<Out>, Error = E> + Unpin + Send + 'static,
E: std::fmt::Debug + std::fmt::Display,
In: serde::Serialize + Send + 'static,
Out: serde::de::DeserializeOwned + Send + 'static
Out: serde::de::DeserializeOwned + Send + 'static,
{
use soketto::handshake::{Client, ServerResponse};
use bincode::Options;
use soketto::handshake::{Client, ServerResponse};
let host = telemetry_uri.host().unwrap_or("127.0.0.1");
let port = telemetry_uri.port_u16().unwrap_or(8000);
@@ -100,9 +113,13 @@ where
let mut client = Client::new(socket.compat(), host, &path);
let (mut ws_to_connection, mut ws_from_connection) = match client.handshake().await? {
ServerResponse::Accepted { .. } => client.into_builder().finish(),
ServerResponse::Redirect { status_code, .. } |
ServerResponse::Rejected { status_code } => {
return Err(anyhow::anyhow!("Failed to connect to {}{}, status code: {}", host, path, status_code));
ServerResponse::Redirect { status_code, .. } | ServerResponse::Rejected { status_code } => {
return Err(anyhow::anyhow!(
"Failed to connect to {}{}, status code: {}",
host,
path,
status_code
));
}
};
@@ -116,7 +133,10 @@ where
if let Err(e) = ws_from_connection.receive_data(&mut data).await {
// Couldn't receive data may mean all senders are gone, so log
// the error and shut this down:
log::error!("Shutting down websocket connection: Failed to receive data: {}", e);
log::error!(
"Shutting down websocket connection: Failed to receive data: {}",
e
);
return;
}
@@ -126,10 +146,13 @@ where
if let Err(e) = tx_to_external.send(Message::Data(msg)).await {
// Failure to send to our loop likely means it's hit an
// issue and shut down, so bail on this loop as well:
log::error!("Shutting down websocket connection: Failed to send data out: {}", e);
log::error!(
"Shutting down websocket connection: Failed to send data out: {}",
e
);
return;
}
},
}
Err(err) => {
// Log the error but otherwise ignore it and keep running:
log::warn!("Failed to decode message from Backend Core: {:?}", err);
@@ -150,11 +173,17 @@ where
// Any errors sending the message leads to this task ending, which should cascade to
// the entire connection being ended.
if let Err(e) = ws_to_connection.send_binary_mut(bytes).await {
log::error!("Shutting down websocket connection: Failed to send data in: {}", e);
log::error!(
"Shutting down websocket connection: Failed to send data in: {}",
e
);
return;
}
if let Err(e) = ws_to_connection.flush().await {
log::error!("Shutting down websocket connection: Failed to flush data: {}", e);
log::error!(
"Shutting down websocket connection: Failed to flush data: {}",
e
);
return;
}
}
@@ -163,4 +192,4 @@ where
// We return a channel that you can send messages down in order to have
// them sent to the telemetry core:
Ok(tx_to_connection)
}
}
backend/shard/src/json_message/hash.rs
+6 -5
View File
@@ -1,7 +1,7 @@
use serde::de::{self, Deserialize, Deserializer, SeqAccess, Unexpected, Visitor};
use serde::ser::{Serialize, Serializer};
use std::fmt::{self, Debug, Display};
use std::str::FromStr;
use serde::ser::{Serialize, Serializer};
use serde::de::{self, Deserialize, Deserializer, Unexpected, Visitor, SeqAccess};
const HASH_BYTES: usize = 32;
@@ -28,7 +28,9 @@ impl<'de> Visitor<'de> for HashVisitor {
type Value = Hash;
fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
formatter.write_str("byte array of length 32, or hexidecimal string of 32 bytes beginning with 0x")
formatter.write_str(
"byte array of length 32, or hexidecimal string of 32 bytes beginning with 0x",
)
}
fn visit_str<E>(self, value: &str) -> Result<Self::Value, E>
@@ -65,7 +67,7 @@ impl<'de> Visitor<'de> for HashVisitor {
for (i, byte) in hash.iter_mut().enumerate() {
match seq.next_element()? {
Some(b) => *byte = b,
None => return Err(de::Error::invalid_length(i, &"an array of 32 bytes"))
None => return Err(de::Error::invalid_length(i, &"an array of 32 bytes")),
}
}
@@ -176,7 +178,6 @@ mod tests {
assert_eq!(hash, DUMMY);
}
#[test]
fn deserialize_json_array_too_short() {
let json = r#"[222,173,190,239,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]"#;
backend/shard/src/json_message/mod.rs
+1 -1
View File
@@ -3,5 +3,5 @@
mod hash;
mod node_message;
pub use hash::Hash;
pub use node_message::*;
pub use hash::Hash;
backend/shard/src/json_message/node_message.rs
+21 -46
View File
@@ -4,9 +4,9 @@
//! compatibility with the input data when we make changes to our internal data
//! structures (for example, to support bincode better).
use super::hash::Hash;
use serde::{Deserialize};
use common::node_message as internal;
use common::node_types;
use serde::Deserialize;
/// This struct represents a telemetry message sent from a node as
/// a JSON payload. Since JSON is self describing, we can use attributes
@@ -36,11 +36,12 @@ pub enum NodeMessage {
impl From<NodeMessage> for internal::NodeMessage {
fn from(msg: NodeMessage) -> Self {
match msg {
NodeMessage::V1 { payload } => {
internal::NodeMessage::V1 { payload: payload.into() }
NodeMessage::V1 { payload } => internal::NodeMessage::V1 {
payload: payload.into(),
},
NodeMessage::V2 { id, payload } => {
internal::NodeMessage::V2 { id, payload: payload.into() }
NodeMessage::V2 { id, payload } => internal::NodeMessage::V2 {
id,
payload: payload.into(),
},
}
}
@@ -80,45 +81,19 @@ pub enum Payload {
impl From<Payload> for internal::Payload {
fn from(msg: Payload) -> Self {
match msg {
Payload::SystemConnected(m) => {
internal::Payload::SystemConnected(m.into())
},
Payload::SystemInterval(m) => {
internal::Payload::SystemInterval(m.into())
},
Payload::BlockImport(m) => {
internal::Payload::BlockImport(m.into())
},
Payload::NotifyFinalized(m) => {
internal::Payload::NotifyFinalized(m.into())
},
Payload::TxPoolImport => {
internal::Payload::TxPoolImport
},
Payload::AfgFinalized(m) => {
internal::Payload::AfgFinalized(m.into())
},
Payload::AfgReceivedPrecommit(m) => {
internal::Payload::AfgReceivedPrecommit(m.into())
},
Payload::AfgReceivedPrevote(m) => {
internal::Payload::AfgReceivedPrevote(m.into())
},
Payload::AfgReceivedCommit(m) => {
internal::Payload::AfgReceivedCommit(m.into())
},
Payload::AfgAuthoritySet(m) => {
internal::Payload::AfgAuthoritySet(m.into())
},
Payload::AfgFinalizedBlocksUpTo => {
internal::Payload::AfgFinalizedBlocksUpTo
},
Payload::AuraPreSealedBlock => {
internal::Payload::AuraPreSealedBlock
},
Payload::PreparedBlockForProposing => {
internal::Payload::PreparedBlockForProposing
},
Payload::SystemConnected(m) => internal::Payload::SystemConnected(m.into()),
Payload::SystemInterval(m) => internal::Payload::SystemInterval(m.into()),
Payload::BlockImport(m) => internal::Payload::BlockImport(m.into()),
Payload::NotifyFinalized(m) => internal::Payload::NotifyFinalized(m.into()),
Payload::TxPoolImport => internal::Payload::TxPoolImport,
Payload::AfgFinalized(m) => internal::Payload::AfgFinalized(m.into()),
Payload::AfgReceivedPrecommit(m) => internal::Payload::AfgReceivedPrecommit(m.into()),
Payload::AfgReceivedPrevote(m) => internal::Payload::AfgReceivedPrevote(m.into()),
Payload::AfgReceivedCommit(m) => internal::Payload::AfgReceivedCommit(m.into()),
Payload::AfgAuthoritySet(m) => internal::Payload::AfgAuthoritySet(m.into()),
Payload::AfgFinalizedBlocksUpTo => internal::Payload::AfgFinalizedBlocksUpTo,
Payload::AuraPreSealedBlock => internal::Payload::AuraPreSealedBlock,
Payload::PreparedBlockForProposing => internal::Payload::PreparedBlockForProposing,
}
}
}
@@ -134,7 +109,7 @@ impl From<SystemConnected> for internal::SystemConnected {
fn from(msg: SystemConnected) -> Self {
internal::SystemConnected {
genesis_hash: msg.genesis_hash.into(),
node: msg.node.into()
node: msg.node.into(),
}
}
}
@@ -243,7 +218,7 @@ impl From<Block> for node_types::Block {
fn from(block: Block) -> Self {
node_types::Block {
hash: block.hash.into(),
height: block.height
height: block.height,
}
}
}
backend/shard/src/main.rs
+36 -44
View File
@@ -1,19 +1,19 @@
mod aggregator;
mod connection;
mod real_ip;
mod json_message;
mod real_ip;
use std::net::IpAddr;
use structopt::StructOpt;
use http::Uri;
use simple_logger::SimpleLogger;
use futures::{StreamExt, SinkExt, channel::mpsc};
use warp::Filter;
use warp::filters::ws;
use aggregator::{Aggregator, FromWebsocket};
use common::{node_message, LogLevel};
use aggregator::{ Aggregator, FromWebsocket };
use futures::{channel::mpsc, SinkExt, StreamExt};
use http::Uri;
use real_ip::real_ip;
use simple_logger::SimpleLogger;
use structopt::StructOpt;
use warp::filters::ws;
use warp::Filter;
const VERSION: &str = env!("CARGO_PKG_VERSION");
const AUTHORS: &str = env!("CARGO_PKG_AUTHORS");
@@ -27,25 +27,17 @@ struct Opts {
/// This is the socket address that this shard is listening to. This is restricted to
/// localhost (127.0.0.1) by default and should be fine for most use cases. If
/// you are using Telemetry in a container, you likely want to set this to '0.0.0.0:8000'
#[structopt(
short = "l",
long = "listen",
default_value = "127.0.0.1:8001",
)]
#[structopt(short = "l", long = "listen", default_value = "127.0.0.1:8001")]
socket: std::net::SocketAddr,
/// The desired log level; one of 'error', 'warn', 'info', 'debug' or 'trace', where
/// 'error' only logs errors and 'trace' logs everything.
#[structopt(
required = false,
long = "log",
default_value = "info",
)]
#[structopt(required = false, long = "log", default_value = "info")]
log_level: LogLevel,
/// Url to the Backend Core endpoint accepting shard connections
#[structopt(
short = "c",
long = "core",
default_value = "ws://127.0.0.1:8000/shard_submit/",
short = "c",
long = "core",
default_value = "ws://127.0.0.1:8000/shard_submit/"
)]
core_url: Uri,
}
@@ -60,10 +52,7 @@ async fn main() {
.init()
.expect("Must be able to start a logger");
log::info!(
"Starting Telemetry Shard version: {}",
VERSION
);
log::info!("Starting Telemetry Shard version: {}", VERSION);
if let Err(e) = start_server(opts).await {
log::error!("Error starting server: {}", e);
@@ -72,26 +61,21 @@ async fn main() {
/// Declare our routes and start the server.
async fn start_server(opts: Opts) -> anyhow::Result<()> {
let aggregator = Aggregator::spawn(opts.core_url).await?;
// Handle requests to /health by returning OK.
let health_route =
warp::path("health")
.map(|| "OK");
let health_route = warp::path("health").map(|| "OK");
// Handle websocket requests to /submit.
let ws_route =
warp::path("submit")
.and(warp::ws())
.and(real_ip())
.map(move |ws: ws::Ws, addr: Option<IpAddr>| {
let ws_route = warp::path("submit").and(warp::ws()).and(real_ip()).map(
move |ws: ws::Ws, addr: Option<IpAddr>| {
// Send messages from the websocket connection to this sink
// to have them pass to the aggregator.
let tx_to_aggregator = aggregator.subscribe_node();
log::info!("Opening /submit connection from {:?}", addr);
ws.on_upgrade(move |websocket| async move {
let (mut tx_to_aggregator, websocket) = handle_websocket_connection(websocket, tx_to_aggregator, addr).await;
let (mut tx_to_aggregator, websocket) =
handle_websocket_connection(websocket, tx_to_aggregator, addr).await;
log::info!("Closing /submit connection from {:?}", addr);
// Tell the aggregator that this connection has closed, so it can tidy up.
let _ = tx_to_aggregator.send(FromWebsocket::Disconnected).await;
@@ -99,7 +83,8 @@ async fn start_server(opts: Opts) -> anyhow::Result<()> {
// a ws::Message using `ws::Message::close_with`, rather than using this method:
let _ = websocket.close().await;
})
});
},
);
// Merge the routes and start our server:
let routes = ws_route.or(health_route);
@@ -108,8 +93,13 @@ async fn start_server(opts: Opts) -> anyhow::Result<()> {
}
/// This takes care of handling messages from an established socket connection.
async fn handle_websocket_connection<S>(mut websocket: ws::WebSocket, mut tx_to_aggregator: S, addr: Option<IpAddr>) -> (S, ws::WebSocket)
where S: futures::Sink<FromWebsocket, Error = anyhow::Error> + Unpin
async fn handle_websocket_connection<S>(
mut websocket: ws::WebSocket,
mut tx_to_aggregator: S,
addr: Option<IpAddr>,
) -> (S, ws::WebSocket)
where
S: futures::Sink<FromWebsocket, Error = anyhow::Error> + Unpin,
{
// This could be a oneshot channel, but it's useful to be able to clone
// messages, and we can't clone oneshot channel senders.
@@ -117,7 +107,7 @@ async fn handle_websocket_connection<S>(mut websocket: ws::WebSocket, mut tx_to_
// Tell the aggregator about this new connection, and give it a way to close this connection:
let init_msg = FromWebsocket::Initialize {
close_connection: close_connection_tx
close_connection: close_connection_tx,
};
if let Err(e) = tx_to_aggregator.send(init_msg).await {
log::error!("Error sending message to aggregator: {}", e);
@@ -179,13 +169,15 @@ async fn handle_websocket_connection<S>(mut websocket: ws::WebSocket, mut tx_to_
/// Deserialize an incoming websocket message, returning an error if something
/// fatal went wrong, [`Some`] message if all went well, and [`None`] if a non-fatal
/// issue was encountered and the message should simply be ignored.
fn deserialize_ws_message(msg: Result<ws::Message, warp::Error>) -> anyhow::Result<Option<node_message::NodeMessage>> {
fn deserialize_ws_message(
msg: Result<ws::Message, warp::Error>,
) -> anyhow::Result<Option<node_message::NodeMessage>> {
// If we see any errors, log them and end our loop:
let msg = match msg {
Err(e) => {
return Err(anyhow::anyhow!("Error in node websocket connection: {}", e));
},
Ok(msg) => msg
}
Ok(msg) => msg,
};
// If the message isn't something we want to handle, just ignore it.
@@ -202,11 +194,11 @@ fn deserialize_ws_message(msg: Result<ws::Message, warp::Error>) -> anyhow::Resu
// let bytes: &[u8] = bytes.get(..512).unwrap_or_else(|| &bytes);
// let msg_start = std::str::from_utf8(bytes).unwrap_or_else(|_| "INVALID UTF8");
// log::warn!("Failed to parse node message ({}): {}", msg_start, e);
return Ok(None)
return Ok(None);
}
};
// Pull relevant details from the message:
let node_message: node_message::NodeMessage = node_message.into();
Ok(Some(node_message))
}
}
backend/shard/src/real_ip.rs
+20 -11
View File
@@ -1,6 +1,6 @@
use std::net::{ SocketAddr, IpAddr };
use warp::filters::header;
use std::net::{IpAddr, SocketAddr};
use warp::filters::addr;
use warp::filters::header;
use warp::Filter;
/**
@@ -24,7 +24,8 @@ If that _still_ doesn't work, fall back to the socket address of the connection.
Return `None` if all of this fails to yield an address.
*/
pub fn real_ip() -> impl warp::Filter<Extract = (Option<IpAddr>,), Error = warp::Rejection> + Clone {
pub fn real_ip() -> impl warp::Filter<Extract = (Option<IpAddr>,), Error = warp::Rejection> + Clone
{
header::optional("forwarded")
.and(header::optional("x-forwarded-for"))
.and(header::optional("x-real-ip"))
@@ -40,12 +41,16 @@ fn pick_best_ip_from_options(
// X-Real-IP header value (if present)
real_ip: Option<String>,
// socket address (if known)
addr: Option<SocketAddr>
addr: Option<SocketAddr>,
) -> Option<IpAddr> {
let realip = forwarded.as_ref().and_then(|val| get_first_addr_from_forwarded_header(val))
let realip = forwarded
.as_ref()
.and_then(|val| get_first_addr_from_forwarded_header(val))
.or_else(|| {
// fall back to X-Forwarded-For
forwarded_for.as_ref().and_then(|val| get_first_addr_from_x_forwarded_for_header(val))
forwarded_for
.as_ref()
.and_then(|val| get_first_addr_from_x_forwarded_for_header(val))
})
.or_else(|| {
// fall back to X-Real-IP
@@ -54,7 +59,8 @@ fn pick_best_ip_from_options(
.and_then(|ip| {
// Try parsing assuming it may have a port first,
// and then assuming it doesn't.
ip.parse::<SocketAddr>().map(|s| s.ip())
ip.parse::<SocketAddr>()
.map(|s| s.ip())
.or_else(|_| ip.parse::<IpAddr>())
.ok()
})
@@ -110,7 +116,10 @@ mod test {
fn get_addr_from_forwarded_rfc_examples() {
let examples = vec![
(r#"for="_gazonk""#, "_gazonk"),
(r#"For="[2001:db8:cafe::17]:4711""#, "[2001:db8:cafe::17]:4711"),
(
r#"For="[2001:db8:cafe::17]:4711""#,
"[2001:db8:cafe::17]:4711",
),
(r#"for=192.0.2.60;proto=http;by=203.0.113.43"#, "192.0.2.60"),
(r#"for=192.0.2.43, for=198.51.100.17"#, "192.0.2.43"),
];
@@ -119,9 +128,9 @@ mod test {
assert_eq!(
get_first_addr_from_forwarded_header(value),
Some(expected),
"Header value: {}", value
"Header value: {}",
value
);
}
}
}
}