From 514da066f4fd503edcfae41761c7e1fbfc6d33ee Mon Sep 17 00:00:00 2001
From: steebchen <contact@luca-steeb.com>
Date: Wed, 11 Sep 2024 17:01:39 -0400
Subject: [PATCH] Refactor secret key handling

Extract secret key parsing to a separate function `get_secret_key` for reuse and clarity. This change ensures consistent secret key retrieval in both `vrf_info` and `stark_vrf` functions. Simplified the secret key documentation by removing type details.
---
 src/main.rs | 21 +++++++++++++--------
 1 file changed, 13 insertions(+), 8 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index f63df6f..92af7f7 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -18,7 +18,7 @@ use tracing::debug;
 #[derive(Parser, Debug)]
 #[command(version, about, long_about = None)]
 struct Args {
-    /// Secret key as a u64
+    /// Secret key
     #[arg(short, long, required = true)]
     secret_key: u64,
 }
@@ -34,10 +34,14 @@ struct InfoResult {
     public_key_y: String,
 }
 
-async fn vrf_info() -> Json<InfoResult> {
+fn get_secret_key() -> String {
     let args = Args::parse();
-    let secret_key = args.secret_key.to_string().parse().unwrap();
-    let public_key = generate_public_key(secret_key);
+    args.secret_key.to_string()
+}
+
+async fn vrf_info() -> Json<InfoResult> {
+    let secret_key = get_secret_key();
+    let public_key = generate_public_key(secret_key.parse().unwrap());
 
     Json(InfoResult {
         public_key_x: format(public_key.x),
@@ -52,9 +56,8 @@ struct JsonResult {
 
 async fn stark_vrf(extract::Json(payload): extract::Json<StarkVrfRequest>) -> Json<JsonResult> {
     debug!("received payload {payload:?}");
-    let args = Args::parse();
-    let secret_key = args.secret_key.to_string().parse().unwrap();
-    let public_key = generate_public_key(secret_key);
+    let secret_key = get_secret_key();
+    let public_key = generate_public_key(secret_key.parse().unwrap());
 
     println!("public key {public_key}");
 
@@ -69,7 +72,9 @@ async fn stark_vrf(extract::Json(payload): extract::Json<StarkVrfRequest>) -> Js
         .collect();
 
     let ecvrf = StarkVRF::new(public_key).unwrap();
-    let proof = ecvrf.prove(&secret_key, seed.as_slice()).unwrap();
+    let proof = ecvrf
+        .prove(&secret_key.parse().unwrap(), seed.as_slice())
+        .unwrap();
     let sqrt_ratio_hint = ecvrf.hash_to_sqrt_ratio_hint(seed.as_slice());
     let rnd = ecvrf.proof_to_hash(&proof).unwrap();