From 457456fbe58dc839ec660e68c99eee1afaeea0a8 Mon Sep 17 00:00:00 2001 From: ghe0 Date: Wed, 5 Feb 2025 02:44:42 +0200 Subject: [PATCH] added admin functionality --- snp.proto | 16 ++++++++- src/data.rs | 20 +++++++++-- src/grpc.rs | 96 ++++++++++++++++++++++++++++++++++++++++++++++++----- 3 files changed, 121 insertions(+), 11 deletions(-) diff --git a/snp.proto b/snp.proto index fa9654e..c1f8ad7 100644 --- a/snp.proto +++ b/snp.proto @@ -191,8 +191,18 @@ message ExtendVmReq { uint64 locked_nano = 3; } +message AirdropReq { + string pubkey = 1; + uint64 tokens = 2; +} + +message ListAccountsResp { + string pubkey = 1; + uint64 balance = 2; + uint64 tmp_locked = 3; +} + service BrainCli { - rpc GetAirdrop (Pubkey) returns (Empty); rpc GetBalance (Pubkey) returns (AccountBalance); rpc NewVm (NewVmReq) returns (NewVmResp); rpc ListContracts (ListContractsReq) returns (stream Contract); @@ -201,4 +211,8 @@ service BrainCli { rpc DeleteVm (DeleteVmReq) returns (Empty); rpc UpdateVm (UpdateVmReq) returns (UpdateVmResp); rpc ExtendVm (ExtendVmReq) returns (Empty); + // admin commands + rpc Airdrop (AirdropReq) returns (Empty); + rpc ListAllContracts (Empty) returns (stream Contract); + rpc ListAccounts (Empty) returns (stream ListAccountsResp); } diff --git a/src/data.rs b/src/data.rs index d286c52..32ba306 100644 --- a/src/data.rs +++ b/src/data.rs @@ -176,8 +176,8 @@ impl BrainData { } } - pub fn get_airdrop(&self, account: &str) { - self.add_nano_to_wallet(account, 1000_000000000); + pub fn give_airdrop(&self, account: &str, tokens: u64) { + self.add_nano_to_wallet(account, tokens.saturating_mul(1000_000000000)); } fn add_nano_to_wallet(&self, account: &str, nano_lp: u64) { @@ -684,6 +684,22 @@ impl BrainData { contracts.iter().cloned().find(|c| c.uuid == uuid) } + pub fn list_all_contracts(&self) -> Vec { + let contracts = self.contracts.read().unwrap(); + contracts.iter().cloned().collect() + } + + pub fn list_accounts(&self) -> Vec { + self.accounts + .iter() + .map(|a| grpc::ListAccountsResp { + pubkey: a.key().to_string(), + balance: a.balance, + tmp_locked: a.tmp_locked, + }) + .collect() + } + pub fn find_contracts_by_admin_pubkey(&self, admin_pubkey: &str) -> Vec { debug!("Searching contracts for admin pubkey {admin_pubkey}"); let contracts: Vec = self diff --git a/src/grpc.rs b/src/grpc.rs index 6d3d102..c89daa6 100644 --- a/src/grpc.rs +++ b/src/grpc.rs @@ -15,6 +15,11 @@ use tokio::sync::mpsc; use tokio_stream::{wrappers::ReceiverStream, Stream, StreamExt}; use tonic::{Request, Response, Status, Streaming}; +const ADMIN_ACCOUNTS: &[&str] = &[ + "x52w7jARC5erhWWK65VZmjdGXzBK6ZDgfv1A283d8XK", + "ThisIsNotARealWallet000000000000000000000000", +]; + pub struct BrainDaemonMock { data: Arc, } @@ -99,7 +104,10 @@ impl BrainDaemon for BrainDaemonMock { let mut req_stream = req.into_inner(); let pubkey: String; if let Some(Ok(msg)) = req_stream.next().await { - log::debug!("demon_messages received the following auth message: {:?}", msg.msg); + log::debug!( + "demon_messages received the following auth message: {:?}", + msg.msg + ); if let Some(daemon_message::Msg::Auth(auth)) = msg.msg { pubkey = auth.pubkey.clone(); check_sig_from_parts( @@ -149,12 +157,6 @@ impl BrainCli for BrainCliMock { Ok(Response::new(self.data.get_balance(&req.pubkey).into())) } - async fn get_airdrop(&self, req: Request) -> Result, Status> { - let req = check_sig_from_req(req)?; - self.data.get_airdrop(&req.pubkey); - Ok(Response::new(Empty {})) - } - async fn new_vm(&self, req: Request) -> Result, Status> { let req = check_sig_from_req(req)?; info!("New VM requested via CLI: {req:?}"); @@ -270,6 +272,54 @@ impl BrainCli for BrainCliMock { Err(e) => Err(Status::not_found(e.to_string())), } } + + async fn airdrop(&self, req: Request) -> Result, Status> { + check_admin_key(&req)?; + let req = check_sig_from_req(req)?; + self.data.give_airdrop(&req.pubkey, req.tokens); + Ok(Response::new(Empty{})) + } + + type ListAllContractsStream = Pin> + Send>>; + async fn list_all_contracts( + &self, + req: Request, + ) -> Result, Status> { + check_admin_key(&req)?; + let _ = check_sig_from_req(req)?; + let contracts = self.data.list_all_contracts(); + let (tx, rx) = mpsc::channel(6); + tokio::spawn(async move { + for contract in contracts { + let _ = tx.send(Ok(contract.into())).await; + } + }); + let output_stream = ReceiverStream::new(rx); + Ok(Response::new( + Box::pin(output_stream) as Self::ListContractsStream + )) + } + + type ListAccountsStream = Pin> + Send>>; + async fn list_accounts( + &self, + req: Request, + ) -> Result, Status> { + check_admin_key(&req)?; + let _ = check_sig_from_req(req)?; + let accounts = self.data.list_accounts(); + let (tx, rx) = mpsc::channel(6); + tokio::spawn(async move { + for account in accounts { + let _ = tx.send(Ok(account.into())).await; + } + }); + let output_stream = ReceiverStream::new(rx); + Ok(Response::new( + Box::pin(output_stream) as Self::ListAccountsStream + )) + } + } trait PubkeyGetter { @@ -324,6 +374,18 @@ impl PubkeyGetter for RegisterNodeReq { } } +impl PubkeyGetter for Empty { + fn get_pubkey(&self) -> Option { + None + } +} + +impl PubkeyGetter for AirdropReq { + fn get_pubkey(&self) -> Option { + None + } +} + fn check_sig_from_req(req: Request) -> Result { let time = match req.metadata().get("timestamp") { Some(t) => t.clone(), @@ -360,7 +422,7 @@ fn check_sig_from_req(req: Request) -> Res let pubkey_value = match req.metadata().get("pubkey") { Some(p) => p.clone(), - None => return Err(Status::unauthenticated("Signature not found in metadata.")), + None => return Err(Status::unauthenticated("pubkey not found in metadata.")), }; let pubkey = ed25519_dalek::VerifyingKey::from_bytes( &bs58::decode(&pubkey_value) @@ -426,3 +488,21 @@ fn check_sig_from_parts(pubkey: &str, time: &str, msg: &str, sig: &str) -> Resul Ok(()) } + +fn check_admin_key(req: &Request) -> Result<(), Status> { + let pubkey = match req.metadata().get("pubkey") { + Some(p) => p.clone(), + None => return Err(Status::unauthenticated("pubkey not found in metadata.")), + }; + let pubkey = pubkey + .to_str() + .map_err(|_| Status::unauthenticated("could not parse pubkey metadata to str"))?; + + if !ADMIN_ACCOUNTS.contains(&pubkey) { + return Err(Status::unauthenticated( + "This operation is reserved to admin accounts", + )); + } + + Ok(()) +}