From 772f1c4b28cfe2a1f25ad7ee4577df6e9d1d89b3 Mon Sep 17 00:00:00 2001
From: red binder <redbinder0526@proton.me>
Date: Tue, 23 Jun 2026 11:43:21 +0200
Subject: [PATCH] update ActStage to derive nex_key on its own

---
 Cargo.lock  |  2 +-
 Cargo.toml  |  2 +-
 grpc.proto  |  7 ++++++-
 src/main.rs | 32 +++++++++++++++++++++++++++-----
 4 files changed, 35 insertions(+), 8 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index e021f70..e7f0f09 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -918,7 +918,7 @@ checksum = "1d87ecb2933e8aeadb3e3a02b828fed80a7528047e68b4f424523a0981a3a084"
 
 [[package]]
 name = "nex-account"
-version = "0.2.0"
+version = "0.2.1"
 dependencies = [
  "hmac",
  "log",
diff --git a/Cargo.toml b/Cargo.toml
index d3207bf..d11502d 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "nex-account"
-version = "0.2.0"
+version = "0.2.1"
 edition = "2024"
 publish = ["spbr"]
 
diff --git a/grpc.proto b/grpc.proto
index dd31fcb..ddbddbf 100644
--- a/grpc.proto
+++ b/grpc.proto
@@ -10,7 +10,12 @@ message ActCreateInfo {
 }
 
 message ActStageInfo {
+  bytes password = 1;
+}
+
+message ActStageReturn {
   bytes nex_key = 1;
+  int32 pid = 2;
 }
 
 message PID {
@@ -24,5 +29,5 @@ message NexKey {
 service NexAccountService {
   rpc CreateNewSequentialOrUpdateAndGetAccount(ActCreateInfo) returns (NexKey);
   rpc GetNexKeyByPid(PID) returns (NexKey);
-  rpc StageNewAccount(ActStageInfo) returns (PID);
+  rpc StageNewAccount(ActStageInfo) returns (ActStageReturn);
 }
diff --git a/src/main.rs b/src/main.rs
index 16553c0..6a59a7a 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -8,9 +8,10 @@ use simplelog::{Config, TerminalMode};
 use sqlx::{PgPool, query};
 use tonic::{Response, transport::Server};
 use crate::grpc::{
-    ActCreateInfo, NexKey, Pid, ActStageInfo,
+    ActCreateInfo, NexKey, Pid, ActStageInfo, ActStageReturn,
     nex_account_service_server::{NexAccountService, NexAccountServiceServer},
 };
+use md5::{Digest, Md5};
 
 pub mod grpc;
 
@@ -50,6 +51,25 @@ fn db_neverfail(sql_err: sqlx::Error) -> tonic::Status {
     tonic::Status::aborted("error in database")
 }
 
+pub fn derive_key(pid: i32, password: &[u8]) -> [u8; 16] {
+    let iteration_count = 65000 + pid % 1024;
+    // we do one iteration out here to ensure the key is always 16 bytes
+
+    let mut key: [u8; 16] = {
+        let mut md5 = Md5::new();
+        md5.update(password);
+        md5.finalize().try_into().unwrap()
+    };
+
+    for _ in 1..iteration_count {
+        let mut md5 = Md5::new();
+        md5.update(key);
+        key = md5.finalize().try_into().unwrap();
+    }
+
+    key
+}
+
 #[tonic::async_trait]
 impl NexAccountService for NexAccountServer {
     async fn create_new_sequential_or_update_and_get_account(
@@ -139,12 +159,14 @@ impl NexAccountService for NexAccountServer {
     async fn stage_new_account(
         &self,
         request: tonic::Request<ActStageInfo>,
-    ) -> Result<tonic::Response<Pid>, tonic::Status> {
+    ) -> Result<tonic::Response<ActStageReturn>, tonic::Status> {
         let request = request.into_inner();
 
         let next_pid = next_pid(&self.pool).await;
 
-        if request.nex_key.len() != 16 {
+        let nex_key = derive_key(next_pid, &request.password);
+
+        if nex_key.len() != 16 {
             return Err(tonic::Status::invalid_argument("invalid key length"));
         }
 
@@ -155,13 +177,13 @@ impl NexAccountService for NexAccountServer {
                 $1, $2
             )",
             next_pid,
-            request.nex_key
+            &nex_key
         )
             .execute(&self.pool)
             .await
             .map_err(db_neverfail)?;
 
-        Ok(Response::new(Pid { pid: next_pid}))
+        Ok(Response::new(ActStageReturn{ pid: next_pid, nex_key: nex_key.into() }))
     }
 }