move invite logic to prime-core crate

Author: JannikSt

Cargo.lock

@@ -9,6 +9,7 @@ workspace = true
 [dependencies]
 p2p = { workspace = true}
 shared = { workspace = true }
+prime-core = { workspace = true }
 
 actix-web = { workspace = true }
 alloy = { workspace = true }

crates/orchestrator/Cargo.toml

@@ -3,19 +3,17 @@ use crate::models::node::OrchestratorNode;
 use crate::p2p::InviteRequest as InviteRequestWithMetadata;
 use crate::store::core::StoreContext;
 use crate::utils::loop_heartbeats::LoopHeartbeats;
-use alloy::primitives::utils::keccak256 as keccak;
-use alloy::primitives::U256;
-use alloy::signers::Signer;
 use anyhow::{bail, Result};
 use futures::stream;
 use futures::StreamExt;
 use log::{debug, error, info, warn};
-use p2p::InviteRequest;
 use p2p::InviteRequestUrl;
+use prime_core::invite::{
+    admin::{generate_invite_expiration, generate_invite_nonce, generate_invite_signature},
+    common::InviteBuilder,
+};
 use shared::web3::wallet::Wallet;
 use std::sync::Arc;
-use std::time::SystemTime;
-use std::time::UNIX_EPOCH;
 use tokio::sync::mpsc::Sender;
 use tokio::time::{interval, Duration};
 
@@ -89,29 +87,15 @@ impl NodeInviter {
         nonce: [u8; 32],
         expiration: [u8; 32],
     ) -> Result<[u8; 65]> {
-        let domain_id: [u8; 32] = U256::from(self.domain_id).to_be_bytes();
-        let pool_id: [u8; 32] = U256::from(self.pool_id).to_be_bytes();
-
-        let digest = keccak(
-            [
-                &domain_id,
-                &pool_id,
-                node.address.as_slice(),
-                &nonce,
-                &expiration,
-            ]
-            .concat(),
-        );
-
-        let signature = self
-            .wallet
-            .signer
-            .sign_message(digest.as_slice())
-            .await?
-            .as_bytes()
-            .to_owned();
-
-        Ok(signature)
+        generate_invite_signature(
+            &self.wallet,
+            self.domain_id,
+            self.pool_id,
+            node.address,
+            nonce,
+            expiration,
+        )
+        .await
     }
 
     async fn send_invite(&self, node: &OrchestratorNode) -> Result<(), anyhow::Error> {
@@ -122,29 +106,21 @@ impl NodeInviter {
         let p2p_addresses = node.worker_p2p_addresses.as_ref().unwrap();
 
         // Generate random nonce and expiration
-        let nonce: [u8; 32] = rand::random();
-        let expiration: [u8; 32] = U256::from(
-            SystemTime::now()
-                .duration_since(UNIX_EPOCH)
-                .map_err(|e| anyhow::anyhow!("System time error: {}", e))?
-                .as_secs()
-                + 1000,
-        )
-        .to_be_bytes();
+        let nonce = generate_invite_nonce();
+        let expiration = generate_invite_expiration(Some(1000))?;
 
         let invite_signature = self.generate_invite(node, nonce, expiration).await?;
-        let payload = InviteRequest {
-            invite: hex::encode(invite_signature),
-            pool_id: self.pool_id,
-            url: self.url.clone(),
-            timestamp: SystemTime::now()
-                .duration_since(UNIX_EPOCH)
-                .map_err(|e| anyhow::anyhow!("System time error: {}", e))?
-                .as_secs(),
-            expiration,
-            nonce,
+
+        // Build the invite request using the builder
+        let builder = match &self.url {
+            InviteRequestUrl::MasterUrl(url) => InviteBuilder::with_url(self.pool_id, url.clone()),
+            InviteRequestUrl::MasterIpPort(ip, port) => {
+                InviteBuilder::with_ip_port(self.pool_id, ip.clone(), *port)
+            }
         };
 
+        let payload = builder.build(invite_signature, nonce, expiration)?;
+
         info!("Sending invite to node: {p2p_id}");
 
         let (response_tx, response_rx) = tokio::sync::oneshot::channel();

crates/orchestrator/src/node/invite.rs

@@ -12,6 +12,7 @@ path = "src/lib.rs"
 
 [dependencies]
 shared = { workspace = true }
+p2p = { workspace = true }
 alloy = { workspace = true }
 alloy-provider = { workspace = true }
 serde = { workspace = true }

crates/prime-core/Cargo.toml

@@ -0,0 +1,81 @@
+use alloy::primitives::utils::keccak256 as keccak;
+use alloy::primitives::{Address, U256};
+use alloy::signers::Signer;
+use anyhow::Result;
+use rand_v8::prelude::*;
+use shared::web3::wallet::Wallet;
+use std::time::{SystemTime, UNIX_EPOCH};
+
+/// Generates an invite signature for a node
+///
+/// This function is used by pool owners/admins to create signed invites
+/// that authorize nodes to join their pool.
+pub async fn generate_invite_signature(
+    wallet: &Wallet,
+    domain_id: u32,
+    pool_id: u32,
+    node_address: Address,
+    nonce: [u8; 32],
+    expiration: [u8; 32],
+) -> Result<[u8; 65]> {
+    let domain_id_bytes: [u8; 32] = U256::from(domain_id).to_be_bytes();
+    let pool_id_bytes: [u8; 32] = U256::from(pool_id).to_be_bytes();
+
+    let digest = keccak(
+        [
+            &domain_id_bytes,
+            &pool_id_bytes,
+            node_address.as_slice(),
+            &nonce,
+            &expiration,
+        ]
+        .concat(),
+    );
+
+    let signature = wallet
+        .signer
+        .sign_message(digest.as_slice())
+        .await?
+        .as_bytes()
+        .to_owned();
+
+    Ok(signature)
+}
+
+/// Generates an invite expiration timestamp
+///
+/// Creates an expiration timestamp for an invite, defaulting to 1000 seconds from now
+pub fn generate_invite_expiration(seconds_from_now: Option<u64>) -> Result<[u8; 32]> {
+    let duration = seconds_from_now.unwrap_or(1000);
+    let expiration = U256::from(
+        SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .map_err(|e| anyhow::anyhow!("System time error: {}", e))?
+            .as_secs()
+            + duration,
+    );
+    Ok(expiration.to_be_bytes())
+}
+
+/// Generates a random nonce for invite
+pub fn generate_invite_nonce() -> [u8; 32] {
+    rand_v8::rngs::OsRng.gen()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_generate_invite_nonce() {
+        let nonce1 = generate_invite_nonce();
+        let nonce2 = generate_invite_nonce();
+        assert_ne!(nonce1, nonce2, "Nonces should be unique");
+    }
+
+    #[test]
+    fn test_generate_invite_expiration() {
+        let expiration = generate_invite_expiration(Some(3600)).unwrap();
+        assert_eq!(expiration.len(), 32);
+    }
+}

crates/prime-core/src/invite/admin.rs

@@ -0,0 +1,108 @@
+use alloy::primitives::Address;
+use anyhow::Result;
+use p2p::{InviteRequest, InviteRequestUrl};
+use std::time::{SystemTime, UNIX_EPOCH};
+
+/// Builder for creating invite requests
+pub struct InviteBuilder {
+    pool_id: u32,
+    url: InviteRequestUrl,
+}
+
+impl InviteBuilder {
+    /// Creates a new InviteBuilder with a master URL
+    pub fn with_url(pool_id: u32, url: String) -> Self {
+        Self {
+            pool_id,
+            url: InviteRequestUrl::MasterUrl(url),
+        }
+    }
+
+    /// Creates a new InviteBuilder with IP and port
+    pub fn with_ip_port(pool_id: u32, ip: String, port: u16) -> Self {
+        Self {
+            pool_id,
+            url: InviteRequestUrl::MasterIpPort(ip, port),
+        }
+    }
+
+    /// Builds an InviteRequest with the given parameters
+    pub fn build(
+        self,
+        invite_signature: [u8; 65],
+        nonce: [u8; 32],
+        expiration: [u8; 32],
+    ) -> Result<InviteRequest> {
+        Ok(InviteRequest {
+            invite: hex::encode(invite_signature),
+            pool_id: self.pool_id,
+            url: self.url,
+            timestamp: SystemTime::now()
+                .duration_since(UNIX_EPOCH)
+                .map_err(|e| anyhow::anyhow!("System time error: {}", e))?
+                .as_secs(),
+            expiration,
+            nonce,
+        })
+    }
+}
+
+/// Metadata for an invite request that includes worker information
+#[derive(Debug, Clone)]
+pub struct InviteMetadata {
+    pub worker_wallet_address: Address,
+    pub worker_p2p_id: String,
+    pub worker_addresses: Vec<String>,
+}
+
+/// Full invite request with metadata
+#[derive(Debug)]
+pub struct InviteWithMetadata {
+    pub metadata: InviteMetadata,
+    pub request: InviteRequest,
+}
+
+/// Helper to parse InviteRequestUrl into a usable endpoint
+pub fn get_endpoint_from_url(url: &InviteRequestUrl, path: &str) -> String {
+    match url {
+        InviteRequestUrl::MasterIpPort(ip, port) => {
+            format!("http://{ip}:{port}/{path}")
+        }
+        InviteRequestUrl::MasterUrl(url) => {
+            let url = url.trim_end_matches('/');
+            format!("{url}/{path}")
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_invite_builder_with_url() {
+        let builder = InviteBuilder::with_url(1, "https://example.com".to_string());
+        let signature = [0u8; 65];
+        let nonce = [1u8; 32];
+        let expiration = [2u8; 32];
+
+        let invite = builder.build(signature, nonce, expiration).unwrap();
+        assert_eq!(invite.pool_id, 1);
+        assert!(matches!(invite.url, InviteRequestUrl::MasterUrl(_)));
+    }
+
+    #[test]
+    fn test_get_endpoint_from_url() {
+        let url = InviteRequestUrl::MasterUrl("https://example.com".to_string());
+        assert_eq!(
+            get_endpoint_from_url(&url, "heartbeat"),
+            "https://example.com/heartbeat"
+        );
+
+        let ip_port = InviteRequestUrl::MasterIpPort("192.168.1.1".to_string(), 8080);
+        assert_eq!(
+            get_endpoint_from_url(&ip_port, "heartbeat"),
+            "http://192.168.1.1:8080/heartbeat"
+        );
+    }
+}

crates/prime-core/src/invite/common.rs

@@ -0,0 +1,7 @@
+pub mod admin;
+pub mod common;
+pub mod worker;
+
+pub use admin::*;
+pub use common::*;
+pub use worker::*;