# Copyright (c) 2024 by the Zeek Project. See LICENSE for details.

module WebSocket;

import spicy;
import zeek;

const OPCODE_CONTINUATION = 0x00;
const OPCODE_TEXT         = 0x01;
const OPCODE_BINARY       = 0x02;
const OPCODE_CLOSE        = 0x08;
const OPCODE_PING         = 0x09;
const OPCODE_PONG         = 0x0a;

public function fast_unmask(
  masking_key_idx: uint64,
  masking_key: vector<uint8>,
  chunk: bytes
): bytes &cxxname="hlt_websocket::WebSocket::fast_unmask";

type Frame = unit(m: Message) {
  var payload_len: uint64;
  var masking_key_idx: uint64;
  var close_data: bytes;
  var effective_opcode: uint8;

  : bitfield(16) {
    fin: 0 &convert=cast<bool>($$);
    rsv: 1..3;
    opcode: 4..7 &convert=cast<uint8>($$);
    mask: 8 &convert=cast<bool>($$);
    payload_len1: 9..15;
  } &bit-order=spicy::BitOrder::MSB0;

  # Verify that this is either a continuation frame, or the Message does not
  # yet have a non-continuation opcode, but this frame does.
  : void &requires=(m.opcode != OPCODE_CONTINUATION && self.opcode == OPCODE_CONTINUATION || m.opcode == OPCODE_CONTINUATION && self.opcode != OPCODE_CONTINUATION);

  # Type/opcode to decide what to do with individual chunks.
  : void {
    self.effective_opcode = m.opcode != OPCODE_CONTINUATION ? m.opcode : self.opcode;
  }

  payload_len2: uint16 if (self.payload_len1 == 126);
  payload_len8: uint64 if (self.payload_len1 == 127);

  : void {
    self.payload_len = self.payload_len1;
    if ( self?.payload_len2 )
      self.payload_len = self.payload_len2;
    else if ( self?.payload_len8 )
      self.payload_len = self.payload_len8;
  }

  # This being an uint8[] allows masking_key[x] indexing, while a bytes
  # object would require *masking_key.at(i) which took roughly 20% more
  # runtime when I tested it.
  masking_key: uint8[] &size=4 if (self.mask);

  chunk: bytes &size=self.payload_len &chunked {
    # Don't use &convert with &chunked: https://github.com/zeek/spicy/issues/1661
    if ( self.mask ) {
      self.chunk = fast_unmask(self.masking_key_idx, self.masking_key, $$);
      self.masking_key_idx += |$$|;
    } else {
      self.chunk = $$;
    }

    # Forward TEXT and BINARY data to dowstream analyzers.
    if ( self.effective_opcode == OPCODE_TEXT || self.effective_opcode == OPCODE_BINARY )
      zeek::protocol_data_in(zeek::is_orig(), $$);

    # Accumulate the unmasked data in close_data if this a close frame
    # so it can be parsed by the outer Message. It's a bit of a hack.
    if ( self.effective_opcode == OPCODE_CLOSE )
      self.close_data += $$;
  }
};

type CloseFrame = unit {
  var status: uint16;
  var reason: bytes;

  : bytes &eod {
    if ( |$$| > 0 ) {
      self.status = cast<uint16>($$.sub(0, 2).to_uint(spicy::ByteOrder::Network));
      self.reason = $$.sub(2, 0);
    }
  }
};


public type Message = unit {
  # transient trickery
  var done: bool = False;

  var opcode: uint8;

  # Keep the first one persistent to have access
  # to the payload if it's a close frame.
  first_frame: Frame(self) {
    self.opcode = $$.opcode;
    self.done = $$.fin;
  }

  : Frame(self)[] &until=(self.done) if (!self.done) foreach {
    self.done = $$.fin;
  }

  : CloseFrame &parse-from=self.first_frame.close_data if (self.opcode == OPCODE_CLOSE);

  on %done {
    spicy::accept_input();
  }
};

public type Messages = unit {
  : Message[];
};
