Skip to content

Commit a095ca9

Browse files
JannikStJannikSt
authored
imp: add nonce to auth middleware (#489)
* add nonce to auth middleware
1 parent 24e7e89 commit a095ca9

File tree

11 files changed

+402
-115
lines changed

11 files changed

+402
-115
lines changed

‎crates/discovery/src/api/routes/node.rs‎

Lines changed: 28 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -225,7 +225,7 @@ mod tests {
225225
use actix_web::App;
226226
use shared::models::node::{ComputeSpecs, CpuSpecs, DiscoveryNode, GpuSpecs};
227227
use shared::security::auth_signature_middleware::{ValidateSignature, ValidatorState};
228-
use shared::security::request_signer::sign_request;
228+
use shared::security::request_signer::sign_request_with_nonce;
229229
use shared::web3::wallet::Wallet;
230230
use std::sync::Arc;
231231
use std::time::SystemTime;
@@ -318,7 +318,7 @@ mod tests {
318318
.await;
319319

320320
let json = serde_json::to_value(node.clone()).unwrap();
321-
let signature = sign_request(
321+
let signed_request = sign_request_with_nonce(
322322
"/nodes",
323323
&Wallet::new(private_key, Url::parse("http://localhost:8080").unwrap()).unwrap(),
324324
Some(&json),
@@ -328,9 +328,9 @@ mod tests {
328328

329329
let req = test::TestRequest::put()
330330
.uri("/nodes")
331-
.set_json(json)
331+
.set_json(signed_request.data.as_ref().unwrap())
332332
.insert_header(("x-address", node.id.clone()))
333-
.insert_header(("x-signature", signature))
333+
.insert_header(("x-signature", signed_request.signature))
334334
.to_request();
335335

336336
let resp = test::call_service(&app, req).await;
@@ -381,7 +381,7 @@ mod tests {
381381
}
382382

383383
let json = serde_json::to_value(node_clone_for_recall.clone()).unwrap();
384-
let signature = sign_request(
384+
let signed_request = sign_request_with_nonce(
385385
"/nodes",
386386
&Wallet::new(private_key, Url::parse("http://localhost:8080").unwrap()).unwrap(),
387387
Some(&json),
@@ -391,9 +391,9 @@ mod tests {
391391

392392
let req = test::TestRequest::put()
393393
.uri("/nodes")
394-
.set_json(json)
394+
.set_json(signed_request.data.as_ref().unwrap())
395395
.insert_header(("x-address", node_clone_for_recall.id.clone()))
396-
.insert_header(("x-signature", signature))
396+
.insert_header(("x-signature", signed_request.signature))
397397
.to_request();
398398

399399
let resp = test::call_service(&app, req).await;
@@ -443,7 +443,7 @@ mod tests {
443443
.await;
444444

445445
let json = serde_json::to_value(node.clone()).unwrap();
446-
let signature = sign_request(
446+
let signed_request = sign_request_with_nonce(
447447
"/nodes",
448448
&Wallet::new(private_key, Url::parse("http://localhost:8080").unwrap()).unwrap(),
449449
Some(&json),
@@ -453,9 +453,9 @@ mod tests {
453453

454454
let req = test::TestRequest::put()
455455
.uri("/nodes")
456-
.set_json(json)
456+
.set_json(signed_request.data.as_ref().unwrap())
457457
.insert_header(("x-address", node.id.clone()))
458-
.insert_header(("x-signature", signature))
458+
.insert_header(("x-signature", signed_request.signature))
459459
.to_request();
460460

461461
let resp = test::call_service(&app, req).await;
@@ -508,7 +508,7 @@ mod tests {
508508
.await;
509509

510510
let json = serde_json::to_value(node.clone()).unwrap();
511-
let signature = sign_request(
511+
let signed_request = sign_request_with_nonce(
512512
"/nodes",
513513
&Wallet::new(private_key, Url::parse("http://localhost:8080").unwrap()).unwrap(),
514514
Some(&json),
@@ -518,9 +518,9 @@ mod tests {
518518

519519
let req = test::TestRequest::put()
520520
.uri("/nodes")
521-
.set_json(json)
521+
.set_json(signed_request.data.as_ref().unwrap())
522522
.insert_header(("x-address", "0x7E5F4552091A69125d5DfCb7b8C2659029395Bdf"))
523-
.insert_header(("x-signature", signature))
523+
.insert_header(("x-signature", signed_request.signature))
524524
.to_request();
525525

526526
let resp = test::call_service(&app, req).await;
@@ -587,7 +587,7 @@ mod tests {
587587
.await;
588588

589589
let json = serde_json::to_value(node.clone()).unwrap();
590-
let signature = sign_request(
590+
let signed_request = sign_request_with_nonce(
591591
"/nodes",
592592
&Wallet::new(private_key, Url::parse("http://localhost:8080").unwrap()).unwrap(),
593593
Some(&json),
@@ -597,9 +597,9 @@ mod tests {
597597

598598
let req = test::TestRequest::put()
599599
.uri("/nodes")
600-
.set_json(json)
600+
.set_json(signed_request.data.as_ref().unwrap())
601601
.insert_header(("x-address", node.id.clone()))
602-
.insert_header(("x-signature", signature))
602+
.insert_header(("x-signature", signed_request.signature))
603603
.to_request();
604604

605605
let resp = test::call_service(&app, req).await;
@@ -669,7 +669,7 @@ mod tests {
669669

670670
// Register first node - should succeed
671671
let json1 = serde_json::to_value(node1.clone()).unwrap();
672-
let signature1 = sign_request(
672+
let signature1 = sign_request_with_nonce(
673673
"/nodes",
674674
&Wallet::new(private_key, Url::parse("http://localhost:8080").unwrap()).unwrap(),
675675
Some(&json1),
@@ -679,17 +679,17 @@ mod tests {
679679

680680
let req1 = test::TestRequest::put()
681681
.uri("/nodes")
682-
.set_json(json1)
682+
.set_json(signature1.data)
683683
.insert_header(("x-address", node1.id.clone()))
684-
.insert_header(("x-signature", signature1))
684+
.insert_header(("x-signature", signature1.signature))
685685
.to_request();
686686

687687
let resp1 = test::call_service(&app, req1).await;
688688
assert_eq!(resp1.status(), StatusCode::OK);
689689

690690
// Try to register same node again - should succeed (update)
691691
let json1_duplicate = serde_json::to_value(node1.clone()).unwrap();
692-
let signature1_duplicate = sign_request(
692+
let signature1_duplicate = sign_request_with_nonce(
693693
"/nodes",
694694
&Wallet::new(private_key, Url::parse("http://localhost:8080").unwrap()).unwrap(),
695695
Some(&json1_duplicate),
@@ -699,17 +699,17 @@ mod tests {
699699

700700
let req1_duplicate = test::TestRequest::put()
701701
.uri("/nodes")
702-
.set_json(json1_duplicate)
702+
.set_json(signature1_duplicate.data)
703703
.insert_header(("x-address", node1.id.clone()))
704-
.insert_header(("x-signature", signature1_duplicate))
704+
.insert_header(("x-signature", signature1_duplicate.signature))
705705
.to_request();
706706

707707
let resp1_duplicate = test::call_service(&app, req1_duplicate).await;
708708
assert_eq!(resp1_duplicate.status(), StatusCode::OK);
709709

710710
// Register second node with different ID - should succeed
711711
let json2 = serde_json::to_value(node2.clone()).unwrap();
712-
let signature2 = sign_request(
712+
let signature2 = sign_request_with_nonce(
713713
"/nodes",
714714
&Wallet::new(private_key_2, Url::parse("http://localhost:8080").unwrap()).unwrap(),
715715
Some(&json2),
@@ -719,9 +719,9 @@ mod tests {
719719

720720
let req2 = test::TestRequest::put()
721721
.uri("/nodes")
722-
.set_json(json2)
722+
.set_json(signature2.data)
723723
.insert_header(("x-address", node2.id.clone()))
724-
.insert_header(("x-signature", signature2))
724+
.insert_header(("x-signature", signature2.signature))
725725
.to_request();
726726

727727
let resp2 = test::call_service(&app, req2).await;
@@ -745,7 +745,7 @@ mod tests {
745745

746746
// Register third node - should fail (exceeds max_nodes_per_ip)
747747
let json3 = serde_json::to_value(node3.clone()).unwrap();
748-
let signature3 = sign_request(
748+
let signature3 = sign_request_with_nonce(
749749
"/nodes",
750750
&Wallet::new(private_key_3, Url::parse("http://localhost:8080").unwrap()).unwrap(),
751751
Some(&json3),
@@ -755,9 +755,9 @@ mod tests {
755755

756756
let req3 = test::TestRequest::put()
757757
.uri("/nodes")
758-
.set_json(json3)
758+
.set_json(signature3.data)
759759
.insert_header(("x-address", node3.id.clone()))
760-
.insert_header(("x-signature", signature3))
760+
.insert_header(("x-signature", signature3.signature))
761761
.to_request();
762762

763763
let resp3 = test::call_service(&app, req3).await;

‎crates/discovery/src/api/server.rs‎

Lines changed: 20 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
use crate::api::routes::get_nodes::{get_node_by_subkey, get_nodes, get_nodes_for_pool};
22
use crate::api::routes::node::node_routes;
33
use crate::store::node_store::NodeStore;
4+
use crate::store::redis::RedisStore;
45
use actix_web::middleware::{Compress, NormalizePath, TrailingSlash};
56
use actix_web::HttpResponse;
67
use actix_web::{
@@ -71,10 +72,12 @@ async fn health_check(app_state: web::Data<AppState>) -> HttpResponse {
7172
}))
7273
}
7374

75+
#[allow(clippy::too_many_arguments)]
7476
pub async fn start_server(
7577
host: &str,
7678
port: u16,
7779
node_store: Arc<NodeStore>,
80+
redis_store: Arc<RedisStore>,
7881
contracts: Contracts<RootProvider>,
7982
platform_api_key: String,
8083
last_chain_sync: Arc<Mutex<Option<SystemTime>>>,
@@ -97,11 +100,23 @@ pub async fn start_server(
97100
max_nodes_per_ip,
98101
};
99102

100-
// it seems we have a validator for the validator
101-
let validator_validator = Arc::new(ValidatorState::new(validators));
102-
103-
// All nodes can register as long as they have a valid signature
104-
let validate_signatures = Arc::new(ValidatorState::new(vec![]).with_validator(move |_| true));
103+
let validator_validator = Arc::new(
104+
ValidatorState::new(validators)
105+
.with_redis(redis_store.client.clone())
106+
.await
107+
.map_err(|e| {
108+
std::io::Error::other(format!("Failed to initialize Redis connection pool: {}", e))
109+
})?,
110+
);
111+
let validate_signatures = Arc::new(
112+
ValidatorState::new(vec![])
113+
.with_redis(redis_store.client.clone())
114+
.await
115+
.map_err(|e| {
116+
std::io::Error::other(format!("Failed to initialize Redis connection pool: {}", e))
117+
})?
118+
.with_validator(move |_| true),
119+
);
105120
let api_key_middleware = Arc::new(ApiKeyMiddleware::new(platform_api_key));
106121

107122
HttpServer::new(move || {

‎crates/discovery/src/main.rs‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -48,8 +48,8 @@ async fn main() -> Result<()> {
4848

4949
let args = Args::parse();
5050

51-
let redis_store = RedisStore::new(&args.redis_url);
52-
let node_store = Arc::new(NodeStore::new(redis_store));
51+
let redis_store = Arc::new(RedisStore::new(&args.redis_url));
52+
let node_store = Arc::new(NodeStore::new(redis_store.as_ref().clone()));
5353
let endpoint = match args.rpc_url.parse() {
5454
Ok(url) => url,
5555
Err(_) => {
@@ -82,6 +82,7 @@ async fn main() -> Result<()> {
8282
"0.0.0.0",
8383
args.port,
8484
node_store,
85+
redis_store,
8586
contracts,
8687
args.platform_api_key,
8788
last_chain_sync,

‎crates/orchestrator/src/api/server.rs‎

Lines changed: 16 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -72,18 +72,22 @@ pub async fn start_server(
7272
});
7373
let node_store = app_state.store_context.node_store.clone();
7474
let node_store_clone = node_store.clone();
75-
let validator_state = Arc::new(ValidatorState::new(vec![]).with_async_validator(
76-
move |address| {
77-
let address = *address;
78-
let node_store = node_store_clone.clone();
79-
Box::pin(async move {
80-
match node_store.get_node(&address).await {
81-
Ok(Some(node)) => node.status != NodeStatus::Ejected,
82-
_ => false,
83-
}
84-
})
85-
},
86-
));
75+
let validator_state = Arc::new(
76+
ValidatorState::new(vec![])
77+
.with_redis(app_state.redis_store.client.clone())
78+
.await
79+
.map_err(|e| anyhow::anyhow!("Failed to initialize Redis connection pool: {}", e))?
80+
.with_async_validator(move |address| {
81+
let address = *address;
82+
let node_store = node_store_clone.clone();
83+
Box::pin(async move {
84+
match node_store.get_node(&address).await {
85+
Ok(Some(node)) => node.status != NodeStatus::Ejected,
86+
_ => false,
87+
}
88+
})
89+
}),
90+
);
8791

8892
let api_key_middleware = Arc::new(ApiKeyMiddleware::new(admin_api_key));
8993
HttpServer::new(move || {

‎crates/orchestrator/src/discovery/monitor.rs‎

Lines changed: 11 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ use log::{error, info};
1111
use serde_json;
1212
use shared::models::api::ApiResponse;
1313
use shared::models::node::DiscoveryNode;
14-
use shared::security::request_signer::sign_request;
14+
use shared::security::request_signer::sign_request_with_nonce;
1515
use shared::web3::wallet::Wallet;
1616
use std::sync::Arc;
1717
use std::time::Duration;
@@ -110,13 +110,14 @@ impl DiscoveryMonitor {
110110
let discovery_route = format!("/api/pool/{}", self.compute_pool_id);
111111
let address = self.coordinator_wallet.address().to_string();
112112

113-
let signature = match sign_request(&discovery_route, &self.coordinator_wallet, None).await {
114-
Ok(sig) => sig,
115-
Err(e) => {
116-
error!("Failed to sign discovery request: {}", e);
117-
return Ok(Vec::new());
118-
}
119-
};
113+
let signature =
114+
match sign_request_with_nonce(&discovery_route, &self.coordinator_wallet, None).await {
115+
Ok(sig) => sig,
116+
Err(e) => {
117+
error!("Failed to sign discovery request: {}", e);
118+
return Ok(Vec::new());
119+
}
120+
};
120121

121122
let mut headers = reqwest::header::HeaderMap::new();
122123
headers.insert(
@@ -125,12 +126,13 @@ impl DiscoveryMonitor {
125126
);
126127
headers.insert(
127128
"x-signature",
128-
reqwest::header::HeaderValue::from_str(&signature)?,
129+
reqwest::header::HeaderValue::from_str(&signature.signature)?,
129130
);
130131

131132
let response = match self
132133
.http_client
133134
.get(format!("{}{}", self.discovery_url, discovery_route))
135+
.query(&[("nonce", signature.nonce)])
134136
.headers(headers)
135137
.send()
136138
.await

‎crates/shared/Cargo.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ actix-web = { workspace = true }
2222
futures-util = { workspace = true }
2323
hex = { workspace = true }
2424
uuid = { workspace = true }
25-
redis = { workspace = true }
25+
redis = { workspace = true, features = ["aio", "tokio-comp"] }
2626
dashmap = "6.1.0"
2727
anyhow = { workspace = true }
2828
nalgebra = { workspace = true }

0 commit comments

Comments
 (0)