Allow adding extra headers

ulrikstrid · Feb 16, 2023 · ae8ded1 · ae8ded1
1 parent 4f0e42a
commit ae8ded1
Show file tree

Hide file tree

Showing 3 changed files with 55 additions and 22 deletions.
diff --git a/jose/Header.ml b/jose/Header.ml
@@ -2,17 +2,26 @@ open Utils
 
 type t = {
   alg : Jwa.alg;
-  jku : string option;
   jwk : Jwk.public Jwk.t option;
   kid : string option;
   x5t : string option;
   x5t256 : string option;
   typ : string option;
   cty : string option;
   enc : Jwa.enc option;
+  extra : (string * Yojson.Safe.t) list option;
 }
 
-let make_header ?typ ?alg ?enc (jwk : Jwk.priv Jwk.t) =
+(* TODO: This is probably very slow *)
+let remove_supported (l : (string * Yojson.Safe.t) list) =
+  l |> List.remove_assoc "alg" |> List.remove_assoc "jwk"
+  |> List.remove_assoc "kid" |> List.remove_assoc "x5t"
+  |> List.remove_assoc "x5t#256"
+  |> List.remove_assoc "typ" |> List.remove_assoc "cty"
+  |> List.remove_assoc "enc"
+
+let make_header ?typ ?alg ?enc ?(extra = []) ?(jwk_header = false)
+    (jwk : Jwk.priv Jwk.t) =
   let alg =
     match alg with
     | Some alg -> alg
@@ -24,26 +33,38 @@ let make_header ?typ ?alg ?enc (jwk : Jwk.priv Jwk.t) =
         | Jwk.Es384_priv _ -> `ES384
         | Jwk.Es512_priv _ -> `ES512)
   in
+  let kid =
+    match List.assoc_opt "kid" extra with
+    | Some kid -> Some (Yojson.Safe.Util.to_string kid)
+    | None -> Jwk.get_kid jwk
+  in
+  let extra = remove_supported extra in
   {
     alg;
-    jku = None;
-    jwk = None;
-    kid = Jwk.get_kid jwk;
+    jwk = (if jwk_header then Some (Jwk.pub_of_priv jwk) else None);
+    kid;
     x5t = None;
     x5t256 = None;
     typ;
     cty = None;
     enc;
+    extra = (match extra with [] -> None | extra -> Some extra);
   }
 
 module Json = Yojson.Safe.Util
 
+let get_extra_headers (json : Yojson.Safe.t) =
+  match json with
+  | `Assoc vals -> (
+      let extra = remove_supported vals in
+      match extra with [] -> None | extra -> Some extra)
+  | _ -> None (* TODO: raise here? *)
+
 let of_json json =
   try
     Ok
       {
         alg = json |> Json.member "alg" |> Jwa.alg_of_json;
-        jku = json |> Json.member "jku" |> Json.to_string_option;
         jwk =
           json |> Json.member "jwk"
           |> Json.to_option (fun jwk_json ->
@@ -57,6 +78,7 @@ let of_json json =
         enc =
           json |> Json.member "enc" |> Json.to_string_option
           |> U_Opt.map Jwa.enc_of_string;
+        extra = get_extra_headers json;
       }
   with Json.Type_error (s, _) -> Error (`Msg s)
 
@@ -65,17 +87,18 @@ let to_json t =
     [
       RJson.to_json_string_opt "typ" t.typ;
       Some ("alg", Jwa.alg_to_json t.alg);
-      RJson.to_json_string_opt "jku" t.jku;
-      U_Opt.map Jwk.to_pub_json t.jwk |> U_Opt.map (fun jwk -> ("jwk", jwk));
       RJson.to_json_string_opt "kid" t.kid;
+      U_Opt.map Jwk.to_pub_json t.jwk |> U_Opt.map (fun jwk -> ("jwk", jwk));
       RJson.to_json_string_opt "x5t" t.x5t;
       RJson.to_json_string_opt "x5t#256" t.x5t256;
       RJson.to_json_string_opt "cty" t.cty;
-      t.enc |> U_Opt.map Jwa.enc_to_string
+      t.enc
+      |> U_Opt.map Jwa.enc_to_string
       |> U_Opt.map (fun enc -> ("enc", `String enc));
     ]
   in
-  `Assoc (U_List.filter_map (fun x -> x) values)
+  let extra = Option.value ~default:[] t.extra in
+  `Assoc (U_List.filter_map (fun x -> x) values @ extra)
 
 let of_string header_str =
   U_Base64.url_decode header_str

diff --git a/jose/Jose.mli b/jose/Jose.mli
@@ -253,29 +253,39 @@ end
 module Header : sig
   type t = {
     alg : Jwa.alg;
-    jku : string option;
     jwk : Jwk.public Jwk.t option;
     kid : string option;
     x5t : string option;
     x5t256 : string option;
     typ : string option;
     cty : string option;
     enc : Jwa.enc option;
+    extra : (string * Yojson.Safe.t) list option;
   }
-  (** The [header] has the following properties: - [alg] Jwa - RS256 and none is
-      currently the only supported algs - [jku] JWK Set URL - [jwk] JSON Web Key
-      - [kid] Key ID - We currently always expect this to be there, this can
-      change in the future - [x5t] X.509 Certificate SHA-1 Thumbprint -
-      [x5t#S256] X.509 Certificate SHA-256 Thumbprint - [typ] Type - [cty]
-      Content Type Not implemented: - [x5u] X.509 URL - [x5c] X.509 Certficate
-      Chain - [crit] Critical
+  (** The [header] has the following properties:
+  
+  - [alg] {! Jwa.alg }
+  - [jwk] JSON Web Key
+  - [kid] Key ID - We currently always expect this to be there, this can change in the future
+  - [x5t] X.509 Certificate SHA-1 Thumbprint -
+  - [x5t#S256] X.509 Certificate SHA-256 Thumbprint
+  - [typ] Type
+  - [cty] Content Type Not implemented
+  
+      {{: https://tools.ietf.org/html/rfc7515#section-4.1 } Link to RFC }
 
-      {{: https://tools.ietf.org/html/rfc7515#section-4.1 } Link to RFC } *)
+      {{: https://www.iana.org/assignments/jose/jose.xhtml#web-signature-encryption-header-parameters } Complete list of registered header parameters} *)
 
   val make_header :
-    ?typ:string -> ?alg:Jwa.alg -> ?enc:Jwa.enc -> Jwk.priv Jwk.t -> t
+    ?typ:string ->
+    ?alg:Jwa.alg ->
+    ?enc:Jwa.enc ->
+    ?extra:(string * Yojson.Safe.t) list ->
+    ?jwk_header:bool ->
+    Jwk.priv Jwk.t ->
+    t
   (** [make_header typ alg enc jwk] if [alg] is not provided it will be derived
-      from [jwk]. *)
+      from [jwk]. [jwk_header] decides if the jwk should be put in the header. *)
 
   val of_string : string -> (t, [> `Msg of string ]) result
   val to_string : t -> string

diff --git a/jose/Jws.ml b/jose/Jws.ml
@@ -3,7 +3,7 @@ open Utils
 type signature = string
 
 type t = {
-  header : Header.t;
+  header : Header.t; (* TODO: This is always treated as protected headers*)
   raw_header : string;
   payload : string;
   signature : signature;