From 4e8797d68d5c4acc14e75d341df3dfe0d0dcf671 Mon Sep 17 00:00:00 2001
From: Pierre HUBERT <pierre.git@communiquons.org>
Date: Sat, 20 Apr 2024 13:36:41 +0200
Subject: [PATCH] Improve API

---
 src/lib.rs | 67 +++++++++++++++++++++++++++++++-----------------------
 1 file changed, 39 insertions(+), 28 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index dd5e1dd..eb791b6 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -35,37 +35,40 @@ pub fn generate_ec384_keypair() -> anyhow::Result<(TokenPubKey, TokenPrivKey)> {
         TokenPrivKey::ES384 { r#priv: priv_pem },
     ))
 }
+impl TokenPrivKey {
+    /// Sign a JWT
+    pub fn sign_jwt<C: Serialize>(&self, claims: &C) -> anyhow::Result<String> {
+        match self {
+            TokenPrivKey::ES384 { r#priv } => {
+                let encoding_key = EncodingKey::from_ec_pem(r#priv.as_bytes())?;
 
-/// Sign JWT with a private key
-pub fn sign_jwt<C: Serialize>(key: &TokenPrivKey, claims: &C) -> anyhow::Result<String> {
-    match key {
-        TokenPrivKey::ES384 { r#priv } => {
-            let encoding_key = EncodingKey::from_ec_pem(r#priv.as_bytes())?;
-
-            Ok(jsonwebtoken::encode(
-                &jsonwebtoken::Header::new(Algorithm::ES384),
-                &claims,
-                &encoding_key,
-            )?)
+                Ok(jsonwebtoken::encode(
+                    &jsonwebtoken::Header::new(Algorithm::ES384),
+                    &claims,
+                    &encoding_key,
+                )?)
+            }
         }
     }
 }
 
-/// Validate a given JWT
-pub fn validate_jwt<E: DeserializeOwned>(key: &TokenPubKey, token: &str) -> anyhow::Result<E> {
-    match key {
-        TokenPubKey::ES384 { r#pub } => {
-            let decoding_key = DecodingKey::from_ec_pem(r#pub.as_bytes())?;
+impl TokenPubKey {
+    /// Validate a given JWT
+    pub fn validate_jwt<E: DeserializeOwned>(&self, jwt: &str) -> anyhow::Result<E> {
+        match self {
+            TokenPubKey::ES384 { r#pub } => {
+                let decoding_key = DecodingKey::from_ec_pem(r#pub.as_bytes())?;
 
-            let validation = Validation::new(Algorithm::ES384);
-            Ok(jsonwebtoken::decode::<E>(token, &decoding_key, &validation)?.claims)
+                let validation = Validation::new(Algorithm::ES384);
+                Ok(jsonwebtoken::decode::<E>(jwt, &decoding_key, &validation)?.claims)
+            }
         }
     }
 }
 
 #[cfg(test)]
 mod test {
-    use crate::{generate_ec384_keypair, sign_jwt, validate_jwt};
+    use crate::generate_ec384_keypair;
     use std::time::{SystemTime, UNIX_EPOCH};
 
     use serde::{Deserialize, Serialize};
@@ -96,8 +99,10 @@ mod test {
     fn jwt_encode_sign_verify_valid() {
         let (pub_key, priv_key) = generate_ec384_keypair().unwrap();
         let claims = Claims::default();
-        let jwt = sign_jwt(&priv_key, &claims).expect("Failed to sign JWT!");
-        let claims_out = validate_jwt::<Claims>(&pub_key, &jwt).expect("Failed to validate JWT!");
+        let jwt = priv_key.sign_jwt(&claims).expect("Failed to sign JWT!");
+        let claims_out = pub_key
+            .validate_jwt::<Claims>(&jwt)
+            .expect("Failed to validate JWT!");
 
         assert_eq!(claims, claims_out)
     }
@@ -107,14 +112,17 @@ mod test {
         let (_pub_key, priv_key) = generate_ec384_keypair().unwrap();
         let (pub_key_2, _priv_key_2) = generate_ec384_keypair().unwrap();
         let claims = Claims::default();
-        let jwt = sign_jwt(&priv_key, &claims).expect("Failed to sign JWT!");
-        validate_jwt::<Claims>(&pub_key_2, &jwt).expect_err("JWT should not have validated!");
+        let jwt = priv_key.sign_jwt(&claims).expect("Failed to sign JWT!");
+        pub_key_2
+            .validate_jwt::<Claims>(&jwt)
+            .expect_err("JWT should not have validated!");
     }
 
     #[test]
     fn jwt_verify_random_string() {
         let (pub_key, _priv_key) = generate_ec384_keypair().unwrap();
-        validate_jwt::<Claims>(&pub_key, "random_string")
+        pub_key
+            .validate_jwt::<Claims>("random_string")
             .expect_err("JWT should not have validated!");
     }
 
@@ -125,16 +133,19 @@ mod test {
             exp: time() - 100,
             ..Default::default()
         };
-        let jwt = sign_jwt(&priv_key, &claims).expect("Failed to sign JWT!");
-        validate_jwt::<Claims>(&pub_key, &jwt).expect_err("JWT should not have validated!");
+        let jwt = priv_key.sign_jwt(&claims).expect("Failed to sign JWT!");
+        pub_key
+            .validate_jwt::<Claims>(&jwt)
+            .expect_err("JWT should not have validated!");
     }
 
     #[test]
     fn jwt_invalid_signature() {
         let (pub_key, priv_key) = generate_ec384_keypair().unwrap();
         let claims = Claims::default();
-        let jwt = sign_jwt(&priv_key, &claims).expect("Failed to sign JWT!");
-        validate_jwt::<Claims>(&pub_key, &format!("{jwt}bad"))
+        let jwt = priv_key.sign_jwt(&claims).expect("Failed to sign JWT!");
+        pub_key
+            .validate_jwt::<Claims>(&format!("{jwt}bad"))
             .expect_err("JWT should not have validated!");
     }
 }