From 74b7b344c6d2ff97a752bbdc8247a82529801832 Mon Sep 17 00:00:00 2001
From: Huang-Ming Huang <huangh@objectcomputing.com>
Date: Sat, 16 Feb 2019 15:36:12 -0600
Subject: [PATCH] Add api to supply secret based on decoded payload

---
 include/jwt/impl/jwt.ipp   | 14 ++++++++++++--
 include/jwt/jwt.hpp        |  5 ++++-
 include/jwt/parameters.hpp | 16 ++++++++++++++++
 tests/test_jwt_es.cc       |  9 +++++++++
 4 files changed, 41 insertions(+), 3 deletions(-)

diff --git a/include/jwt/impl/jwt.ipp b/include/jwt/impl/jwt.ipp
index 5e17936..197a614 100644
--- a/include/jwt/impl/jwt.ipp
+++ b/include/jwt/impl/jwt.ipp
@@ -544,6 +544,14 @@ void jwt_object::set_decode_params(DecodeParams& dparams, params::detail::secret
   jwt_object::set_decode_params(dparams, std::forward<Rest>(args)...);
 }
 
+template <typename DecodeParams, typename T, typename... Rest>
+void jwt_object::set_decode_params(DecodeParams& dparams, params::detail::secret_function_param<T>&& s, Rest&&... args)
+{
+  dparams.secret = s.get(*dparams.payload_ptr);
+  dparams.has_secret = true;
+  jwt_object::set_decode_params(dparams, std::forward<Rest>(args)...);
+}
+
 template <typename DecodeParams, typename... Rest>
 void jwt_object::set_decode_params(DecodeParams& dparams, params::detail::leeway_param l, Rest&&... args)
 {
@@ -650,10 +658,11 @@ jwt_object decode(const jwt::string_view enc_str,
 
     //Validate JTI
     bool validate_jti = false;
+    const jwt_payload* payload_ptr = 0;
   };
 
   decode_params dparams{};
-  jwt_object::set_decode_params(dparams, std::forward<Args>(args)...);
+  
 
   //Signature must have atleast 2 dots
   auto dot_cnt = std::count_if(std::begin(enc_str), std::end(enc_str),
@@ -695,7 +704,8 @@ jwt_object decode(const jwt::string_view enc_str,
     return obj;
   }
   obj.payload(std::move(payload));
-
+  dparams.payload_ptr = & obj.payload();
+  jwt_object::set_decode_params(dparams, std::forward<Args>(args)...);
   if (dparams.verify) {
     try {
       ec = obj.verify(dparams, algos);
diff --git a/include/jwt/jwt.hpp b/include/jwt/jwt.hpp
index c96edcb..36dac5f 100644
--- a/include/jwt/jwt.hpp
+++ b/include/jwt/jwt.hpp
@@ -82,7 +82,7 @@ inline jwt::string_view type_to_str(SCOPED_ENUM type typ)
     case type::JWT: return "JWT";
     default:        assert (0 && "Unknown type");
   };
-
+  __builtin_unreachable();
   assert (0 && "Code not reached");
 }
 
@@ -1121,6 +1121,9 @@ public: //TODO: Not good
   template <typename DecodeParams, typename... Rest>
   static void set_decode_params(DecodeParams& dparams, params::detail::secret_param s, Rest&&... args);
 
+  template <typename DecodeParams, typename T, typename... Rest>
+  static void set_decode_params(DecodeParams& dparams, params::detail::secret_function_param<T>&& s, Rest&&... args);
+
   template <typename DecodeParams, typename... Rest>
   static void set_decode_params(DecodeParams& dparams, params::detail::leeway_param l, Rest&&... args);
 
diff --git a/include/jwt/parameters.hpp b/include/jwt/parameters.hpp
index b444e6e..9b94171 100644
--- a/include/jwt/parameters.hpp
+++ b/include/jwt/parameters.hpp
@@ -84,6 +84,15 @@ struct secret_param
   string_view secret_;
 };
 
+template <typename T>
+struct secret_function_param
+{
+  T get() const { return fun_; }
+  template <typename U>
+  std::string get(U&& u) const { return fun_(u);}
+  T fun_;
+};
+
 /**
  * Parameter for providing the algorithm to use.
  * The parameter can accept either the string representation
@@ -297,6 +306,13 @@ inline detail::secret_param secret(const string_view sv)
   return { sv };
 }
 
+template <typename T>
+inline std::enable_if_t<!std::is_convertible<T, string_view>::value, detail::secret_function_param<T>>  
+secret(T&& fun)
+{
+  return detail::secret_function_param<T>{ fun };
+}
+
 /**
  */
 inline detail::algorithm_param algorithm(const string_view sv)
diff --git a/tests/test_jwt_es.cc b/tests/test_jwt_es.cc
index 9064d2d..085cfd1 100644
--- a/tests/test_jwt_es.cc
+++ b/tests/test_jwt_es.cc
@@ -136,6 +136,15 @@ TEST (ESAlgo, ES384EncodingDecodingValidTest)
   EXPECT_EQ (dec_obj.header().algo(), jwt::algorithm::ES384);
   EXPECT_TRUE (dec_obj.has_claim("exp"));
   EXPECT_TRUE (obj.payload().has_claim_with_value("exp", 4682665886));
+
+  std::map<std::string, std::string> keystore{{"arun.muralidharan", key}};
+
+  auto l = [&keystore](const jwt::jwt_payload& payload){ 
+    auto iss = payload.get_claim_value<std::string>("iss");
+    return keystore[iss];
+  };
+  auto dec_obj2 = jwt::decode(enc_str, algorithms({"es384"}), verify(true), secret(l));
+  EXPECT_EQ (dec_obj2.header().algo(), jwt::algorithm::ES384);
 }
 
 int main(int argc, char* argv[]) {