# See the file "COPYING" in the main distribution directory for copyright.
# Copyright (c) 2023, NCC Group / Fox-IT. See COPYING for details.

module QUIC;

import spicy;
import zeek;

# The interface to the C++ code that handles the decryption of the INITIAL packet payload using well-known keys
public function decrypt_crypto_payload(
  version: uint32,
  data: bytes,
  connection_id: bytes,
  encrypted_offset: uint64,
  payload_offset: uint64,
  from_client: bool
): bytes &cxxname="QUIC_decrypt_crypto_payload";


# Can we decrypt?
function can_decrypt(long_header: LongHeaderPacket, context: Context, crypto: CryptoSinkUnit&): bool {

  if ( ! long_header.is_initial )
    return False;

  if ( crypto == Null )
    return False;

  # Can only decrypt the responder if we've seen the initial destination conn id.
  if ( ! crypto.is_orig && ! context.initial_destination_conn_id )
    return False;

  # Only attempt decryption if we haven't flushed some SSL data yet.
  return ! crypto.finished;
}

function reset_crypto(context: Context&) {
  # Recreate all the crypto state on the next %init of Packet.
  zeek::protocol_handle_close(context.ssl_handle);
  unset context.ssl_handle;
  context.client_crypto = Null;
  context.server_crypto = Null;
  context.client_sink = Null;
  context.server_sink = Null;
  context.initial_destination_conn_id = Null;
}

# This unit is connected with the server and client sinks receiving
# CRYPTO frames and forwards data to the SSL handle in the context.
type CryptoSinkUnit = unit(is_orig: bool, context: Context&) {
  var buffered: bytes;
  var length: uint32 = 0;
  var is_orig: bool = is_orig;
  var finished: bool;

  # The first 4 bytes of crypto data contain the expected tag and a
  # 24bit length from the TLS HandshakeMessage. Extract the length
  # so we can determine when all CRYPTO frames have arrived.
  #
  # https://datatracker.ietf.org/doc/html/rfc8446#section-4
  #
  #  struct {
  #    HandshakeType msg_type;    /* handshake type */
  #    uint24 length;             /* remaining bytes in message */
  #    ...
  #
  : uint8 {
    self.buffered += $$;
  }

  len: uint8[3] {
    self.length = (cast<uint32>($$[0]) << 16) + (cast<uint32>($$[1]) << 8) + cast<uint32>($$[2]) + 4;

    self.buffered += $$[0];
    self.buffered += $$[1];
    self.buffered += $$[2];
  }

  : void &requires=(self.length <= 2**14 + 256) {  # The length MUST NOT exceed 2^14 + 256 bytes (RFC 8446)

    # The client or server hello data is forwarded to the SSL analyzer as a
    # TLSPlaintext record with legacy_record_version set to \x03\x03 (1.3).
    #
    # enum {
    #     invalid(0),
    #     change_cipher_spec(20),
    #     alert(21),
    #     handshake(22),
    #     application_data(23),
    #     (255)
    # } ContentType;
    #
    # struct {
    #     ContentType type;
    #     ProtocolVersion legacy_record_version;
    #     uint16 length;
    #     opaque fragment[TLSPlaintext.length];
    # } TLSPlaintext;
    #
    # https://datatracker.ietf.org/doc/html/rfc8446#section-5.1
    local length_bytes = pack(cast<uint16>(self.length), spicy::ByteOrder::Big);
    zeek::protocol_data_in(is_orig, b"\x16\x03\x03" + length_bytes + self.buffered, context.ssl_handle);
  }

  : bytes &chunked &size=(self.length - 4) {
    zeek::protocol_data_in(is_orig, $$, context.ssl_handle);
  }

  : void {
    self.finished = True;
  }
};

##############
## Context
##############
type Context = struct {
  client_cid_len: uint8;
  server_cid_len: uint8;

  # The DCID used by the client is employed by client and
  # server for packet protection. Packet re-ordering
  # will make life miserable.
  #
  # https://quicwg.org/base-drafts/rfc9001.html#appendix-A
  initial_destination_conn_id: optional<bytes>;

  # Track crypto state.
  client_crypto: CryptoSinkUnit&;
  client_sink: sink&;
  server_crypto: CryptoSinkUnit&;
  server_sink: sink&;

  ssl_handle: zeek::ProtocolHandle &optional;
};

##############
# Definitions
##############
const VersionDraft22: uint32 = 0xff000016;
const VersionDraft23: uint32 = 0xff000017;
const VersionDraft24: uint32 = 0xff000018;
const VersionDraft25: uint32 = 0xff000019;
const VersionDraft26: uint32 = 0xff00001a;
const VersionDraft27: uint32 = 0xff00001b;
const VersionDraft28: uint32 = 0xff00001c;
const VersionDraft29: uint32 = 0xff00001d;
const VersionDraft30: uint32 = 0xff00001e;
const VersionDraft31: uint32 = 0xff00001f;
const VersionDraft32: uint32 = 0xff000020;
const VersionDraft33: uint32 = 0xff000021;
const VersionDraft34: uint32 = 0xff000022;
const VersionFace001: uint32 = 0xfaceb001;
const VersionFace002: uint32 = 0xfaceb002;
const VersionFace00e: uint32 = 0xfaceb00e;
const VersionFace011: uint32 = 0xfaceb011;
const VersionFace012: uint32 = 0xfaceb012;
const VersionFace013: uint32 = 0xfaceb013;
const Version1: uint32 = 0x00000001;
const Version2: uint32 = 0x6b3343cf;

type LongPacketTypeV1 = enum {
  INITIAL         = 0,
  ZERO_RTT        = 1,
  HANDSHAKE       = 2,
  RETRY           = 3,
};

# V2 changed packet types to avoid ossification.
#
# https://www.rfc-editor.org/rfc/rfc9369.html#name-long-header-packet-types
type LongPacketTypeV2 = enum {
  INITIAL         = 1,
  ZERO_RTT        = 2,
  HANDSHAKE       = 3,
  RETRY           = 0,
};

type HeaderForm = enum {
  SHORT           = 0,
  LONG            = 1,
};

type FrameType = enum {
  PADDING                 = 0x00,
  PING                    = 0x01,
  ACK1                    = 0x02,
  ACK2                    = 0x03,
  RESET_STREAM            = 0x04,
  STOP_SENDING            = 0x05,
  CRYPTO                  = 0x06,
  NEW_TOKEN               = 0x07,
  STREAM1                 = 0x08,
  STREAM2                 = 0x09,
  STREAM3                 = 0x0a,
  STREAM4                 = 0x0b,
  STREAM5                 = 0x0c,
  STREAM6                 = 0x0d,
  STREAM7                 = 0x0e,
  STREAM8                 = 0x0f,
  MAX_DATA                = 0x10,
  MAX_STREAM_DATA         = 0x11,
  MAX_STREAMS1            = 0x12,
  MAX_STREAMS2            = 0x13,
  DATA_BLOCKED            = 0x14,
  STREAM_DATA_BLOCKED     = 0x15,
  STREAMS_BLOCKED1        = 0x16,
  STREAMS_BLOCKED2        = 0x17,
  NEW_CONNECTION_ID       = 0x18,
  RETIRE_CONNECTION_ID    = 0x19,
  PATH_CHALLENGE          = 0x1a,
  PATH_RESPONSE           = 0x1b,
  CONNECTION_CLOSE1       = 0x1c,
  CONNECTION_CLOSE2       = 0x1d,
  HANDSHAKE_DONE          = 0x1e,
};

##############
# Helper units
##############

type VariableLengthInteger = unit {
  var bytes_to_parse: uint64;
  var result_: uint64;

  # Value of the two most significant bits indicates number of bytes
  # to parse for the variable length integer.
  #
  # https://datatracker.ietf.org/doc/rfc9000/
  # Section 16 and Appendix A
  : uint8 {
    self.bytes_to_parse = 2**((0xC0 & $$) >> 6);
    self.result_ = $$ & 0x3F;
  }

  : uint8[self.bytes_to_parse - 1] if (self.bytes_to_parse > 1) foreach {
    self.result_ = (self.result_ << 8) | $$;
  }
};

##############
# Long packets
# Generic units
##############
public type LongHeaderPacketV1 = unit(inout outer: LongHeaderPacket) {
  switch ( LongPacketTypeV1(outer.first_byte.packet_type) ) {
    LongPacketTypeV1::INITIAL -> initial_hdr : InitialPacket(outer) {
      outer.is_initial = True;
      outer.encrypted_offset = outer.offset() +
                               self.initial_hdr.length.bytes_to_parse +
                               self.initial_hdr.token_length.bytes_to_parse +
                               self.initial_hdr.token_length.result_;
      outer.payload_length = self.initial_hdr.length.result_;
    }

    LongPacketTypeV1::ZERO_RTT -> zerortt_hdr : ZeroRTTPacket(outer);
    LongPacketTypeV1::HANDSHAKE -> handshake_hdr : HandshakePacket(outer);
    LongPacketTypeV1::RETRY -> retry_hdr : RetryPacket(outer) {
        outer.is_retry = True;
    }
  };
};

public type LongHeaderPacketV2 = unit(inout outer: LongHeaderPacket) {
  switch ( LongPacketTypeV2(outer.first_byte.packet_type) ) {
    LongPacketTypeV2::INITIAL -> initial_hdr : InitialPacket(outer) {
      outer.is_initial = True;
      outer.encrypted_offset = outer.offset() +
                                self.initial_hdr.length.bytes_to_parse +
                                self.initial_hdr.token_length.bytes_to_parse +
                                self.initial_hdr.token_length.result_;
      outer.payload_length = self.initial_hdr.length.result_;
    }

    LongPacketTypeV2::ZERO_RTT -> zerortt_hdr : ZeroRTTPacket(outer);
    LongPacketTypeV2::HANDSHAKE -> handshake_hdr : HandshakePacket(outer);
    LongPacketTypeV2::RETRY -> retry_hdr : RetryPacket(outer) {
        outer.is_retry = True;
    }
  };
};

# Just eat the data for event raising.
public type UnhandledVersion = unit(header: LongHeaderPacket) {
  var header: LongHeaderPacket = header;
  payload: skip bytes &eod;
};

public type LongHeaderPacket = unit {
  var encrypted_offset: uint64;
  var payload_length: uint64;
  var client_conn_id_length: uint8;
  var server_conn_id_length: uint8;
  var is_initial: bool;
  var is_retry: bool;

  first_byte: bitfield(8) {
      header_form: 7 &convert=cast<HeaderForm>(cast<uint8>($$));
      fixed_bit: 6;
      packet_type: 4..5;
      type_specific_bits: 0..3 &convert=cast<uint8>($$);
  };

  version: uint32;
  dest_conn_id_len: uint8 { self.server_conn_id_length = $$; }
  dest_conn_id: bytes &size=self.server_conn_id_length;
  src_conn_id_len: uint8 { self.client_conn_id_length = $$; }
  src_conn_id: bytes &size=self.client_conn_id_length;

  switch ( self.version ) {
    VersionDraft22,
    VersionDraft23,
    VersionDraft24,
    VersionDraft25,
    VersionDraft26,
    VersionDraft27,
    VersionDraft28,
    VersionDraft29,
    VersionDraft30,
    VersionDraft31,
    VersionDraft32,
    VersionDraft33,
    VersionDraft34,
    VersionFace001,
    VersionFace002,
    VersionFace00e,
    VersionFace011,
    VersionFace012,
    VersionFace013,
    Version1 -> v1: LongHeaderPacketV1(self);
    Version2 -> v2: LongHeaderPacketV2(self);
    * -> unknown: UnhandledVersion(self) {
      throw "unhandled QUIC version 0x%x" % self.version;
    }
  };
};

# A QUIC Frame.
public type Frame = unit(header: LongHeaderPacket, from_client: bool, crypto: CryptoSinkUnit, crypto_sink: sink&) {
  frame_type : uint8 &convert=cast<FrameType>($$);

  # TODO: add other FrameTypes as well
  switch ( self.frame_type ) {
    FrameType::ACK1 -> a: ACKPayload(FrameType::ACK1);
    FrameType::ACK2 -> b: ACKPayload(FrameType::ACK2);
    FrameType::CRYPTO -> c: CRYPTOPayload(from_client) {
      # Have the sink re-assemble potentially out-of-order cryptodata
      crypto_sink.write(self.c.cryptodata, self.c.offset.result_);

      # If the crypto unit has determined a valid length, ensure we
      # don't attempt to write more bytes into the sink. If it doesn't,
      # use 2000 bytes as an arbitrary limit required to observe the
      # length of the contained Client Hello or Server Hello.
      if ( crypto.length > 0 ) {
        if ( |crypto_sink| > crypto.length )
          throw "too much crypto data received %s > %s" % ( |crypto_sink|, crypto.length);
      } else {
        if ( |crypto_sink| > 2000 )
          throw "too much crypto data without length received %s" % |crypto_sink|;
      }
    }
    FrameType::CONNECTION_CLOSE1 -> : ConnectionClosePayload(header);
    FrameType::PADDING -> : skip /\x00*/;  # eat the padding
    FrameType::PING -> : void;
    * -> : void {
      throw "unhandled frame type %s in %s" % (self.frame_type, header.first_byte.packet_type);
    }
  };
};

type CRYPTOPayload = unit(from_client: bool) {
  offset: VariableLengthInteger;
  length: VariableLengthInteger;
  cryptodata: bytes &size=self.length.result_;
};

# https://datatracker.ietf.org/doc/html/rfc9000#ack-ranges
type ACKRange = unit {
  gap: VariableLengthInteger;
  ack_range_length: VariableLengthInteger;
};

type ACKECNCounts = unit {
  ect0: VariableLengthInteger;
  ect1: VariableLengthInteger;
  ecn_ce: VariableLengthInteger;
};

# https://datatracker.ietf.org/doc/html/rfc9000#name-ack-frames
type ACKPayload = unit(frame_type: FrameType) {
  latest_ack: VariableLengthInteger;
  ack_delay: VariableLengthInteger;
  ack_range_count: VariableLengthInteger;
  first_ack_range: VariableLengthInteger;
  ack_ranges: ACKRange[self.ack_range_count.result_];
  ecn_counts: ACKECNCounts if(frame_type == FrameType::ACK2);
};

type ConnectionClosePayload = unit(header: LongHeaderPacket) {
  var header: LongHeaderPacket = header;
  error_code: VariableLengthInteger;
  switch {
    -> unknown_frame_type: b"\x00";
    -> frame_type: VariableLengthInteger;
  };
  reason_phrase_length: VariableLengthInteger;
  reason_phrase: bytes &size=self.reason_phrase_length.result_;
};


##############
# Long packets
# Specific long packet type units
##############

# Remainder of an Initial packet
type InitialPacket = unit(header: LongHeaderPacket) {
  var header: LongHeaderPacket = header;
  token_length: VariableLengthInteger;
  token: bytes &size=self.token_length.result_;

  # 5.4.2. Header Protection Sample
  #
  #   That is, in sampling packet ciphertext for header
  #   protection, the Packet Number field is assumed to
  #   be 4 bytes long (its maximum possible encoded length).
  #
  # Enforce 4 bytes Packet Number length + 16 bytes sample
  # ciphertext available.
  length: VariableLengthInteger &requires=self.length.result_ >= 20;

  # Consume the remainder of payload. This
  # includes the packet number field, but we
  # do not know its length yet. We need the
  # payload for sampling, however.
  payload: skip bytes &size=self.length.result_;
};

type ZeroRTTPacket = unit(header: LongHeaderPacket) {
  var header: LongHeaderPacket = header;
  length: VariableLengthInteger;
  payload: skip bytes &size=self.length.result_;
};

type HandshakePacket = unit(header: LongHeaderPacket) {
  var header: LongHeaderPacket = header;
  length: VariableLengthInteger;
  payload: skip bytes &size=self.length.result_;
};


type RetryPacket = unit(header: LongHeaderPacket) {
  var header: LongHeaderPacket = header;
  var retry_token: bytes;
  var integrity_tag: bytes;

  # A retry packet ends with a 128bit / 16 byte integrity
  # tag, but otherwise we do not know anything about the
  # size of the retry_token. Slurp the whole datagram and
  # post split it into the distinct parts.
  data: bytes &eod {
    self.retry_token = self.data.sub(0, |self.data| - 16);
    self.integrity_tag = self.data.sub(|self.data| - 16, |self.data|);
  }
};

##############
# Short packets
##############

# TODO: implement
public type ShortHeader = unit(dest_conn_id_length: uint8) {
  first_byte:  bitfield(8) {
      header_form: 7 &convert=cast<HeaderForm>(cast<uint8>($$));
      fixed_bit: 6;
      spin_bit: 5;
      todo: 0..4;
    };
    dest_conn_id: bytes &size=dest_conn_id_length;
};

# TODO: investigate whether we can parse something useful out of this
public type ShortPacketPayload = unit {
  payload: skip bytes &eod;
};

##############
# QUIC packet parsing
#
# A UDP datagram contains one or more QUIC packets.
##############
type Packet = unit(from_client: bool, context: Context&) {
  var decrypted_data: bytes;
  var packet_size: uint64 = 0;
  var start: iterator<stream>;

  var crypto: CryptoSinkUnit&;
  var crypto_sink: sink&;

  # Attach an SSL analyzer to this connection once.
  on %init {
    if ( ! context?.ssl_handle ) {
      context.ssl_handle = zeek::protocol_handle_get_or_create("SSL");
    }

    self.start = self.input();

    # Initialize crypto state in context for both sides if not already done.
    if ( context.client_crypto == Null ) {
      assert ! context.server_crypto;
      context.client_crypto = new CryptoSinkUnit(True, context);
      context.client_sink = new sink;
      context.client_sink.connect(context.client_crypto);

      context.server_crypto = new CryptoSinkUnit(False, context);
      context.server_sink = new sink;
      context.server_sink.connect(context.server_crypto);
    }

    if ( from_client ) {
      self.crypto = context.client_crypto;
      self.crypto_sink = context.client_sink;
    } else {
      self.crypto = context.server_crypto;
      self.crypto_sink = context.server_sink;
    }
  }

  # Peek into the first byte and determine the header type.
  first_byte: bitfield(8) {
    header_form: 7 &convert=HeaderForm($$);
  };

  # TODO: Consider bitfield based look-ahead-parsing in the switch below
  #       to avoid this rewinding here. It's a hack.
  : void {
    self.set_input(self.start);  # rewind
  }

  # Depending on the header, parse it and update the src/dest ConnectionID's
  switch ( self.first_byte.header_form ) {
    HeaderForm::SHORT -> short_header: ShortHeader(context.client_cid_len);
    HeaderForm::LONG -> long_header: LongHeaderPacket {
        # For now, only allow a change of src/dest ConnectionID's for INITIAL packets.

        # If we see a retry packet from the responder, reset the decryption
        # context such that the next DCID from the client is used for decryption.
        if ( self.long_header.is_retry ) {
          reset_crypto(context);

          self.crypto = Null;
          self.crypto_sink = Null;
        }
     }
  };

  : void {
      if ( self?.long_header && can_decrypt(self.long_header, context, self.crypto ) )
          # If we have parsed an initial packet that we can decrypt the payload,
          # determine the size to store into a buffer.
          self.packet_size = self.offset();
  }

  # Buffer the whole packet if we determined we have a chance to decrypt.
  packet_data: bytes &parse-at=self.start &size=self.packet_size if ( self.packet_size > 0 ) {

    if ( from_client ) {
      context.server_cid_len = self.long_header.dest_conn_id_len;
      context.client_cid_len = self.long_header.src_conn_id_len;

      # This is the first INITIAL packet we attempt to decrypt and it is
      # coming from the client. Use its destination connection ID for
      # decryption purposes.
      if ( ! context.initial_destination_conn_id ) {
        context.initial_destination_conn_id = self.long_header.dest_conn_id;
      }

      # This means that here, we can try to decrypt the initial packet!
      # All data is accessible via the `long_header` unit
      self.decrypted_data = decrypt_crypto_payload(
        self.long_header.version,
        self.packet_data,
        *context.initial_destination_conn_id,
        self.long_header.encrypted_offset,
        self.long_header.payload_length,
        from_client
      );

    } else {
      context.server_cid_len = self.long_header.src_conn_id_len;
      context.client_cid_len = self.long_header.dest_conn_id_len;

      self.decrypted_data = decrypt_crypto_payload(
        self.long_header.version,
        self.packet_data,
        *context.initial_destination_conn_id,
        self.long_header.encrypted_offset,
        self.long_header.payload_length,
        from_client
      );
    }

    # We attempted decryption, but it failed. Just reject the
    # input and assume Zeek will disable the analyzer for this
    # connection.
    if ( |self.decrypted_data| == 0 )
      throw "decryption failed";

    # We were able to decrypt the INITIAL packet. Confirm QUIC!
    spicy::accept_input();
  }

  # If this packet has a SHORT header, consume until &eod, there's nothing
  # we can do with it anyhow.
  : ShortPacketPayload if (self.first_byte.header_form == HeaderForm::SHORT);

  # If this was packet with a long header and decrypted data exists, attempt
  # to parse the plain QUIC frames from it.
  frames: Frame(self.long_header, from_client, self.crypto, self.crypto_sink)[] &parse-from=self.decrypted_data if (self.first_byte.header_form == HeaderForm::LONG && |self.decrypted_data| > 0);
};

##############
# Entrypoints
##############
public type RequestFrame = unit {
  %context = Context;
  : Packet(True, self.context())[];
};

public type ResponseFrame = unit {
  %context = Context;
  : Packet(False, self.context())[];
};
