patch(hermes): improve ws reliability

- Add max message size for incoming messages
- Add sent message rate limit and ip whitelisting

Author: ali-behjati

hermes/Cargo.lock

@@ -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"
 
@@ -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"] }

hermes/Cargo.toml

@@ -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,
@@ -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)),
         }
     }
 }
@@ -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.
@@ -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.