// tests for SHA-1 / base64 / the WebSocket codec / the handshake parser
#include <cstdlib>
#include <cstring>

#include "ws_codec.h"
#include "test.h"  // CHECK / CHECK_EQ / CHECK_STR / report()

using namespace kr;

namespace {

void hexDigest(const uint8_t d[20], char out[41]) {
  static const char kHex[] = "0123456789abcdef";
  for (int i = 0; i < 20; i++) {
    out[2 * i] = kHex[d[i] >> 4];
    out[2 * i + 1] = kHex[d[i] & 0xF];
  }
  out[40] = '\0';
}

// Build a client frame the way a browser does: masked payload, mask key appended
// in front of it. Returns the frame length.
size_t buildClientFrame(uint8_t *out, size_t cap, uint8_t opcode, const char *payload,
                        size_t len, bool mask = true, bool fin = true, uint8_t rsv = 0,
                        uint64_t declaredLen = 0xFFFFFFFFull) {
  if (declaredLen == 0xFFFFFFFFull) declaredLen = len;
  // A test that asks for a frame larger than its own buffer is a broken test, not
  // a stack smash: say so and stop.
  if (2 + 8 + 4 + len > cap) {
    std::fprintf(stderr, "buildClientFrame: %zu bytes needed, buffer is %zu\n", 2 + 8 + 4 + len,
                 cap);
    std::abort();
  }
  size_t o = 0;
  out[o++] = static_cast<uint8_t>((fin ? 0x80 : 0x00) | (rsv << 4) | (opcode & 0x0F));
  const uint8_t m = mask ? 0x80 : 0x00;
  if (declaredLen < 126) {
    out[o++] = static_cast<uint8_t>(m | declaredLen);
  } else if (declaredLen <= 0xFFFF) {
    out[o++] = m | 126;
    out[o++] = static_cast<uint8_t>(declaredLen >> 8);
    out[o++] = static_cast<uint8_t>(declaredLen);
  } else {
    out[o++] = m | 127;
    for (int i = 0; i < 8; i++) out[o++] = static_cast<uint8_t>(declaredLen >> (56 - 8 * i));
  }
  uint8_t maskKey[4] = {0xDE, 0xAD, 0xBE, 0xEF};
  if (mask) {
    memcpy(out + o, maskKey, 4);
    o += 4;
  }
  for (size_t i = 0; i < len; i++) {
    out[o++] = mask ? static_cast<uint8_t>(payload[i] ^ maskKey[i & 3])
                    : static_cast<uint8_t>(payload[i]);
  }
  return o;
}

}  // namespace

int main() {
  char hex[64];

  CASE("SHA-1 against the published vectors");
  {
    uint8_t d[20];
    sha1("", 0, d);
    hexDigest(d, hex);
    CHECK_STR(hex, "da39a3ee5e6b4b0d3255bfef95601890afd80709");

    sha1("abc", 3, d);
    hexDigest(d, hex);
    CHECK_STR(hex, "a9993e364706816aba3e25717850c26c9cd0d89d");

    // 56 bytes is the padding boundary: with 0x80 + the 8-byte length the block
    // overflows, so the count lands in a second block. Classic SHA-1 bug, and the
    // reason this vector is here (digest cross-checked against Python hashlib).
    char fiftySix[57];
    memset(fiftySix, 'a', 56);
    fiftySix[56] = '\0';
    sha1(fiftySix, 56, d);
    hexDigest(d, hex);
    CHECK_STR(hex, "c2db330f6083854c99d4b5bfb6e8f29f201be699");

    // NIST's two-block message digest
    sha1("abcdbcdecdefdefgefghfghighijhijkijkljklmklmnlmnomnopnopq", 56, d);
    hexDigest(d, hex);
    CHECK_STR(hex, "84983e441c3bd26ebaae4aa1f95129e5e54670f1");

    // the canonical million-character vector: exercises many blocks
    {
      Sha1 many;
      char chunk[64];
      memset(chunk, 'a', sizeof(chunk));
      for (int i = 0; i < 1000000 / 64; i++) many.update(chunk, sizeof(chunk));
      many.finish(d);
      hexDigest(d, hex);
      CHECK_STR(hex, "34aa973cd4c4daa4f61eeb2bdbad27316534016f");
    }
  }

  CASE("SHA-1 incremental equals one-shot");
  {
    const char *text = "the quick brown fox jumps over the lazy dog, twice over";
    uint8_t one[20], many[20];
    sha1(text, strlen(text), one);
    Sha1 s;
    for (size_t i = 0; i < strlen(text); i++) s.update(text + i, 1);
    s.finish(many);
    CHECK(memcmp(one, many, 20) == 0);

    // reuse after finish() must start clean
    uint8_t again[20];
    s.update("abc", 3);
    s.finish(again);
    sha1("abc", 3, one);
    CHECK(memcmp(again, one, 20) == 0);
  }

  CASE("base64 against RFC 4648");
  {
    char out[32];
    CHECK_EQ(base64Encode(reinterpret_cast<const uint8_t *>("\0"), 0, out, sizeof(out)), 0u);
    CHECK_STR(out, "");
    CHECK_EQ(base64Encode(reinterpret_cast<const uint8_t *>("f"), 1, out, sizeof(out)), 4u);
    CHECK_STR(out, "Zg==");
    CHECK_EQ(base64Encode(reinterpret_cast<const uint8_t *>("fo"), 2, out, sizeof(out)), 4u);
    CHECK_STR(out, "Zm8=");
    CHECK_EQ(base64Encode(reinterpret_cast<const uint8_t *>("foo"), 3, out, sizeof(out)), 4u);
    CHECK_STR(out, "Zm9v");
    CHECK_EQ(base64Encode(reinterpret_cast<const uint8_t *>("foob"), 4, out, sizeof(out)), 8u);
    CHECK_STR(out, "Zm9vYg==");
    CHECK_EQ(base64Encode(reinterpret_cast<const uint8_t *>("fooba"), 5, out, sizeof(out)), 8u);
    CHECK_STR(out, "Zm9vYmE=");
    CHECK_EQ(base64Encode(reinterpret_cast<const uint8_t *>("foobar"), 6, out, sizeof(out)), 8u);
    CHECK_STR(out, "Zm9vYmFy");

    // one byte short of the requirement must refuse, not overwrite
    char tiny[4] = {'x', 'x', 'x', '\0'};
    CHECK_EQ(base64Encode(reinterpret_cast<const uint8_t *>("foo"), 3, tiny, 4), 0u);
    CHECK_STR(tiny, "xxx");
  }

  CASE("the RFC 6455 handshake example");
  {
    const char *key = "dGhlIHNhbXBsZSBub25jZQ==";
    char accept[64] = "";
    CHECK(ws::acceptFor(key, strlen(key), accept, sizeof(accept)));
    CHECK_STR(accept, "s3pPLMBiTxaQ9kYGzzhZRbK+xOo=");

    // a different key must produce a different accept, and short buffers refuse
    char other[64] = "";
    CHECK(ws::acceptFor("anotherKey12345=", strlen("anotherKey12345="), other, sizeof(other)));
    CHECK(std::strcmp(accept, other) != 0);
    CHECK(!ws::acceptFor("", 0, accept, sizeof(accept)));
    CHECK(!ws::acceptFor(key, strlen(key), accept, 8));
  }

  CASE("outbound frame encoding");
  {
    uint8_t buf[64];
    size_t n = ws::encodeText("A1", 2, buf, sizeof(buf));
    CHECK_EQ(n, 4u);
    CHECK_EQ(buf[0], 0x81);                  // FIN + text
    CHECK_EQ(buf[1], 2);                     // unmasked, length 2
    CHECK_EQ(buf[2], 'A');
    CHECK_EQ(buf[3], '1');

    n = ws::encodeText("", 0, buf, sizeof(buf));
    CHECK_EQ(n, 2u);
    CHECK_EQ(buf[1], 0);

    // 125 bytes is the last length that fits the 7-bit field, 126 the first that
    // needs the 16-bit extension
    char blob[200];
    memset(blob, 'x', sizeof(blob));
    uint8_t big[400];
    n = ws::encodeText(blob, 125, big, sizeof(big));
    CHECK_EQ(n, 127u);
    CHECK_EQ(big[1], 125);
    CHECK_EQ(big[2], 'x');

    n = ws::encodeText(blob, 126, big, sizeof(big));
    CHECK_EQ(n, 130u);
    CHECK_EQ(big[1], 126);
    CHECK_EQ(big[2], 0);
    CHECK_EQ(big[3], 126);
    CHECK_EQ(big[4], 'x');

    // a buffer that is one byte too small refuses without writing the frame
    uint8_t tight[10] = {0};
    CHECK_EQ(ws::encodeText("0123456789ab", 12, tight, sizeof(tight)), 0u);
    CHECK_EQ(tight[0], 0);
  }

  CASE("inbound frame decodes, byte by byte");
  {
    uint8_t wire[128], payloadBuf[64];
    const size_t n = buildClientFrame(wire, sizeof(wire), ws::kOpText, "S=0a:04", 7);
    CHECK_EQ(n, 13u);                        // 2 header + 4 mask + 7 payload

    ws::Reader r(payloadBuf, sizeof(payloadBuf));
    ws::In res = ws::In::NeedMore;
    for (size_t i = 0; i < n; i++) {
      res = r.feed(wire[i]);
      if (i + 1 < n) CHECK(res == ws::In::NeedMore);
    }
    CHECK(res == ws::In::Text);
    CHECK(r.finished());
    CHECK_EQ(r.payloadLen(), 7u);
    CHECK(memcmp(r.payload(), "S=0a:04", 7) == 0);
    CHECK_EQ(r.error(), ws::Err::None);

    // feeding more after completion keeps reporting the same frame until reset()
    CHECK(r.feed(0x00) == ws::In::Text);
    r.reset();
    CHECK(!r.finished());
  }

  CASE("inbound frame decodes from a whole buffer");
  {
    uint8_t wire[256], payloadBuf[64];
    const size_t a = buildClientFrame(wire, sizeof(wire), ws::kOpText, "A1", 2);
    const size_t b = buildClientFrame(wire + a, sizeof(wire) - a, ws::kOpText, "K", 1);

    ws::Reader r(payloadBuf, sizeof(payloadBuf));
    size_t consumed = 0;
    CHECK(r.feedAll(wire, a + b, &consumed) == ws::In::Text);
    CHECK_EQ(consumed, a);
    CHECK(memcmp(r.payload(), "A1", 2) == 0);
    r.reset();
    CHECK(r.feedAll(wire + consumed, a + b - consumed, &consumed) == ws::In::Text);
    CHECK_EQ(r.payloadLen(), 1u);
    CHECK_EQ(r.payload()[0], 'K');
  }

  CASE("empty frames and the 16-bit length field");
  {
    uint8_t wire[512];
    char blob[130];
    for (size_t i = 0; i < sizeof(blob); i++) blob[i] = static_cast<char>('a' + (i % 26));

    uint8_t payloadBuf[200];
    ws::Reader r(payloadBuf, sizeof(payloadBuf));

    const size_t n = buildClientFrame(wire, sizeof(wire), ws::kOpText, nullptr, 0);
    CHECK_EQ(n, 2u + 4u);                    // a masked frame sends its mask key anyway
    CHECK(r.feedAll(wire, n) == ws::In::Text);
    CHECK_EQ(r.payloadLen(), 0u);

    r.reset();
    const size_t m = buildClientFrame(wire, sizeof(wire), ws::kOpText, blob, 126);
    CHECK_EQ(m, 4u + 4u + 126u);              // 2 + 2 extended length, mask, payload
    CHECK(r.feedAll(wire, m) == ws::In::Text);
    CHECK_EQ(r.payloadLen(), 126u);
    CHECK(memcmp(r.payload(), blob, 126) == 0);

    // A 64-bit length field is parsed (8 extension bytes) and then refused, because
    // the payload it announces cannot fit the caller's buffer. The decoder never
    // buffers beyond the cap it was constructed with.
    r.reset();
    const size_t k = buildClientFrame(wire, sizeof(wire), ws::kOpText, blob, 2, true, true,
                                      0, 0x10000);
    CHECK_EQ(k, 2u + 8u + 4u + 2u);
    CHECK(r.feedAll(wire, k) == ws::In::Error);
    CHECK(r.error() == ws::Err::TooLong);
  }

  CASE("control frames");
  {
    uint8_t wire[64], payloadBuf[64];
    ws::Reader r(payloadBuf, sizeof(payloadBuf));
    size_t n = buildClientFrame(wire, sizeof(wire), ws::kOpPing, "hi", 2);
    CHECK(r.feedAll(wire, n) == ws::In::Ping);
    CHECK_EQ(r.payloadLen(), 2u);

    r.reset();
    n = buildClientFrame(wire, sizeof(wire), ws::kOpPong, nullptr, 0);
    CHECK(r.feedAll(wire, n) == ws::In::Pong);

    r.reset();
    n = buildClientFrame(wire, sizeof(wire), ws::kOpClose, "\x03\xe8", 2);
    CHECK(r.feedAll(wire, n) == ws::In::Close);
    CHECK_EQ(r.payload()[0], 3);
    CHECK_EQ(r.payload()[1], 232);
  }

  CASE("protocol deviations are hard errors, not graceful surprises");
  {
    uint8_t wire[600], payloadBuf[64];
    ws::Reader r(payloadBuf, sizeof(payloadBuf));

    // unmasked client frame: RFC 6455 5.1 says the server must fail the connection
    size_t n = buildClientFrame(wire, sizeof(wire), ws::kOpText, "A1", 2, /*mask=*/false);
    CHECK(r.feedAll(wire, n) == ws::In::Error);
    CHECK(r.error() == ws::Err::NotMasked);

    r.reset();
    n = buildClientFrame(wire, sizeof(wire), ws::kOpBinary, "A1", 2);
    CHECK(r.feedAll(wire, n) == ws::In::Error);
    CHECK(r.error() == ws::Err::BadOpcode);

    r.reset();
    n = buildClientFrame(wire, sizeof(wire), ws::kOpText, "A1", 2, true, /*fin=*/false);
    CHECK(r.feedAll(wire, n) == ws::In::Error);
    CHECK(r.error() == ws::Err::Fragmented);

    r.reset();
    n = buildClientFrame(wire, sizeof(wire), ws::kOpText, "A1", 2, true, true, /*rsv=*/1);
    CHECK(r.feedAll(wire, n) == ws::In::Error);
    CHECK(r.error() == ws::Err::ReservedBits);

    char blob[200];
    for (size_t i = 0; i < sizeof(blob); i++) blob[i] = static_cast<char>('A' + (i % 26));

    r.reset();
    n = buildClientFrame(wire, sizeof(wire), ws::kOpPing, blob, 126);
    CHECK(r.feedAll(wire, n) == ws::In::Error);
    CHECK(r.error() == ws::Err::ControlTooLong);

    // a text frame larger than the reader's buffer is refused outright
    r.reset();
    n = buildClientFrame(wire, sizeof(wire), ws::kOpText, blob, 123);
    CHECK(r.feedAll(wire, n) == ws::In::Error);
    CHECK(r.error() == ws::Err::TooLong);

    // 64-bit length field declaring more than the buffer can hold
    r.reset();
    n = buildClientFrame(wire, sizeof(wire), ws::kOpText, "A1", 2, true, true, 0,
                         /*declaredLen=*/1u << 20);
    CHECK(r.feedAll(wire, n) == ws::In::Error);
    CHECK(r.error() == ws::Err::TooLong);
  }

  CASE("handshake request parsing");
  {
    const char *req =
        "GET /ws?token=hunter2 HTTP/1.1\r\n"
        "Host: keyrelay.local\r\n"
        "Upgrade: websocket\r\n"
        "Connection: Upgrade\r\n"
        "Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n"
        "Sec-WebSocket-Version: 13\r\n"
        "\r\n";
    ws::Request rq;
    CHECK_EQ(ws::parseRequest(req, strlen(req), rq), 1);
    CHECK(rq.complete);
    CHECK_STR(rq.method, "GET");
    CHECK_STR(rq.path, "/ws");
    CHECK_STR(rq.query, "token=hunter2");
    CHECK(rq.upgradeWebsocket);
    CHECK_STR(rq.wsKey, "dGhlIHNhbXBsZSBub25jZQ==");

    // and the accept value for the key we just parsed is the RFC's own example
    char accept[64] = "";
    CHECK(ws::acceptFor(rq.wsKey, strlen(rq.wsKey), accept, sizeof(accept)));
    CHECK_STR(accept, "s3pPLMBiTxaQ9kYGzzhZRbK+xOo=");
  }

  CASE("handshake request parsing is case-insensitive and whitespace-tolerant");
  {
    const char *req =
        "GET / HTTP/1.1\r\n"
        "upGRADE:   WebSocket\r\n"
        "sec-WEBSOCKET-key:  dGhlIHNhbXBsZSBub25jZQ==  \r\n"
        "\r\n";
    ws::Request rq;
    CHECK_EQ(ws::parseRequest(req, strlen(req), rq), 1);
    CHECK(rq.upgradeWebsocket);
    CHECK_STR(rq.wsKey, "dGhlIHNhbXBsZSBub25jZQ==");
  }

  CASE("plain HTTP request is not an upgrade");
  {
    const char *req = "GET /status?token=abc HTTP/1.0\r\nHost: x\r\n\r\n";
    ws::Request rq;
    CHECK_EQ(ws::parseRequest(req, strlen(req), rq), 1);
    CHECK(!rq.upgradeWebsocket);
    CHECK_STR(rq.path, "/status");
    CHECK_EQ(rq.wsKey[0], 0);

    char token[64] = "";
    CHECK(ws::tokenFromQuery(rq.query, strlen(rq.query), token, sizeof(token)));
    CHECK_STR(token, "abc");
  }

  CASE("incomplete or malformed requests are distinguishable");
  {
    ws::Request rq;
    const char *partial = "GET / HTTP/1.1\r\nUpgrade: websocket\r\n";
    CHECK_EQ(ws::parseRequest(partial, strlen(partial), rq), 0);   // need more bytes
    const char *noSpace = "GARBAGE\r\n\r\n";
    CHECK_EQ(ws::parseRequest(noSpace, strlen(noSpace), rq), -1);
    const char *noLineEnd = "GET /";
    CHECK_EQ(ws::parseRequest(noLineEnd, strlen(noLineEnd), rq), 0);
  }

  CASE("query token extraction");
  {
    char token[64];
    struct { const char *q; bool ok; const char *value; } cases[] = {
        {"token=abc", true, "abc"},
        {"a=1&token=abc", true, "abc"},
        {"token=abc&b=2", true, "abc"},
        {"a=1", false, ""},
        {"token=", false, ""},
        {"xtoken=abc", false, ""},          // must be a whole parameter
        {"a=1xtoken=abc", false, ""},
        {"", false, ""},
    };
    for (const auto &c : cases) {
      token[0] = '\0';
      const bool got = ws::tokenFromQuery(c.q, strlen(c.q), token, sizeof(token));
      CHECK(got == c.ok);
      if (c.ok) CHECK_STR(token, c.value);
    }
    // a token longer than the caller's buffer is refused, not truncated
    char tiny[4] = "";
    CHECK(!ws::tokenFromQuery("token=abcdefghij", strlen("token=abcdefghij"), tiny,
                              sizeof(tiny)));
  }

  return report("ws_codec");
}