// Copyright 2026 Leo Cheng
// SPDX-License-Identifier: Apache-2.0

///|
/// A client's handshake state (RFC 8446 Appendix A.1), after the ClientHello
/// has been sent.
///
/// Sending the ClientHello is the client's own action, not a received message,
/// so `START` is not a state here — a client that has not sent one has no
/// handshake to be in.
#derive(Eq, Debug)
pub(all) enum Client {
  WaitServerHello
  WaitEncryptedExtensions
  WaitCertOrRequest
  WaitCert
  WaitCertVerify
  WaitFinished
  Connected
}

///|
pub extend Client with Eq::{equal, not_equal}

///|
pub extend Client with Debug::{to_repr}

///|
/// The state a client is in having just sent its ClientHello.
pub fn Client::new() -> Client {
  WaitServerHello
}

///|
/// The next client state on receiving a message of `kind` (RFC 8446 A.1).
///
/// A message that does not belong in the current state raises
/// `UnexpectedMessage` (§6.2) — which is what the peer is owed, rather than
/// being ignored into an ambiguous state.
///
/// `resumed` is a handshake authenticated by a PSK rather than a certificate: Appendix A.1 marks
/// the Certificate and CertificateVerify of such a handshake as absent, so the server's Finished
/// is what follows EncryptedExtensions — or a CertificateRequest, which a resumed handshake may
/// still carry.
pub fn Client::recv(
  self : Client,
  kind : @msg.Kind,
  resumed~ : Bool = false,
) -> Client raise @alert.Alert {
  match (self, kind) {
    (WaitServerHello, ServerHello) => WaitEncryptedExtensions
    (WaitCertOrRequest | WaitCert, Finished) if resumed => Connected
    (WaitEncryptedExtensions, EncryptedExtensions) => WaitCertOrRequest
    // After EncryptedExtensions the server may ask for a client certificate,
    // or send its own straight away.
    (WaitCertOrRequest, CertificateRequest) => WaitCert
    (WaitCertOrRequest, Certificate) => WaitCertVerify
    (WaitCert, Certificate) => WaitCertVerify
    // RFC 8879 §4: a compressed Certificate stands where the Certificate would.
    (WaitCertOrRequest | WaitCert, CompressedCertificate) => WaitCertVerify
    (WaitCertVerify, CertificateVerify) => WaitFinished
    (WaitFinished, Finished) => Connected
    _ => raise @alert.fatal(UnexpectedMessage)
  }
}

///|
/// A server's handshake state (RFC 8446 Appendix A.2).
#derive(Eq, Debug)
pub(all) enum Server {
  Start
  RecvdClientHello
  WaitCert
  WaitCertVerify
  WaitFinished
  Connected
}

///|
pub extend Server with Eq::{equal, not_equal}

///|
pub extend Server with Debug::{to_repr}

///|
/// The state a server is in awaiting the ClientHello.
pub fn Server::new() -> Server {
  Start
}

///|
/// The next server state on receiving a message of `kind` (RFC 8446 A.2): the
/// ClientHello that opens the handshake, then the client's second flight.
pub fn Server::recv(
  self : Server,
  kind : @msg.Kind,
) -> Server raise @alert.Alert {
  match (self, kind) {
    (Start, ClientHello) => RecvdClientHello
    (WaitCert, Certificate) => WaitCertVerify
    (WaitCertVerify, CertificateVerify) => WaitFinished
    (WaitFinished, Finished) => Connected
    _ => raise @alert.fatal(UnexpectedMessage)
  }
}

///|
/// The server's own move after the ClientHello: it negotiates and sends its
/// whole flight, then waits for the client's Certificate when one was asked
/// for, or straight for the client's Finished.
///
/// Sending the flight in any other state is this endpoint's bug and not the
/// peer's, so it raises `InternalError` rather than an alert that blames the
/// other side.
pub fn Server::sent_flight(
  self : Server,
  client_cert~ : Bool = false,
) -> Server raise @alert.Alert {
  match self {
    RecvdClientHello => if client_cert { WaitCert } else { WaitFinished }
    _ => raise @alert.fatal(InternalError)
  }
}

// ----------------------------------------------------------------- negotiation

///|
/// The named groups this build can run a key exchange for.
///
/// x25519 alone: offering a group the key exchange cannot honour would be a lie
/// the handshake discovers too late. It is a parameter on [`negotiate`] rather
/// than a fixed list, so a build with more groups says so.
pub let groups : Array[Int32] = [@ext.x25519]

///|
/// The signature schemes this build can verify a CertificateVerify under.
pub let schemes : Array[Int32] = [@ext.ecdsa_secp256r1_sha256]

///|
/// The x25519 public key length (RFC 7748 §5): a Montgomery-u coordinate.
let x25519_size : Int32 = 32

///|
/// What a server settled on after reading a ClientHello.
#derive(Eq, Debug)
pub(all) enum Choice {
  /// The client offered a group this build runs and a key share for it: the
  /// group, the client's public key, and the scheme its CertificateVerify will
  /// be signed under.
  Chosen(group~ : Int32, key~ : Bytes, scheme~ : Int32)
  /// The client offered a group this build runs but no key share for it, so it
  /// has to send the ClientHello again with one (RFC 8446 §4.1.4).
  Retry(Int32)
}

///|
pub extend Choice with Eq::{equal, not_equal}

///|
pub extend Choice with Debug::{to_repr}

///|
/// The first of `offered` that also appears in `ours`, honouring the client's
/// preference order as RFC 8446 §4.2.7 allows a server to.
fn first_common(offered : ArrayView[Int32], ours : ArrayView[Int32]) -> Int32? {
  offered.iter().find_first(fn(v) { ours.iter().any(fn(o) { o == v }) })
}

///|
/// Negotiate a decoded ClientHello (RFC 8446 §4.1.1).
///
/// Checks it offers TLS 1.3, picks the first group and signature scheme it
/// lists that this build runs, and takes its key share for the chosen group.
/// Raises the alert §6 names for each way that fails — a missing mandatory
/// extension, a client that does not speak 1.3, nothing in common, or a key
/// share contradicting the rest of the message — and answers `Retry` when the
/// chosen group is one the client offered but sent no share for.
///
/// `groups` and `schemes` are what this endpoint can do; they default to what
/// this build implements. `datagram` says the hello arrived over a datagram, where the version
/// codes are DTLS's own (RFC 9147 §5.3) rather than TLS's.
pub fn negotiate(
  ch : @msg.Hello,
  groups~ : ArrayView[Int32] = groups[:],
  schemes~ : ArrayView[Int32] = schemes[:],
  datagram~ : Bool = false,
) -> Choice raise @alert.Alert {
  // §4.2.1: without supported_versions naming TLS 1.3 the client is not offering it, which is a
  // protocol_version matter before any 1.3 extension can be missing. A legacy_version of SSL 3.0
  // or below is refused whatever the extension says, as OpenSSL and tlsfuzzer have it.
  let versions = match @ext.find(ch.extensions[:], SupportedVersions) {
    Some(e) => @alert.decoded(@ext.read_versions(e.data[:]))
    None => []
  }
  // RFC 9147 §4.2.1: a datagram hello names DTLS 1.3 by its own number, and a TLS 1.3 number in
  // `supported_versions` over a datagram is a version this endpoint does not speak.
  let code13 = if datagram { @ext.version_13_datagram } else { @ext.version_13 }
  if !versions.contains(code13) || ch.version <= 0x0300 {
    raise @alert.fatal(ProtocolVersion)
  }
  // §4.1.2: a TLS 1.3 ClientHello offers the null compression method and nothing else.
  if ch.compression != b"\x00" {
    raise @alert.fatal(IllegalParameter)
  }
  // §9.2: a ClientHello offering 1.3 without any of these cannot be a 1.3 handshake at all.
  let offered = @alert.decoded(
    @ext.read_groups(required(ch, SupportedGroups)[:]),
  )
  // §9.2: `signature_algorithms` is how a certificate is judged, so it is owed by a handshake
  // that authenticates with one. A client that offers a ticket proves itself with the binder
  // instead and may name no scheme at all; what this build verifies with stands in, and is only
  // reached if the ticket is refused — where the server then needs the client's own list and
  // refuses the handshake rather than sign under a scheme that was never offered.
  let psk = @ext.find(ch.extensions[:], PreSharedKey) is Some(_)
  let accepted = match @ext.find(ch.extensions[:], SignatureAlgorithms) {
    Some(e) => @alert.decoded(@ext.read_schemes(e.data[:]))
    None => if psk { [] } else { raise @alert.fatal(MissingExtension) }
  }
  let shares = @alert.decoded(@ext.read_shares(required(ch, KeyShare)[:]))
  // §4.2.7: groups of earlier versions are not to be offered for 1.3 at all, so a ClientHello
  // offering nothing else broke the rule rather than missed ours.
  guard first_common(offered[:], groups) is Some(fallback) else {
    raise @alert.fatal(
      if offered.iter().all(obsolete) {
        IllegalParameter
      } else {
        HandshakeFailure
      },
    )
  }
  let scheme = match first_common(accepted[:], schemes) {
    Some(scheme) => scheme
    None =>
      // A PSK offer with no scheme list has nothing to choose from; the first of ours is a
      // placeholder that no signature is made under.
      if psk && accepted.length() == 0 && schemes.length() > 0 {
        schemes[0]
      } else {
        raise @alert.fatal(HandshakeFailure)
      }
  }
  // §4.2.8: a share for a group the client did not list in supported_groups, or two shares for
  // one group, contradicts its own offer, and a server that notices must say so rather than
  // quietly use one.
  for i, s in shares {
    if !offered.contains(s.0) || shares[:i].iter().any(t => t.0 == s.0) {
      raise @alert.fatal(IllegalParameter)
    }
  }
  // A common group the client already sent a share for saves a round trip, so it wins over an
  // earlier one without; the client's order decides within each. OpenSSL and BoringSSL choose
  // the same way.
  let shared = for g in offered {
    if groups.contains(g) && shares.iter().find_first(s => s.0 == g) is Some(s) {
      break Some(s)
    }
  } nobreak {
    None
  }
  match shared {
    Some((group, key)) => {
      if group == @ext.x25519 && key.length() != x25519_size {
        raise @alert.fatal(IllegalParameter)
      }
      Chosen(group~, key~, scheme~)
    }
    // §4.1.4: a group is usable and only the share is missing — worth one more round trip.
    None => Retry(fallback)
  }
}

///|
/// A group TLS 1.3 retired (§4.2.7): the obsolete_RESERVED ranges, and the brainpool code points
/// of RFC 7027 that RFC 8734 replaced for 1.3.
fn obsolete(group : Int32) -> Bool {
  group is (0x0001..<=0x0016 | 0x001a..<=0x001c | 0xff01..<=0xff02)
}

///|
/// The payload of an extension a 1.3 ClientHello must carry (§9.2).
fn required(ch : @msg.Hello, kind : @ext.Kind) -> Bytes raise @alert.Alert {
  match @ext.find(ch.extensions[:], kind) {
    Some(e) => e.data
    None => raise @alert.fatal(MissingExtension)
  }
}

///|
/// Negotiate a raw ClientHello handshake message: unframe it, decode it, and
/// run [`negotiate`].
///
/// A message that is not a well-formed ClientHello raises `DecodeError` (§6.2)
/// rather than vanishing into a `None`.
pub fn negotiate_message(
  message : BytesView,
  groups~ : ArrayView[Int32] = groups[:],
  schemes~ : ArrayView[Int32] = schemes[:],
  datagram~ : Bool = false,
) -> Choice raise @alert.Alert {
  guard @msg.unframe(message) is Some((kind, body)) else {
    raise @alert.fatal(DecodeError)
  }
  if kind != ClientHello {
    raise @alert.fatal(UnexpectedMessage)
  }
  guard @msg.read_hello(body[:]) is Some(ch) else {
    raise @alert.fatal(DecodeError)
  }
  negotiate(ch, groups~, schemes~, datagram~)
}

///|
/// A ServerHello answering a negotiated ClientHello (RFC 8446 §4.1.3).
///
/// The `supported_versions` extension naming TLS 1.3 is prepended, because
/// §4.1.3 puts the negotiated version there rather than in the message's frozen
/// `legacy_version`. Over a datagram both differ: `datagram` names DTLS 1.3 as 0xfefc and leaves
/// 0xfefd in the frozen field (RFC 9147 §5.4).
pub fn server_hello(
  random : BytesView,
  session_id : BytesView,
  suite : Int32,
  extensions~ : ArrayView[@ext.Ext] = [][:],
  datagram~ : Bool = false,
) -> @msg.Server {
  let exts : Array[@ext.Ext] = [
    {
      kind: SupportedVersions,
      data: @ext.selected_version(
        if datagram {
          @ext.version_13_datagram
        } else {
          @ext.version_13
        },
      ),
    },
  ]
  for e in extensions {
    exts.push(e)
  }
  {
    // RFC 9147 §4.3: over a datagram the frozen field keeps DTLS 1.2's number, and 1.3 is named
    // in supported_versions alone — the same rule the ClientHello follows.
    version: if datagram {
      0xfefd
    } else {
      @wire.legacy_version
    },
    random: random.to_owned(),
    session_id: session_id.to_owned(),
    suite,
    compression: 0,
    extensions: exts,
  }
}

// ------------------------------------------------------------- hello retry

///|
/// The `ServerHello.random` that marks a HelloRetryRequest (RFC 8446 §4.1.3):
/// the SHA-256 of `"HelloRetryRequest"`.
///
/// A HelloRetryRequest travels as a ServerHello — same type, same framing —
/// distinguished only by this sentinel, so a receiver must compare against it
/// before reading the message as a real ServerHello. One that does not will try
/// a key exchange against a `key_share` carrying no key.
pub let retry_random : Bytes = b"\xcf\x21\xad\x74\xe5\x9a\x61\x11\xbe\x1d\x8c\x02\x1e\x65\xb8\x91\xc2\xa2\x11\x16\x7a\xbb\x8c\x5e\x07\x9e\x09\xe2\xc8\xa8\x33\x9c"

///|
/// A HelloRetryRequest's `key_share` payload (RFC 8446 §4.2.8): the selected
/// group alone. There is no key — asking for one is the whole point.
pub fn retry_share(group : Int32) -> Bytes {
  let out = Buffer()
  @wire.u16(out, group)
  out.to_bytes()
}

///|
/// The group such a payload names, or `None` if it is not two octets.
pub fn read_retry_share(view : BytesView) -> Int32? {
  if view.length() != 2 {
    return None
  }
  Some(@wire.read_u16(view, 0))
}

///|
/// A HelloRetryRequest asking the client to come back with a key share for
/// `group` (RFC 8446 §4.1.4): the sentinel random, the client's `session_id`
/// echoed back untouched, the chosen suite, `supported_versions`, and the
/// group.
///
/// It encodes through `@msg.server_hello` like any other ServerHello.
pub fn retry(
  session_id : BytesView,
  suite : Int32,
  group : Int32,
  datagram~ : Bool = false,
) -> @msg.Server {
  retry_message(session_id, suite, Some(group), None, datagram)
}

///|
/// A HelloRetryRequest that asks the client to echo `cookie` (RFC 8446 §4.2.2).
///
/// A datagram server sends one before it keeps any state for a client, so a
/// forged source address buys nothing (RFC 9147 §4.2). It may ask for a key
/// share in the same message — `group` — because a second retry is forbidden,
/// and a hello that needs both would otherwise have nowhere to go.
pub fn retry_cookie(
  session_id : BytesView,
  suite : Int32,
  group : Int32?,
  cookie : BytesView,
  datagram~ : Bool = false,
) -> @msg.Server {
  retry_message(session_id, suite, group, Some(cookie), datagram)
}

///|
/// The message both constructors above write: a sentinel random, a suite, and
/// whichever request the retry carries. A retry with neither would be the
/// ClientHello over again, which §4.1.4 forbids, so it is this code's own bug
/// and not the peer's.
fn retry_message(
  session_id : BytesView,
  suite : Int32,
  group : Int32?,
  cookie : BytesView?,
  datagram : Bool,
) -> @msg.Server {
  let exts : Array[@ext.Ext] = []
  match group {
    Some(group) => exts.push({ kind: KeyShare, data: retry_share(group), })
    None => ()
  }
  match cookie {
    Some(value) => exts.push({ kind: Cookie, data: @ext.cookie(value), })
    None => ()
  }
  server_hello(
    retry_random[:],
    session_id,
    suite,
    extensions=exts[:],
    datagram~,
  )
}

///|
/// The cookie a HelloRetryRequest asks the client to echo, or `None` if it
/// carries none (RFC 8446 §4.2.2).
pub fn read_retry_cookie(sh : @msg.Server) -> Bytes? {
  if !is_retry(sh) {
    return None
  }
  match @ext.find(sh.extensions[:], Cookie) {
    Some(e) => @ext.read_cookie(e.data[:])
    None => None
  }
}

///|
/// Whether a ServerHello is really a HelloRetryRequest.
pub fn is_retry(sh : @msg.Server) -> Bool {
  sh.random == retry_random
}

///|
/// The group a HelloRetryRequest asks for a share of, or `None` if the message
/// is an ordinary ServerHello or carries no `key_share`.
pub fn retry_group(sh : @msg.Server) -> Int32? {
  if !is_retry(sh) {
    return None
  }
  match @ext.find(sh.extensions[:], KeyShare) {
    Some(e) => read_retry_share(e.data[:])
    None => None
  }
}

// ------------------------------------------------------------------ the driver

///|
/// A server-side handshake in progress.
///
/// It drives the state machine over a byte stream: [`Handshake::feed`] appends
/// received octets, splits off every complete message, adds each to the running
/// transcript and advances the state, buffering a trailing partial message for
/// the next feed. The server's own flight folds in through
/// [`Handshake::sent`], so the transcript stays in message order.
///
/// A pure core: no keys and no socket. Whoever has those wraps it.
pub struct Handshake {
  mut state : Server
  transcript : @keys.Transcript
  mut buffered : Bytes
  mut hello : Bytes
}

///|
/// A fresh handshake, awaiting the ClientHello.
pub fn Handshake::new(digest~ : @keys.Digest = @keys.digest) -> Handshake {
  {
    state: Start,
    transcript: @keys.Transcript::new(digest~),
    buffered: b"",
    hello: b"",
  }
}

///|
/// The state the handshake is in.
pub fn Handshake::state(self : Handshake) -> Server {
  self.state
}

///|
/// Whether the handshake has completed.
pub fn Handshake::is_connected(self : Handshake) -> Bool {
  self.state == Connected
}

///|
/// The transcript hash over every message seen so far, in order.
pub fn Handshake::transcript(self : Handshake) -> Bytes {
  self.transcript.hash()
}

///|
/// The raw ClientHello the handshake received, empty until one arrives.
///
/// A server needs it to pull the client's `key_share` and run the ECDHE the
/// handshake secrets come from.
pub fn Handshake::hello(self : Handshake) -> Bytes {
  self.hello
}

///|
/// Negotiate the ClientHello this handshake received.
///
/// Raises `DecodeError` before one has arrived at all, which is what an empty
/// message decodes to.
pub fn Handshake::negotiate(
  self : Handshake,
  groups~ : ArrayView[Int32] = groups[:],
  schemes~ : ArrayView[Int32] = schemes[:],
) -> Choice raise @alert.Alert {
  negotiate_message(self.hello[:], groups~, schemes~)
}

///|
/// Feed received octets: process every complete handshake message now
/// available — add it to the transcript and advance the state — buffering any
/// trailing partial message. Answers the message types processed, in order.
pub fn Handshake::feed(
  self : Handshake,
  octets : BytesView,
) -> Array[@msg.Kind] raise @alert.Alert {
  let joined = Buffer()
  joined.write_bytes(self.buffered)
  joined.write_bytesview(octets)
  let all = joined.to_bytes()
  let processed : Array[@msg.Kind] = []
  let mut at = 0
  for ;; {
    match @msg.unframe(all[at:]) {
      Some((kind, body)) => {
        let used = 4 + body.length()
        let message = all[at:at + used].to_owned()
        if kind == ClientHello {
          self.hello = message
        }
        self.transcript.add(message[:])
        self.state = self.state.recv(kind)
        processed.push(kind)
        at = at + used
      }
      None => break
    }
  }
  self.buffered = all[at:].to_owned()
  processed
}

///|
/// Fold a message the server sends — ServerHello, EncryptedExtensions,
/// Certificate and the rest — into the transcript, keeping message order.
pub fn Handshake::sent(self : Handshake, message : BytesView) -> Unit {
  self.transcript.add(message)
}

///|
/// Advance the state past the server's own flight, waiting for a client
/// certificate when one was asked for.
pub fn Handshake::sent_flight(
  self : Handshake,
  client_cert~ : Bool = false,
) -> Unit raise @alert.Alert {
  self.state = self.state.sent_flight(client_cert~)
}