From 30b53b379a7b4e032970dd5cf08d1244f534c343 Mon Sep 17 00:00:00 2001
From: DefiCake <rarore92@gmail.com>
Date: Fri, 1 Dec 2023 17:09:28 +0100
Subject: [PATCH] feat: add GET /session endpoint

---
 Cargo.toml     |  2 +-
 src/lib.rs     |  4 ++-
 src/routes.rs  | 66 ++++++++++++++++++++++++++++++++++++++++++++------
 src/session.rs |  4 +++
 4 files changed, 67 insertions(+), 9 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 0ede984..1d263db 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -28,9 +28,9 @@ serde_json = { version = "1.0" }
 tokio = { version = "1.0", features = ["full"] }
 tower = { version = "0.4", features = ["buffer", "limit", "load-shed", "util", "timeout"] }
 tower-http = { version = "0.2.5", features = ["cors", "trace", "set-header", "fs"] }
+rand = "0.8.5"
 tracing = "0.1"
 tracing-subscriber = { version = "0.3", features = ["env-filter", "json"] }
-rand = "0.8.5"
 
 [dev-dependencies]
 fuel-core = { version = "0.21.0", default-features = false }
diff --git a/src/lib.rs b/src/lib.rs
index f2ce06e..84f22c2 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -192,6 +192,8 @@ pub async fn start_server(
                 )
                 .into_inner(),
         )
+        .route("/session", get(routes::get_session))
+        .layer(Extension(sessions.clone()))
         .route("/session", post(routes::create_session))
         .layer(
             ServiceBuilder::new()
@@ -201,7 +203,7 @@ pub async fn start_server(
                 .concurrency_limit(MAX_CONCURRENT_REQUESTS)
                 .timeout(Duration::from_secs(60))
                 .layer(TraceLayer::new_for_http())
-                .layer(Extension(sessions))
+                .layer(Extension(sessions.clone()))
                 .into_inner(),
         );
 
diff --git a/src/routes.rs b/src/routes.rs
index 804eb14..a8a12a1 100644
--- a/src/routes.rs
+++ b/src/routes.rs
@@ -5,6 +5,7 @@ use crate::{
     CoinOutput, SharedConfig, SharedFaucetState, SharedNetworkConfig, SharedWallet,
 };
 use axum::{
+    extract::Query,
     response::{Html, IntoResponse, Response},
     routing::{get_service, MethodRouter},
     Extension, Json,
@@ -23,8 +24,10 @@ use fuels_core::types::{
 };
 use fuels_core::types::{input::Input, transaction_builders::ScriptTransactionBuilder};
 use handlebars::Handlebars;
+use hex::FromHexError;
 use reqwest::StatusCode;
 use secrecy::ExposeSecret;
+use serde::Deserialize;
 use serde_json::json;
 use std::sync::{Arc, Mutex};
 use std::{
@@ -67,13 +70,11 @@ pub fn serve_worker() -> MethodRouter {
     let template = concat!(env!("OUT_DIR"), "/worker.js");
 
     async fn handle_error(_err: io::Error) -> impl IntoResponse {
-        dbg!(_err);
         (
             StatusCode::INTERNAL_SERVER_ERROR,
             "Could not serve worker.js",
         )
     }
-    dbg!(template);
     get_service(ServeFile::new(template)).handle_error(handle_error)
 }
 
@@ -335,20 +336,71 @@ pub async fn create_session(
     } else {
         return Err(CreateSessionError {
             status: StatusCode::BAD_REQUEST,
-            error: "invalid address".to_string(),
+            error: "invalid address".to_owned(),
         });
     }?;
 
-    // Access and modify the session map within the handler
-    let mut map = sessions.lock().unwrap();
+    let mut map = match sessions.lock() {
+        Ok(val) => val,
+        Err(_) => {
+            return Err(CreateSessionError {
+                status: StatusCode::INTERNAL_SERVER_ERROR,
+                error: "Could not acquire sessions lock".to_owned(),
+            })
+        }
+    };
     let salt = Salt::random();
 
     map.insert(salt.clone(), address);
 
-    dbg!(map.get(&salt));
-
     Ok(CreateSessionResponse {
         status: "Success".to_owned(),
         salt: hex::encode(salt.as_bytes()),
     })
 }
+
+#[derive(Deserialize)]
+pub struct SessionQuery {
+    salt: String,
+}
+
+pub async fn get_session(
+    query: Query<SessionQuery>,
+    Extension(sessions): Extension<Arc<Mutex<SessionMap>>>,
+) -> Response {
+    let salt: Result<[u8; 32], _> = hex::decode(&query.salt).and_then(|value| {
+        value
+            .try_into()
+            .map_err(|_| FromHexError::InvalidStringLength)
+    });
+
+    match salt {
+        Ok(value) => value,
+        Err(_) => {
+            return (
+                StatusCode::BAD_REQUEST,
+                Json(json!({"error": "Invalid salt"})),
+            )
+                .into_response()
+        }
+    };
+
+    let map = match sessions.lock() {
+        Ok(val) => val,
+        Err(_) => {
+            return (
+                StatusCode::INTERNAL_SERVER_ERROR,
+                Json(json!({"error": "Could not acquire sessions lock"})),
+            )
+                .into_response()
+        }
+    };
+
+    let result = map.get(&Salt::new(salt.unwrap()));
+
+    match result {
+        Some(address) => (StatusCode::OK, Json(json!({"address": address}))),
+        None => (StatusCode::NOT_FOUND, Json(json!({}))),
+    }
+    .into_response()
+}
diff --git a/src/session.rs b/src/session.rs
index 4caadfb..85c641b 100644
--- a/src/session.rs
+++ b/src/session.rs
@@ -6,6 +6,10 @@ use std::collections::HashMap;
 pub struct Salt([u8; 32]);
 
 impl Salt {
+    pub fn new(bytes: [u8; 32]) -> Self {
+        Salt(bytes)
+    }
+
     pub fn random() -> Self {
         let mut rng = rand::thread_rng();
         let mut bytes = [0u8; 32];