add verification key to permit checking

proto/src/connection/auth/broker.rs

@@ -7,7 +7,7 @@ use tracing::error;
 use crate::{
     bail,
     connection::protocols::{Connection, Protocol},
-    crypto::Serializable,
+    crypto::{self, Serializable},
     error::{Error, Result},
     fail_verification_with_message,
     message::{AuthenticateResponse, Message},
@@ -35,15 +35,22 @@ where
     SignatureScheme::VerificationKey: Serializable,
     SignatureScheme::SigningKey: Serializable,
 {
+    /// We want to return the user's verification key
+    type Return = SignatureScheme::VerificationKey;
+
     /// The authentication implementation for a broker to a user. We take the following steps:
     /// 1. Receive a permit from the user
     /// 2. Validate and remove the permit from `Redis`
-    /// 3. Send a response
+    /// 3. Validate the signature
+    /// 4. Send a response
     ///
     /// # Errors
     /// - If authentication fails
     /// - If our connection fails
-    async fn authenticate(&mut self, connection: &ProtocolType::Connection) -> Result<()> {
+    async fn authenticate(
+        &mut self,
+        connection: &ProtocolType::Connection,
+    ) -> Result<Self::Return> {
         // Receive the permit
         let auth_message = bail!(
             connection.recv_message().await,
@@ -58,20 +65,24 @@ where
         };
 
         // Check the permit with `Redis`
-        match self
+        let serialized_verification_key = match self
             .redis_client
             .validate_permit(&self.identifier, auth_message.permit)
             .await
         {
-            Ok(false) => {
-                fail_verification_with_message!(connection, "malformed verification key");
+            // The permit did not exist
+            Ok(None) => {
+                fail_verification_with_message!(connection, "invalid or expired permit");
             }
+
+            // We failed to contact `Redis`
             Err(err) => {
                 error!("failed to validate permit with Redis: {err}");
                 fail_verification_with_message!(connection, "internal server error");
             }
 
-            Ok(true) => (),
+            // The permit existed, return the associated verification key
+            Ok(Some(serialized_verification_key)) => serialized_verification_key,
         };
 
         // Form the response message
@@ -83,6 +94,14 @@ where
         // Send the successful response to the user
         let _ = connection.send_message(response_message).await;
 
-        Ok(())
+        // Serialize the verification key
+        let verification_key = bail!(
+            crypto::deserialize(&serialized_verification_key),
+            Crypto,
+            "failed to deserialize verification key"
+        );
+
+        // Return the verification key
+        Ok(verification_key)
     }
 }

proto/src/connection/auth/marshal.rs

@@ -1,6 +1,6 @@
 //! In this crate we deal with the authentication flow as a marshal.
 
-use std::time::{SystemTime, UNIX_EPOCH};
+use std::time::{Duration, SystemTime, UNIX_EPOCH};
 
 use async_trait::async_trait;
 use jf_primitives::signatures::SignatureScheme as JfSignatureScheme;
@@ -35,6 +35,9 @@ where
     SignatureScheme::VerificationKey: Serializable,
     SignatureScheme::SigningKey: Serializable,
 {
+    /// We have no auxiliary data to return
+    type Return = ();
+
     /// The authentication implementation for a marshal to a user. We take the following steps:
     /// 1. Receive a signed message from the user
     /// 2. Validate and remove the message
@@ -44,7 +47,10 @@ where
     /// # Errors
     /// - If authentication fails
     /// - If our connection fails
-    async fn authenticate(&mut self, connection: &ProtocolType::Connection) -> Result<()> {
+    async fn authenticate(
+        &mut self,
+        connection: &ProtocolType::Connection,
+    ) -> Result<Self::Return> {
         // Receive the signed message from the user
         let auth_message = bail!(
             connection.recv_message().await,
@@ -102,9 +108,15 @@ where
             };
 
         // Generate and issue a permit for said broker
+        // TODO: add bounds check for verification key. There's the possibility it could be too big, if
+        // verify does not check that.
         let permit = match self
             .redis_client
-            .issue_permit(&broker_with_least_connections)
+            .issue_permit(
+                &broker_with_least_connections,
+                Duration::from_secs(5),
+                auth_message.verification_key,
+            )
             .await
         {
             Ok(broker) => broker,

proto/src/connection/auth/mod.rs

@@ -42,7 +42,11 @@ pub trait AuthenticationFlow<
     ProtocolType: Protocol,
 >: Send + Sync + 'static + Clone
 {
+    /// Used for if we want to return an authentication `Result`. Perhaps something like the
+    /// verification key.
+    type Return;
+
     /// This is where we request or verify authentication. We pass in a connection, which,
     /// if it is verified, returns a resulting, successful connection.
-    async fn authenticate(&mut self, connection: &ProtocolType::Connection) -> Result<()>;
+    async fn authenticate(&mut self, connection: &ProtocolType::Connection) -> Result<Self::Return>;
 }

proto/src/connection/auth/user.rs

@@ -49,6 +49,9 @@ where
     SignatureScheme::VerificationKey: Serializable,
     SignatureScheme::SigningKey: Serializable,
 {
+    /// We have no auxiliary data to return
+    type Return = ();
+
     /// The authentication steps on `UserToBrokerToMarshal`'s connection:
     /// 1. Sign the timestamp with our private key
     /// 2. Send a signed message to the marshal

proto/src/redis.rs

@@ -47,7 +47,7 @@ impl Client {
                 user_advertise_address: String::new(),
                 broker_advertise_address: String::new(),
             },
-            |identifier| identifier
+            |identifier| identifier,
         );
 
         // Return the thinly wrapped `Self`.
@@ -230,15 +230,22 @@ impl Client {
     ///
     /// # Errors
     /// - If the `Redis` connection fails
-    pub async fn issue_permit(&mut self, broker: &BrokerIdentifier) -> Result<u64> {
+    pub async fn issue_permit(
+        &mut self,
+        for_broker: &BrokerIdentifier,
+        expiry: Duration,
+        verification_key: Vec<u8>,
+    ) -> Result<u64> {
         // Create random permit number
         // TODO: figure out if it makes sense to initialize this somewhere else
         let permit = StdRng::from_entropy().next_u64();
 
         // Issue the permit
         bail!(
-            redis::cmd("SADD")
-                .arg(&[format!("{broker}/permits"), permit.to_string()])
+            redis::cmd("SET")
+                .arg(&[format!("{for_broker}/permits/{permit}")])
+                .arg(verification_key)
+                .arg(&["EX", &expiry.as_secs().to_string()])
                 .query_async(&mut self.underlying_connection)
                 .await,
             Connection,
@@ -250,28 +257,24 @@ impl Client {
     }
 
     /// Validate and remove a permit belonging to a particular broker.
-    /// Returns `true` if validation was successful, and `false` if not.
+    /// Returns `Some(validation_key)` if successful, and `None` if not.
     ///
     /// # Errors
     /// - If the `Redis` connection fails
     pub async fn validate_permit(
         &mut self,
         broker: &BrokerIdentifier,
         permit: u64,
-    ) -> Result<bool> {
+    ) -> Result<Option<Vec<u8>>> {
         // Remove the permit
-        match bail!(
-            redis::cmd("SREM")
-                .arg(&[format!("{broker}/permits"), permit.to_string()])
+        Ok(bail!(
+            redis::cmd("GETDEL")
+                .arg(format!("{broker}/permits/{permit}"))
                 .query_async(&mut self.underlying_connection)
                 .await,
             Connection,
             "failed to connect to Redis"
-        ) {
-            0 => Ok(false),
-            1 => Ok(true),
-            _ => Err(Error::Parse("unexpected Redis response".to_string())),
-        }
+        ))
     }
 }