snp.proto

@@ -191,8 +191,18 @@ message ExtendVmReq {
   uint64 locked_nano = 3;
 }
 
+message AirdropReq {
+  string pubkey = 1;
+  uint64 tokens = 2;
+}
+
+message Account {
+  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 Account);
 }

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(1_000_000_000));
     }
 
     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<Contract> {
+        let contracts = self.contracts.read().unwrap();
+        contracts.iter().cloned().collect()
+    }
+
+    pub fn list_accounts(&self) -> Vec<grpc::Account> {
+        self.accounts
+            .iter()
+            .map(|a| grpc::Account {
+                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<Contract> {
         debug!("Searching contracts for admin pubkey {admin_pubkey}");
         let contracts: Vec<Contract> = self

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",
+    "FHuecMbeC1PfjkW2JKyoicJAuiU7khgQT16QUB3Q1XdL",
+];
+
 pub struct BrainDaemonMock {
     data: Arc<BrainData>,
 }
@@ -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<Pubkey>) -> Result<Response<Empty>, 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<NewVmReq>) -> Result<Response<NewVmResp>, 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<AirdropReq>) -> Result<Response<Empty>, 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<Box<dyn Stream<Item = Result<Contract, Status>> + Send>>;
+    async fn list_all_contracts(
+        &self,
+        req: Request<Empty>,
+    ) -> Result<Response<Self::ListContractsStream>, 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<Box<dyn Stream<Item = Result<Account, Status>> + Send>>;
+    async fn list_accounts(
+        &self,
+        req: Request<Empty>,
+    ) -> Result<Response<Self::ListAccountsStream>, 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<String> {
+        None
+    }
+}
+
+impl PubkeyGetter for AirdropReq {
+    fn get_pubkey(&self) -> Option<String> {
+        None
+    }
+}
+
 fn check_sig_from_req<T: std::fmt::Debug + PubkeyGetter>(req: Request<T>) -> Result<T, Status> {
     let time = match req.metadata().get("timestamp") {
         Some(t) => t.clone(),
@@ -360,7 +422,7 @@ fn check_sig_from_req<T: std::fmt::Debug + PubkeyGetter>(req: Request<T>) -> 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<T>(req: &Request<T>) -> 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(())
+}