Skip to content

Commit

Permalink
patch(hermes): improve ws reliability
Browse files Browse the repository at this point in the history
- Add max message size for incoming messages
- Add sent message rate limit and ip whitelisting
  • Loading branch information
ali-bahjati committed Oct 5, 2023
1 parent 5fdc0d2 commit 1a64d58
Show file tree
Hide file tree
Showing 5 changed files with 182 additions and 33 deletions.
69 changes: 68 additions & 1 deletion hermes/Cargo.lock

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

5 changes: 4 additions & 1 deletion hermes/Cargo.toml
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[package]
name = "hermes"
version = "0.2.0"
version = "0.2.1"
description = "Hermes is an agent that provides Verified Prices from the Pythnet Pyth Oracle."
edition = "2021"

Expand All @@ -19,10 +19,13 @@ env_logger = { version = "0.10.0" }
futures = { version = "0.3.28" }
hex = { version = "0.4.3", features = ["serde"] }
humantime = { version = "2.1.0" }
ipnet = { version = "2.8.0" }
governor = { version = "0.6.0" }
lazy_static = { version = "1.4.0" }
libc = { version = "0.2.140" }
log = { version = "0.4.17" }
mock_instant = { version = "0.3.1", features = ["sync"] }
nonzero_ext = { version = "0.3.0" }
prometheus-client = { version = "0.21.1" }
pyth-sdk = { version = "0.8.0" }
pythnet-sdk = { path = "../pythnet/pythnet_sdk/", version = "2.0.0", features = ["strum"] }
Expand Down
18 changes: 11 additions & 7 deletions hermes/src/api.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,10 +11,14 @@ use {
routing::get,
Router,
},
ipnet::IpNet,
serde_qs::axum::QsQueryConfig,
std::sync::{
atomic::Ordering,
Arc,
std::{
net::SocketAddr,
sync::{
atomic::Ordering,
Arc,
},
},
tokio::{
signal,
Expand All @@ -36,10 +40,10 @@ pub struct ApiState {
}

impl ApiState {
pub fn new(state: Arc<State>) -> Self {
pub fn new(state: Arc<State>, ws_whitelist: Vec<IpNet>) -> Self {
Self {
state,
ws: Arc::new(ws::WsState::new()),
ws: Arc::new(ws::WsState::new(ws_whitelist)),
}
}
}
Expand Down Expand Up @@ -84,7 +88,7 @@ pub async fn run(
)]
struct ApiDoc;

let state = ApiState::new(state);
let state = ApiState::new(state, opts.rpc.ws_whitelist);

// Initialize Axum Router. Note the type here is a `Router<State>` due to the use of the
// `with_state` method which replaces `Body` with `State` in the type signature.
Expand Down Expand Up @@ -131,7 +135,7 @@ pub async fn run(
// Binds the axum's server to the configured address and port. This is a blocking call and will
// not return until the server is shutdown.
axum::Server::try_bind(&opts.rpc.addr)?
.serve(app.into_make_service())
.serve(app.into_make_service_with_connect_info::<SocketAddr>())
.with_graceful_shutdown(async {
// Ignore Ctrl+C errors, either way we need to shut down. The main Ctrl+C handler
// should also have triggered so we will let that one print the shutdown warning.
Expand Down
Loading

0 comments on commit 1a64d58

Please sign in to comment.