jwt_io: parse header for key id

The key id should not be in the body, but in the header.
This is also the default case for auth0.

Author: Robin de Rooij

prolog/jwt_io.pl

@@ -26,7 +26,7 @@
   OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */
 
-:- module(jwt_io, [jwt_encode/3, jwt_decode/3]).
+:- module(jwt_io, [jwt_encode/3, jwt_decode/3, jwt_decode_head/2]).
 /** <module> Json Web Tokens implementation
 
 Generates and verifies Json Web Tokens.
@@ -44,7 +44,7 @@
     * =kid=: key id for identifying the key to use
     * =type=: type of the key, one of HMAC, RSA or ECDSA.
     * =algorithm=: algorithm to use, one of HS256, HS384, HS512, RS256, RS384, RS512, ES256, ES384 or ES512.
-    * =key=: private key to use - string for HMAC, private key file for RSA and private PEM file for ECDSA.
+    * =key=: private key to use - string for HMAC, private key file for RSA and private PEM file for ECDSA. Optional for decoding, mandatory for encoding.
     * =public_key=: public key to use - irrelevant for HMAC, public key file for RSA and public PEM file for ECDSA.
 
 RSA keys can be generated by:
@@ -102,9 +102,9 @@
   claims_with_aud(Claims, ClaimsWithAud),
   claims_with_iat(ClaimsWithAud, ClaimsWithIat),
   claims_with_iss(ClaimsWithIat, Key, ClaimsWithIss),
-  ClaimsNew = ClaimsWithIss.put(_{kid: Kid, jti: Jti}),
+  ClaimsNew = ClaimsWithIss.put(_{jti: Jti}),
   atom_json_dict(Data,ClaimsNew,[as(atom)]),
-  jwt_encode_from_string(Data, Token, Key.key, Key.algorithm).
+  jwt_encode_from_string(Data, Token, Key.key, Key.algorithm, Kid).
 
 %! jwt_decode(+Data: atom, -Payload:dict, +Options:options) is semidet
 %
@@ -126,18 +126,22 @@
 % @arg Options options
 %
 jwt_decode(Data, Payload, Options) :-
-  jwt_decode_from_string(Data, PayloadFirst, _, ''),
-  atom_json_dict(PayloadFirst, PayloadDict, [as(string)]),
-  atom_string(Kid, PayloadDict.kid),
+  jwt_decode_head(Data, PayloadFirst),
+  atom_json_dict(PayloadFirst, PayloadHeader, [as(string)]),
+  atom_string(Kid, PayloadHeader.kid),
   get_key_from_settings(Kid, KeyDict),
   jwt_decode_from_string(Data, Payload, KeyDict),
+  atom_json_dict(Payload, PayloadDict, [as(string)]),
   jwt_jti_valid(PayloadDict),
   jwt_exp_valid(PayloadDict),
   jwt_nbf_valid(PayloadDict),
   jwt_iss_valid(PayloadDict, Options),
   jwt_aud_valid(PayloadDict, Options),
   jwt_iat_valid(PayloadDict).
 
+jwt_decode_head(Data, Payload) :-
+  jwt_parse_head(Data, Payload).
+
 get_jti(Jti) :-
   setting(jti_generator, Generator),
   call(Generator, Jti).
@@ -203,6 +207,7 @@
 
 jwt_decode_from_string(Data, Payload, Key) :-
   member(Key.type, ['RSA', 'ECDSA']),
+  !,
   jwt_decode_from_string(Data, Payload, Key.algorithm, Key.public_key),
   !.
 
@@ -268,12 +273,21 @@
   !.
 read_key_file(Key,KeyRead) :-
   member(Key.type, ['RSA', 'ECDSA']),
-  read_file_to_string(Key.key, KeyStr, []),
+  get_dict('key', Key, PrivateKeyFile),
+  !,
+  read_file_to_string(PrivateKeyFile, KeyStr, []),
   read_file_to_string(Key.public_key, PublicKeyStr, []),
   atom_string(KeyAtom, KeyStr),
   atom_string(PublicKeyAtom, PublicKeyStr),
   KeyRead = Key.put(_{key: KeyAtom, public_key: PublicKeyAtom}),
   !.
+read_key_file(Key, KeyRead) :-
+  member(Key.type, ['RSA', 'ECDSA']),
+  read_file_to_string(Key.public_key, PublicKeyStr, []),
+  atom_string(PublicKeyAtom, PublicKeyStr),
+  KeyRead = Key.put(_{public_key: PublicKeyAtom}),
+  !.
+
 
 get_key_file(File, Key) :-
   read_file_to_string(File, KeyStr, []),

src/jwt_io.c

@@ -19,17 +19,19 @@ get_pl_arg_str(char *predicate_name, char *term_name, term_t term, char **out)
 }
 
 static		foreign_t
-pl_jwt_encode(term_t in, term_t out, term_t key_term, term_t algorithm)
+pl_jwt_encode(term_t in, term_t out, term_t key_term, term_t algorithm, term_t key_id)
 {
 	jwt_t          *jwt;
-	char           *result, *grants, *alg_str, *key;
+	char           *result, *grants, *alg_str, *key, *kid;
 	int		rval;
 
-	get_pl_arg_str("jwt_encode_from_string/4", "in", in, &grants);
-	get_pl_arg_str("jwt_encode_from_string/4", "key", key_term, &key);
-	get_pl_arg_str("jwt_encode_from_string/4", "algorithm", algorithm, &alg_str);
+	get_pl_arg_str("jwt_encode_from_string/5", "in", in, &grants);
+	get_pl_arg_str("jwt_encode_from_string/5", "key", key_term, &key);
+        get_pl_arg_str("jwt_encode_from_string/5", "kid", key_id, &kid);
+	get_pl_arg_str("jwt_encode_from_string/5", "algorithm", algorithm, &alg_str);
 	jwt_new(&jwt);
 	jwt_add_grants_json(jwt, grants);
+        jwt_add_header(jwt, "kid", kid);
 	jwt_set_alg(jwt, jwt_str_alg(alg_str), (const unsigned char *)key, strlen(key));
 	result = jwt_encode_str(jwt);
 	rval = PL_unify_atom_chars(out, result);
@@ -38,6 +40,26 @@ pl_jwt_encode(term_t in, term_t out, term_t key_term, term_t algorithm)
 	return rval;
 }
 
+static		foreign_t
+pl_jwt_parse_head(term_t in, term_t head_term) {
+    char        *input, *head_payload;
+    jwt_t       *jwt;
+    int		jwt_result;
+    get_pl_arg_str("jwt_encode_from_string/4", "in", in, &input);
+
+    jwt_result = jwt_decode(&jwt, input, NULL, 0);
+    if (!jwt_result) {
+        head_payload = jwt_get_headers_json(jwt, NULL);
+        (void) PL_unify_atom_chars(head_term, head_payload);
+        free(head_payload);
+        jwt_free(jwt);
+        PL_succeed;
+    }
+    else {
+        PL_fail;
+    }
+}
+
 static		foreign_t
 pl_jwt_decode(term_t in, term_t out_payload, term_t out_algorithm, term_t in_key)
 {
@@ -62,14 +84,14 @@ pl_jwt_decode(term_t in, term_t out_payload, term_t out_algorithm, term_t in_key
 		jwt_free(jwt);
 		PL_succeed;
 	} else {
-    PL_warning("cannot decode jwt: %s", strerror(jwt_result));
     PL_fail;
   }
 }
 
 install_t
 install(void)
 {
-	PL_register_foreign("jwt_encode_from_string", 4, pl_jwt_encode, 0);
+	PL_register_foreign("jwt_parse_head", 2, pl_jwt_parse_head, 0);
+	PL_register_foreign("jwt_encode_from_string", 5, pl_jwt_encode, 0);
 	PL_register_foreign("jwt_decode_from_string", 4, pl_jwt_decode, 0);
 }