// Copyright (c) 2012 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "net/websockets/websocket_frame.h"

#include <algorithm>
#include <vector>

#include "base/basictypes.h"
#include "base/command_line.h"
#include "base/logging.h"
#include "base/memory/aligned_memory.h"
#include "base/string_number_conversions.h"
#include "base/stringprintf.h"
#include "base/time.h"
#include "net/base/net_errors.h"
#include "testing/gtest/include/gtest/gtest.h"

// Run
//   out/Release/net_unittests --websocket-mask-iterations=100000
//      --gtest_filter='WebSocketFrameTestMaskBenchmark.*'
// to benchmark the MaskWebSocketFramePayload() function.
static const char kBenchmarkIterations[] = "websocket-mask-iterations";
static const int kDefaultIterations = 10;
static const int kLongPayloadSize = 1 << 16;

namespace net {

TEST(WebSocketFrameHeaderTest, FrameLengths) {
  struct TestCase {
    const char* frame_header;
    size_t frame_header_length;
    uint64 frame_length;
  };
  static const TestCase kTests[] = {
    { "\x81\x00", 2, GG_UINT64_C(0) },
    { "\x81\x7D", 2, GG_UINT64_C(125) },
    { "\x81\x7E\x00\x7E", 4, GG_UINT64_C(126) },
    { "\x81\x7E\xFF\xFF", 4, GG_UINT64_C(0xFFFF) },
    { "\x81\x7F\x00\x00\x00\x00\x00\x01\x00\x00", 10, GG_UINT64_C(0x10000) },
    { "\x81\x7F\x7F\xFF\xFF\xFF\xFF\xFF\xFF\xFF", 10,
      GG_UINT64_C(0x7FFFFFFFFFFFFFFF) }
  };
  static const int kNumTests = ARRAYSIZE_UNSAFE(kTests);

  for (int i = 0; i < kNumTests; ++i) {
    WebSocketFrameHeader header;
    header.final = true;
    header.reserved1 = false;
    header.reserved2 = false;
    header.reserved3 = false;
    header.opcode = WebSocketFrameHeader::kOpCodeText;
    header.masked = false;
    header.payload_length = kTests[i].frame_length;

    std::vector<char> expected_output(
        kTests[i].frame_header,
        kTests[i].frame_header + kTests[i].frame_header_length);
    std::vector<char> output(expected_output.size());
    EXPECT_EQ(static_cast<int>(expected_output.size()),
              WriteWebSocketFrameHeader(header, NULL, &output.front(),
                                        output.size()));
    EXPECT_EQ(expected_output, output);
  }
}

TEST(WebSocketFrameHeaderTest, FrameLengthsWithMasking) {
  static const char kMaskingKey[] = "\xDE\xAD\xBE\xEF";
  COMPILE_ASSERT(ARRAYSIZE_UNSAFE(kMaskingKey) - 1 ==
                 WebSocketFrameHeader::kMaskingKeyLength,
                 incorrect_masking_key_size);

  struct TestCase {
    const char* frame_header;
    size_t frame_header_length;
    uint64 frame_length;
  };
  static const TestCase kTests[] = {
    { "\x81\x80\xDE\xAD\xBE\xEF", 6, GG_UINT64_C(0) },
    { "\x81\xFD\xDE\xAD\xBE\xEF", 6, GG_UINT64_C(125) },
    { "\x81\xFE\x00\x7E\xDE\xAD\xBE\xEF", 8, GG_UINT64_C(126) },
    { "\x81\xFE\xFF\xFF\xDE\xAD\xBE\xEF", 8, GG_UINT64_C(0xFFFF) },
    { "\x81\xFF\x00\x00\x00\x00\x00\x01\x00\x00\xDE\xAD\xBE\xEF", 14,
      GG_UINT64_C(0x10000) },
    { "\x81\xFF\x7F\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xDE\xAD\xBE\xEF", 14,
      GG_UINT64_C(0x7FFFFFFFFFFFFFFF) }
  };
  static const int kNumTests = ARRAYSIZE_UNSAFE(kTests);

  WebSocketMaskingKey masking_key;
  std::copy(kMaskingKey, kMaskingKey + WebSocketFrameHeader::kMaskingKeyLength,
            masking_key.key);

  for (int i = 0; i < kNumTests; ++i) {
    WebSocketFrameHeader header;
    header.final = true;
    header.reserved1 = false;
    header.reserved2 = false;
    header.reserved3 = false;
    header.opcode = WebSocketFrameHeader::kOpCodeText;
    header.masked = true;
    header.payload_length = kTests[i].frame_length;

    std::vector<char> expected_output(
        kTests[i].frame_header,
        kTests[i].frame_header + kTests[i].frame_header_length);
    std::vector<char> output(expected_output.size());
    EXPECT_EQ(static_cast<int>(expected_output.size()),
              WriteWebSocketFrameHeader(header, &masking_key,
                                        &output.front(), output.size()));
    EXPECT_EQ(expected_output, output);
  }
}

TEST(WebSocketFrameHeaderTest, FrameOpCodes) {
  struct TestCase {
    const char* frame_header;
    size_t frame_header_length;
    WebSocketFrameHeader::OpCode opcode;
  };
  static const TestCase kTests[] = {
    { "\x80\x00", 2, WebSocketFrameHeader::kOpCodeContinuation },
    { "\x81\x00", 2, WebSocketFrameHeader::kOpCodeText },
    { "\x82\x00", 2, WebSocketFrameHeader::kOpCodeBinary },
    { "\x88\x00", 2, WebSocketFrameHeader::kOpCodeClose },
    { "\x89\x00", 2, WebSocketFrameHeader::kOpCodePing },
    { "\x8A\x00", 2, WebSocketFrameHeader::kOpCodePong },
    // These are undefined opcodes, but the builder should accept them anyway.
    { "\x83\x00", 2, WebSocketFrameHeader::OpCode(0x3) },
    { "\x84\x00", 2, WebSocketFrameHeader::OpCode(0x4) },
    { "\x85\x00", 2, WebSocketFrameHeader::OpCode(0x5) },
    { "\x86\x00", 2, WebSocketFrameHeader::OpCode(0x6) },
    { "\x87\x00", 2, WebSocketFrameHeader::OpCode(0x7) },
    { "\x8B\x00", 2, WebSocketFrameHeader::OpCode(0xB) },
    { "\x8C\x00", 2, WebSocketFrameHeader::OpCode(0xC) },
    { "\x8D\x00", 2, WebSocketFrameHeader::OpCode(0xD) },
    { "\x8E\x00", 2, WebSocketFrameHeader::OpCode(0xE) },
    { "\x8F\x00", 2, WebSocketFrameHeader::OpCode(0xF) }
  };
  static const int kNumTests = ARRAYSIZE_UNSAFE(kTests);

  for (int i = 0; i < kNumTests; ++i) {
    WebSocketFrameHeader header;
    header.final = true;
    header.reserved1 = false;
    header.reserved2 = false;
    header.reserved3 = false;
    header.opcode = kTests[i].opcode;
    header.masked = false;
    header.payload_length = 0;

    std::vector<char> expected_output(
        kTests[i].frame_header,
        kTests[i].frame_header + kTests[i].frame_header_length);
    std::vector<char> output(expected_output.size());
    EXPECT_EQ(static_cast<int>(expected_output.size()),
              WriteWebSocketFrameHeader(header, NULL,
                                        &output.front(), output.size()));
    EXPECT_EQ(expected_output, output);
  }
}

TEST(WebSocketFrameHeaderTest, FinalBitAndReservedBits) {
  struct TestCase {
    const char* frame_header;
    size_t frame_header_length;
    bool final;
    bool reserved1;
    bool reserved2;
    bool reserved3;
  };
  static const TestCase kTests[] = {
    { "\x81\x00", 2, true, false, false, false },
    { "\x01\x00", 2, false, false, false, false },
    { "\xC1\x00", 2, true, true, false, false },
    { "\xA1\x00", 2, true, false, true, false },
    { "\x91\x00", 2, true, false, false, true },
    { "\x71\x00", 2, false, true, true, true },
    { "\xF1\x00", 2, true, true, true, true }
  };
  static const int kNumTests = ARRAYSIZE_UNSAFE(kTests);

  for (int i = 0; i < kNumTests; ++i) {
    WebSocketFrameHeader header;
    header.final = kTests[i].final;
    header.reserved1 = kTests[i].reserved1;
    header.reserved2 = kTests[i].reserved2;
    header.reserved3 = kTests[i].reserved3;
    header.opcode = WebSocketFrameHeader::kOpCodeText;
    header.masked = false;
    header.payload_length = 0;

    std::vector<char> expected_output(
        kTests[i].frame_header,
        kTests[i].frame_header + kTests[i].frame_header_length);
    std::vector<char> output(expected_output.size());
    EXPECT_EQ(static_cast<int>(expected_output.size()),
              WriteWebSocketFrameHeader(header, NULL,
                                        &output.front(), output.size()));
    EXPECT_EQ(expected_output, output);
  }
}

TEST(WebSocketFrameHeaderTest, InsufficientBufferSize) {
  struct TestCase {
    uint64 payload_length;
    bool masked;
    size_t expected_header_size;
  };
  static const TestCase kTests[] = {
    { GG_UINT64_C(0), false, 2u },
    { GG_UINT64_C(125), false, 2u },
    { GG_UINT64_C(126), false, 4u },
    { GG_UINT64_C(0xFFFF), false, 4u },
    { GG_UINT64_C(0x10000), false, 10u },
    { GG_UINT64_C(0x7FFFFFFFFFFFFFFF), false, 10u },
    { GG_UINT64_C(0), true, 6u },
    { GG_UINT64_C(125), true, 6u },
    { GG_UINT64_C(126), true, 8u },
    { GG_UINT64_C(0xFFFF), true, 8u },
    { GG_UINT64_C(0x10000), true, 14u },
    { GG_UINT64_C(0x7FFFFFFFFFFFFFFF), true, 14u }
  };
  static const int kNumTests = ARRAYSIZE_UNSAFE(kTests);

  for (int i = 0; i < kNumTests; ++i) {
    WebSocketFrameHeader header;
    header.final = true;
    header.reserved1 = false;
    header.reserved2 = false;
    header.reserved3 = false;
    header.opcode = WebSocketFrameHeader::kOpCodeText;
    header.masked = kTests[i].masked;
    header.payload_length = kTests[i].payload_length;

    char dummy_buffer[14];
    // Set an insufficient size to |buffer_size|.
    EXPECT_EQ(ERR_INVALID_ARGUMENT,
              WriteWebSocketFrameHeader(header, NULL, dummy_buffer,
                                        kTests[i].expected_header_size - 1));
  }
}

TEST(WebSocketFrameTest, MaskPayload) {
  struct TestCase {
    const char* masking_key;
    uint64 frame_offset;
    const char* input;
    const char* output;
    size_t data_length;
  };
  static const TestCase kTests[] = {
    { "\xDE\xAD\xBE\xEF", 0, "FooBar", "\x98\xC2\xD1\xAD\xBF\xDF", 6 },
    { "\xDE\xAD\xBE\xEF", 1, "FooBar", "\xEB\xD1\x80\x9C\xCC\xCC", 6 },
    { "\xDE\xAD\xBE\xEF", 2, "FooBar", "\xF8\x80\xB1\xEF\xDF\x9D", 6 },
    { "\xDE\xAD\xBE\xEF", 3, "FooBar", "\xA9\xB1\xC2\xFC\x8E\xAC", 6 },
    { "\xDE\xAD\xBE\xEF", 4, "FooBar", "\x98\xC2\xD1\xAD\xBF\xDF", 6 },
    { "\xDE\xAD\xBE\xEF", 42, "FooBar", "\xF8\x80\xB1\xEF\xDF\x9D", 6 },
    { "\xDE\xAD\xBE\xEF", 0, "", "", 0 },
    { "\xDE\xAD\xBE\xEF", 0, "\xDE\xAD\xBE\xEF", "\x00\x00\x00\x00", 4 },
    { "\xDE\xAD\xBE\xEF", 0, "\x00\x00\x00\x00", "\xDE\xAD\xBE\xEF", 4 },
    { "\x00\x00\x00\x00", 0, "FooBar", "FooBar", 6 },
    { "\xFF\xFF\xFF\xFF", 0, "FooBar", "\xB9\x90\x90\xBD\x9E\x8D", 6 },
  };
  static const int kNumTests = ARRAYSIZE_UNSAFE(kTests);

  for (int i = 0; i < kNumTests; ++i) {
    WebSocketMaskingKey masking_key;
    std::copy(kTests[i].masking_key,
              kTests[i].masking_key + WebSocketFrameHeader::kMaskingKeyLength,
              masking_key.key);
    std::vector<char> frame_data(kTests[i].input,
                                 kTests[i].input + kTests[i].data_length);
    std::vector<char> expected_output(kTests[i].output,
                                      kTests[i].output + kTests[i].data_length);
    MaskWebSocketFramePayload(masking_key,
                              kTests[i].frame_offset,
                              frame_data.empty() ? NULL : &frame_data.front(),
                              frame_data.size());
    EXPECT_EQ(expected_output, frame_data);
  }
}

// Check that all combinations of alignment, frame offset and chunk size work
// correctly for MaskWebSocketFramePayload(). This is mainly used to ensure that
// vectorisation optimisations don't break anything. We could take a "white box"
// approach and only test the edge cases, but since the exhaustive "black box"
// approach runs in acceptable time, we don't have to take the risk of being
// clever.
//
// This brute-force approach runs in O(N^3) time where N is the size of the
// maximum vector size we want to test again. This might need reconsidering if
// MaskWebSocketFramePayload() is ever optimised for a dedicated vector
// architecture.
TEST(WebSocketFrameTest, MaskPayloadAlignment) {
  // This reflects what might be implemented in the future, rather than
  // the current implementation. FMA3 and FMA4 support 256-bit vector ops.
  static const size_t kMaxVectorSizeInBits = 256;
  static const size_t kMaxVectorSize = kMaxVectorSizeInBits / 8;
  static const size_t kMaxVectorAlignment = kMaxVectorSize;
  static const size_t kMaskingKeyLength =
      WebSocketFrameHeader::kMaskingKeyLength;
  static const size_t kScratchBufferSize =
      kMaxVectorAlignment + kMaxVectorSize * 2;
  static const char kTestMask[] = "\xd2\xba\x5a\xbe";
  // We use 786 bits of random input to reduce the risk of correlated errors.
  static const char kTestInput[] =
      { "\x3d\x77\x1d\x1b\x19\x8c\x48\xa3\x19\x6d\xf7\xcc\x39\xe7\x57\x0b"
        "\x69\x8c\xda\x4b\xfc\xac\x2c\xd3\x49\x96\x6e\x8a\x7b\x5a\x32\x76"
        "\xd0\x11\x43\xa0\x89\xfc\x76\x2b\x10\x2f\x4c\x7b\x4f\xa6\xdd\xe4"
        "\xfc\x8e\xd8\x72\xcf\x7e\x37\xcd\x31\xcd\xc1\xc0\x89\x0c\xa7\x4c"
        "\xda\xa8\x4b\x75\xa1\xcb\xa9\x77\x19\x4d\x6e\xdf\xc8\x08\x1c\xb6"
        "\x6d\xfb\x38\x04\x44\xd5\xba\x57\x9f\x76\xb0\x2e\x07\x91\xe6\xa8"
      };
  static const size_t kTestInputSize = arraysize(kTestInput) - 1;
  static const char kTestOutput[] =
      { "\xef\xcd\x47\xa5\xcb\x36\x12\x1d\xcb\xd7\xad\x72\xeb\x5d\x0d\xb5"
        "\xbb\x36\x80\xf5\x2e\x16\x76\x6d\x9b\x2c\x34\x34\xa9\xe0\x68\xc8"
        "\x02\xab\x19\x1e\x5b\x46\x2c\x95\xc2\x95\x16\xc5\x9d\x1c\x87\x5a"
        "\x2e\x34\x82\xcc\x1d\xc4\x6d\x73\xe3\x77\x9b\x7e\x5b\xb6\xfd\xf2"
        "\x08\x12\x11\xcb\x73\x71\xf3\xc9\xcb\xf7\x34\x61\x1a\xb2\x46\x08"
        "\xbf\x41\x62\xba\x96\x6f\xe0\xe9\x4d\xcc\xea\x90\xd5\x2b\xbc\x16"
      };
  COMPILE_ASSERT(arraysize(kTestInput) == arraysize(kTestOutput),
                 output_and_input_arrays_have_the_same_length);
  scoped_ptr_malloc<char, base::ScopedPtrAlignedFree> scratch(
      static_cast<char*>(base::AlignedAlloc(kScratchBufferSize,
                                            kMaxVectorAlignment)));
  WebSocketMaskingKey masking_key;
  std::copy(kTestMask, kTestMask + kMaskingKeyLength, masking_key.key);
  for (size_t frame_offset = 0;
       frame_offset < kMaskingKeyLength;
       ++frame_offset) {
    for (size_t alignment = 0; alignment < kMaxVectorAlignment; ++alignment) {
      char* const aligned_scratch = scratch.get() + alignment;
      const size_t aligned_len =
          std::min(kScratchBufferSize - alignment,
                   kTestInputSize - frame_offset);
      for (size_t chunk_size = 1; chunk_size < kMaxVectorSize; ++chunk_size) {
        memcpy(aligned_scratch, kTestInput + frame_offset, aligned_len);
        for (size_t chunk_start = 0;
             chunk_start < aligned_len;
             chunk_start += chunk_size) {
          const size_t this_chunk_size = std::min(chunk_size,
                                                  aligned_len - chunk_start);
          MaskWebSocketFramePayload(masking_key,
                                    frame_offset + chunk_start,
                                    aligned_scratch + chunk_start,
                                    this_chunk_size);
        }
        // Stop the test if it fails, since we don't want to spew thousands of
        // failures.
        ASSERT_TRUE(std::equal(aligned_scratch, aligned_scratch + aligned_len,
                               kTestOutput + frame_offset))
            << "Output failed to match for frame_offset="
            << frame_offset
            << ", alignment="
            << alignment
            << ", chunk_size="
            << chunk_size;
      }
    }
  }
}

class WebSocketFrameTestMaskBenchmark : public testing::Test {
 public:
  WebSocketFrameTestMaskBenchmark()
      : iterations_(kDefaultIterations) {}

  void SetUp() {
    std::string iterations(
        CommandLine::ForCurrentProcess()->GetSwitchValueASCII(
            kBenchmarkIterations));
    int benchmark_iterations = 0;
    if (!iterations.empty() && base::StringToInt(iterations,
                                                 &benchmark_iterations)) {
      iterations_ = benchmark_iterations;
    }
  }

  void Benchmark(const char* const payload, size_t size) {
    std::vector<char> scratch(payload, payload + size);
    static const char kMaskingKey[] = "\xFE\xED\xBE\xEF";
    COMPILE_ASSERT(arraysize(kMaskingKey) ==
                   WebSocketFrameHeader::kMaskingKeyLength + 1,
                   incorrect_masking_key_size);
    WebSocketMaskingKey masking_key;
    std::copy(kMaskingKey, kMaskingKey +
              WebSocketFrameHeader::kMaskingKeyLength,
              masking_key.key);
    LOG(INFO) << "Benchmarking MaskWebSocketFramePayload() for "
              << iterations_ << " iterations";
    using base::TimeTicks;
    TimeTicks start = TimeTicks::HighResNow();
    for (int x = 0; x < iterations_; ++x) {
      MaskWebSocketFramePayload(masking_key, x % size, &scratch.front(),
                                scratch.size());
    }
    double total_time_ms =
        1000 * (TimeTicks::HighResNow() - start).InMillisecondsF() /
        iterations_;
    LOG(INFO) << "Payload size " << size
              << StringPrintf(" took %.03f microseconds per iteration",
                              total_time_ms);
  }

 private:
  int iterations_;

  DISALLOW_COPY_AND_ASSIGN(WebSocketFrameTestMaskBenchmark);
};

TEST_F(WebSocketFrameTestMaskBenchmark, BenchmarkMaskShortPayload) {
  static const char kShortPayload[] = "Short Payload";
  Benchmark(kShortPayload, arraysize(kShortPayload));
}

TEST_F(WebSocketFrameTestMaskBenchmark, BenchmarkMaskLongPayload) {
  scoped_array<char> payload(new char[kLongPayloadSize]);
  std::fill(payload.get(), payload.get() + kLongPayloadSize, 'a');
  Benchmark(payload.get(), kLongPayloadSize);
}

}  // namespace net