Source file noise.ml

module Error = Error

type keypair = Dh.keypair = { priv : string; pub : string }
type role = Handshake_state.role = Initiator | Responder

let generate_keypair () = Dh.x25519.Dh.generate ()
let keypair_of_secret priv = Dh.x25519.Dh.of_secret priv

module Handshake = struct
  type t = { mutable hs : Handshake_state.t; mutable consumed : bool }

  let create ~protocol_name ~role ?prologue ?s ?e ?rs ?re ?psks () =
    match
      Handshake_state.create ~protocol_name ~role ?prologue ?s ?e ?rs ?re ?psks
        ()
    with
    | Ok hs -> Ok { hs; consumed = false }
    | Error e -> Error e

  let write_message t payload =
    match Handshake_state.write_message t.hs payload with
    | Ok (message, hs) ->
        t.hs <- hs;
        Ok message
    | Error e -> Error e

  let read_message t message =
    match Handshake_state.read_message t.hs message with
    | Ok (payload, hs) ->
        t.hs <- hs;
        Ok payload
    | Error e -> Error e

  let is_finished t = Handshake_state.is_finished t.hs
  let handshake_hash t = Handshake_state.handshake_hash t.hs
  let remote_static t = Handshake_state.remote_static t.hs

  let split_raw t =
    if t.consumed then Error (Error.Invalid_state "transport already extracted")
    else
      let ( let* ) = Result.bind in
      let* cs1, cs2 = Handshake_state.split t.hs in
      match (cs1.Cipher_state.k, cs2.Cipher_state.k) with
      | Some k1, Some k2 ->
          t.consumed <- true;
          Ok (k1, k2)
      | _ -> Error (Error.Invalid_state "split produced empty key")
end

module Transport = struct
  type t = { mutable tr : Transport.t }

  let of_handshake (h : Handshake.t) =
    if h.Handshake.consumed then
      Error (Error.Invalid_state "transport already extracted")
    else
      let ( let* ) = Result.bind in
      let* cs1, cs2 = Handshake_state.split h.Handshake.hs in
      let role = Handshake_state.role_of h.Handshake.hs in
      let cipher = Handshake_state.cipher_of h.Handshake.hs in
      let one_way = Handshake_state.one_way_of h.Handshake.hs in
      h.Handshake.consumed <- true;
      Ok { tr = Transport.create ~role ~cipher ~one_way (cs1, cs2) }

  let send t plaintext =
    match Transport.send t.tr plaintext with
    | Ok (ciphertext, tr) ->
        t.tr <- tr;
        Ok ciphertext
    | Error e -> Error e

  let recv t ciphertext =
    match Transport.recv t.tr ciphertext with
    | Ok (plaintext, tr) ->
        t.tr <- tr;
        Ok plaintext
    | Error e -> Error e

  let rekey_send t =
    match Transport.rekey_send t.tr with
    | Ok tr ->
        t.tr <- tr;
        Ok ()
    | Error e -> Error e

  let rekey_recv t =
    match Transport.rekey_recv t.tr with
    | Ok tr ->
        t.tr <- tr;
        Ok ()
    | Error e -> Error e
end