From 0ebfcc20bede86957deb863afd36c760ea391cf4 Mon Sep 17 00:00:00 2001
From: Jos Dehaes <jos.dehaes@gmail.com>
Date: Fri, 11 Apr 2025 11:56:25 +0200
Subject: [PATCH] feat: allow passing in 16 bytes for the JWT secret

---
 cli/Cargo.toml                |  3 +++
 cli/src/commands/start.rs     | 26 ++++++++++++++++++++++++--
 node/rest/src/helpers/auth.rs | 10 +++++++---
 3 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/cli/Cargo.toml b/cli/Cargo.toml
index e19cb7c70a..81066d141c 100644
--- a/cli/Cargo.toml
+++ b/cli/Cargo.toml
@@ -40,6 +40,9 @@ workspace = true
 [dependencies.anstyle]
 version = "1"
 
+[dependencies.base64]
+version = "0.22"
+
 [dependencies.anyhow]
 version = "1.0.79"
 
diff --git a/cli/src/commands/start.rs b/cli/src/commands/start.rs
index 5b76a62304..45764799bb 100644
--- a/cli/src/commands/start.rs
+++ b/cli/src/commands/start.rs
@@ -34,6 +34,7 @@ use snarkvm::{
 
 use aleo_std::StorageMode;
 use anyhow::{Result, bail, ensure};
+use base64::prelude::*;
 use clap::Parser;
 use colored::Colorize;
 use core::str::FromStr;
@@ -169,6 +170,14 @@ pub struct Start {
     /// If development mode is enabled, specify the custom bonded balances as a JSON object (default: None)
     #[clap(long)]
     pub dev_bonded_balances: Option<BondedBalances>,
+    /// Pass in an optional jwt secret for the node instance (16 bytes, base64 encoded) for keeping
+    /// the JWT constant
+    #[clap(long)]
+    pub jwt_secret: Option<String>,
+    /// Pass in an optional jwt creation timestamp for keeping the JWT constant. Can be any time in
+    /// the last 10 years
+    #[clap(long)]
+    pub jwt_timestamp: Option<i64>,
 }
 
 impl Start {
@@ -569,11 +578,24 @@ impl Start {
             );
 
             // If the node is running a REST server, print the REST IP and JWT.
-            if node_type.is_validator() {
+            if node_type.is_validator() || node_type.is_client() {
                 if let Some(rest_ip) = rest_ip {
                     println!("🌐 Starting the REST server at {}.\n", rest_ip.to_string().bold());
 
-                    if let Ok(jwt_token) = snarkos_node_rest::Claims::new(account.address()).to_jwt_string() {
+                    let jwt_secret = if let Some(jwt_b64) = &self.jwt_secret {
+                        if self.jwt_timestamp.is_none() {
+                            bail!("The '--jwt-timestamp' flag must be set if the '--jwt-secret' flag is set");
+                        }
+                        let jwt_bytes = BASE64_STANDARD.decode(jwt_b64).map_err(|_| anyhow::anyhow!("Invalid JWT secret"))?;
+                        if jwt_bytes.len() != 16 {
+                            bail!("The JWT secret must be 16 bytes long");
+                        }
+                        Some(jwt_bytes)
+                    } else {
+                        None
+                    };
+
+                    if let Ok(jwt_token) = snarkos_node_rest::Claims::new(account.address(), jwt_secret, self.jwt_timestamp).to_jwt_string() {
                         println!("🔑 Your one-time JWT token is {}\n", jwt_token.dimmed());
                     }
                 }
diff --git a/node/rest/src/helpers/auth.rs b/node/rest/src/helpers/auth.rs
index fe1ee0d491..26402815d4 100644
--- a/node/rest/src/helpers/auth.rs
+++ b/node/rest/src/helpers/auth.rs
@@ -32,12 +32,13 @@ use jsonwebtoken::{Algorithm, DecodingKey, EncodingKey, Header, Validation, deco
 use once_cell::sync::OnceCell;
 use serde::{Deserialize, Serialize};
 
+/// The JWT secret.
+static SECRET: OnceCell<Vec<u8>> = OnceCell::new();
 /// The time a jwt token is valid for.
 pub const EXPIRATION: i64 = 10 * 365 * 24 * 60 * 60; // 10 years.
 
 /// Returns the JWT secret for the node instance.
 fn jwt_secret() -> &'static Vec<u8> {
-    static SECRET: OnceCell<Vec<u8>> = OnceCell::new();
     SECRET.get_or_init(|| {
         let seed: [u8; 16] = ::rand::thread_rng().gen();
         seed.to_vec()
@@ -56,8 +57,11 @@ pub struct Claims {
 }
 
 impl Claims {
-    pub fn new<N: Network>(address: Address<N>) -> Self {
-        let issued_at = OffsetDateTime::now_utc().unix_timestamp();
+    pub fn new<N: Network>(address: Address<N>, jwt_secret: Option<Vec<u8>>, jwt_timestamp: Option<i64>) -> Self {
+        if let Some(secret) = jwt_secret {
+            SECRET.set(secret).expect("Failed to set JWT secret: already initialized");
+        }
+        let issued_at = jwt_timestamp.unwrap_or(OffsetDateTime::now_utc().unix_timestamp());
         let expiration = issued_at.saturating_add(EXPIRATION);
 
         Self { sub: address.to_string(), iat: issued_at, exp: expiration }