Add cookies and authentication with jwto

examples/09-jwt/foo.json

@@ -0,0 +1,2 @@
+{ "username": "dinosaure",
+  "password": "foo" }

examples/09-jwt/main.ml

@@ -0,0 +1,97 @@
+#require "vif" ;;
+#require "jwto" ;;
+
+open Vif ;;
+
+type user =
+  { username : string }
+;;
+
+type cfg =
+  { secret : string }
+;;
+
+let jwt = Vif.Ms.make ~name:"jwt" @@ fun req target server { secret } ->
+  Logs.debug (fun m -> m "Search vif-token cookie");
+  match Cookie.get server req ~name:"vif-token" with
+  | Error err ->
+      Logs.err (fun m -> m "jwt: %a" Cookie.pp_error err);
+      None
+  | Ok token ->
+      Logs.debug (fun m -> m "Token found: %S" token);
+      let ( let* ) = Option.bind in
+      let* token = Result.to_option (Jwto.decode_and_verify secret token) in
+      let* username = List.assoc_opt "username" (Jwto.get_payload token) in
+      Some { username }
+;;
+
+type credential =
+  { username : string
+  ; password : string }
+;;
+
+let credential =
+  let open Json_encoding in
+  let username = req "username" string in
+  let password = req "password" string in
+  let credential = obj2 username password in
+  let prj { username; password } =
+    (username, password) in
+  let inj (username, password) =
+    { username; password } in
+  conv prj inj credential
+;;
+
+let users =
+  [ "dinosaure", "foo" ]
+;;
+
+let login req server { secret }=
+  match Vif.Request.of_json req with
+  | Ok { username; password } ->
+      begin match List.assoc_opt username users with
+      | Some p' when password = p' ->
+          let token = Jwto.encode HS512 secret [ "username", username ] in
+          let token = Result.get_ok token in
+          let* () = Vif.Cookie.set ~name:"vif-token" server req token in
+          let field = "content-type" in
+          let* () = Response.add ~field "text/plain; charset=utf-8" in
+          let* () = Response.with_string req "Authenticated!\n" in
+          Response.respond `OK
+      | _ ->
+          let field = "content-type" in
+          let* () = Response.add ~field "text/plain; charset= utf-8" in
+          let* () = Response.with_string req "Bad credentials\n" in
+          Response.respond `Unauthorized end
+  | Error _ ->
+      let field = "content-type" in
+      let* () = Response.add ~field "text/plain; charset=utf-8" in
+      let* () = Response.with_string req "Invalid JSON\n" in
+      Response.respond (`Code 422)
+;;
+
+let default req target server _cfg =
+  match Request.get jwt req with
+  | None ->
+      let field = "content-type" in
+      let* () = Response.add ~field "text/plain; charset=utf-8" in
+      let* () = Response.with_string req "Unauthorized place" in
+      Response.respond `Unauthorized
+  | Some { username } ->
+      let field = "content-type" in
+      let* () = Response.add ~field "text/plain; charset=utf-8" in
+      let str = Fmt.str "Connected as %S\n" username in
+      let* () = Response.with_string req str in
+      Response.respond `OK
+;;
+
+let routes =
+  let open Vif.U in
+  let open Vif.R in
+  let open Vif.Content_type in
+  [ post (json_encoding credential) (rel / "login" /?? nil) --> login ]
+;;
+
+let () = Miou_unix.run @@ fun () ->
+  let secret = "deadbeef" in
+  Vif.run ~default ~middlewares:Ms.[ jwt ] routes { secret } ;;

lib/vif/vif.ml

@@ -192,7 +192,7 @@ let rec user's_functions daemon =
       Miou.call ~orphans:daemon.orphans @@ fun () ->
       let response = fn daemon.server daemon.user's_value in
       match Vif_response.(run req0 empty) response with
-      | Vif_response.Sent, () -> ()
+      | Vif_response.Sent, () -> Vif_request0.close req0
     in
     ()
   in

lib/vif/vif.mli

@@ -281,11 +281,14 @@ module Cookie : sig
        ?encrypted:bool
     -> name:string
     -> S.t
-    -> ('c, 'a) Request.t
+    -> Request.request
     -> ( string
        , [> `Invalid_encrypted_cookie | `Msg of string | `Not_found ] )
        result
 
+  val pp_error :
+    [ `Invalid_encrypted_cookie | `Msg of string | `Not_found ] Fmt.t
+
   val set :
        ?encrypt:bool
     -> ?cfg:config

lib/vif/vif_cookie.ml

@@ -1,3 +1,7 @@
+let src = Logs.Src.create "vif.cookie"
+
+module Log = (val Logs.src_log src : Logs.LOG)
+
 let prefix req0 =
   let target = Vif_request0.target req0 in
   let secure =
@@ -10,8 +14,7 @@ let prefix req0 =
 
 let is_cookie key = String.lowercase_ascii key = "cookie"
 
-let all_cookies req =
+let all_cookies hdrs =
-  let hdrs = Vif_request0.headers req in
   let cookies = List.filter (fun (k, _) -> is_cookie k) hdrs in
   let cookies = List.map snd cookies in
   let cookies = List.map (String.split_on_char ';') cookies in
@@ -29,10 +32,16 @@ let all_cookies req =
 let guard error fn = if fn () then Ok () else Error error
 let err_cookie = `Invalid_encrypted_cookie
 
+let pp_error ppf = function
+  | `Invalid_encrypted_cookie -> Fmt.string ppf "Invalid encrypted cookie"
+  | `Not_found -> Fmt.string ppf "Cookie not found"
+  | `Msg str -> Fmt.string ppf str
+
 let get ?(encrypted = true) ~name server req0 =
+  let hdrs = Vif_request0.headers req0 in
   let prefix = prefix req0 in
   let name = prefix ^ name in
-  match List.assoc_opt name (all_cookies req0) with
+  match List.assoc_opt name (all_cookies hdrs) with
   | None -> Error `Not_found
   | Some value when encrypted ->
       let ( let* ) = Result.bind in
@@ -52,9 +61,6 @@ let get ?(encrypted = true) ~name server req0 =
       Ok (Option.get value)
   | Some value -> Ok value
 
-let get ?encrypted ~name server req =
-  get ?encrypted ~name server req.Vif_request.request
-
 type config = {
     expires: float option
   ; max_age: float option
@@ -118,11 +124,11 @@ let set ?(encrypt = true) ?(cfg = default_config) ?(path = "/") ~name server
     let key = Vif_s.cookie_key server in
     let nonce = random 12 in
     let adata = "vif.cookie-" ^ name in
-    let value = "\x00" ^ nonce ^ value in
     let value =
       Mirage_crypto.AES.GCM.authenticate_encrypt ~key ~nonce ~adata value
     in
     let alphabet = Base64.uri_safe_alphabet in
+    let value = "\x00" ^ nonce ^ value in
     let value = Base64.encode_exn ~pad:false ~alphabet value in
     let value = set_cookie cfg ~path name value in
     Vif_response.add ~field:"set-cookie" value

lib/vif/vif_m.ml

@@ -12,7 +12,7 @@ type 'cfg m = [] : 'cfg m | ( :: ) : ('cfg, 'a) t * 'cfg m -> 'cfg m
 
 type ('value, 'a, 'c) ctx = {
     server: Vif_s.t
-  ; request: Vif_request0.t
+  ; req0: Vif_request0.t
   ; target: string
   ; user's_value: 'value
 }
@@ -27,7 +27,7 @@ let rec run : type v. v m -> (v, 'a, 'c) ctx -> Hmap.t -> Hmap.t =
   match lst with
   | [] -> env
   | Middleware (fn, key) :: r -> begin
-      match fn ctx.request ctx.target ctx.server ctx.user's_value with
+      match fn ctx.req0 ctx.target ctx.server ctx.user's_value with
       | Some value -> run r ctx (Hmap.add key value env)
       | None -> run r ctx env
       | exception _exn -> run r ctx env

lib/vif/vif_request0.ml

@@ -5,6 +5,7 @@ type t = {
   ; socket: socket
   ; on_localhost: bool
   ; stream: string Stream.stream
+  ; body: [ `V1 of H1.Body.Reader.t | `V2 of H2.Body.Reader.t ]
 }
 
 and reqd = Httpcats.Server.reqd
@@ -35,10 +36,10 @@ let to_stream = function
       |> Stream.Stream.of_bqueue
 
 let of_reqd socket reqd =
-  let request =
+  let request, body =
     match reqd with
-    | `V1 reqd -> V1 (H1.Reqd.request reqd)
+    | `V1 reqd -> (V1 (H1.Reqd.request reqd), `V1 (H1.Reqd.request_body reqd))
-    | `V2 reqd -> V2 (H2.Reqd.request reqd)
+    | `V2 reqd -> (V2 (H2.Reqd.request reqd), `V2 (H2.Reqd.request_body reqd))
   in
   let tls =
     match socket with `Tls tls -> Tls_miou_unix.epoch tls | `Tcp _ -> None
@@ -58,7 +59,7 @@ let of_reqd socket reqd =
         || inet_addr = Unix.inet6_addr_loopback
   in
   let stream = to_stream reqd in
-  { request; tls; reqd; socket; on_localhost; stream }
+  { request; tls; reqd; socket; on_localhost; stream; body }
 
 let headers { request; _ } =
   match request with
@@ -85,3 +86,8 @@ let tls { tls; _ } = tls
 let on_localhost { on_localhost; _ } = on_localhost
 let reqd { reqd; _ } = reqd
 let stream { stream; _ } = stream
+
+let close { body; _ } =
+  match body with
+  | `V1 body -> H1.Body.Reader.close body
+  | `V2 body -> H2.Body.Reader.close body