/*
 * Copyright (C) 2021 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "test/gtest_and_gmock.h"

#include "perfetto/protozero/message.h"
#include "perfetto/protozero/scattered_heap_buffer.h"
#include "src/protozero/filtering/message_tokenizer.h"

namespace protozero {

using proto_utils::ProtoWireType;
using ::testing::ElementsAre;
using Token = MessageTokenizer::Token;

// For ASSERT_THAT(ElementsAre(...))
inline std::ostream& operator<<(std::ostream& stream, const Token& t) {
  stream << "{" << t.field_id << ", ";
  switch (t.type) {
    case ProtoWireType::kVarInt:
      stream << "varint, ";
      break;
    case ProtoWireType::kFixed32:
      stream << "fixed32, ";
      break;
    case ProtoWireType::kFixed64:
      stream << "fixed64, ";
      break;
    case ProtoWireType::kLengthDelimited:
      stream << "lendelim, ";
      break;
    default:
      stream << "???, ";
      break;
  }
  stream << t.value << "}";
  return stream;
}

namespace {

TEST(MessageTokenizerTest, FlatMessage) {
  HeapBuffered<Message> msg;
  msg->AppendVarInt(/*field_id*/ 1, 42u);
  msg->AppendVarInt(/*field_id*/ 1, 1000u);
  msg->AppendVarInt(/*field_id*/ 2, 1000000000ull);
  msg->AppendVarInt(/*field_id*/ 3, 0xFF001234DEADBEEFull);
  msg->AppendString(/*field_id*/ 4, "foo");
  msg->AppendFixed(/*field_id*/ 5, 0xFFAAFFFFu);
  msg->AppendString(/*field_id*/ 4, "foobar");
  msg->AppendFixed(/*field_id*/ 6, uint64_t(1ull << 63));
  msg->AppendVarInt(/*field_id*/ 1000, 1001ull);
  msg->AppendVarInt(/*field_id*/ 1000000, 1000001ull);
  msg->AppendVarInt(/*field_id*/ 1 << 28, uint64_t(1ull << 63));

  // Treat all len-delimited fields as strings/bytes and just eat their payload.
  MessageTokenizer tokenizer;
  std::vector<Token> tokens;
  size_t eat_bytes = 0;
  for (uint8_t octet : msg.SerializeAsArray()) {
    if (eat_bytes > 0) {
      --eat_bytes;
      continue;
    }
    auto token = tokenizer.Push(octet);
    if (token.valid())
      tokens.emplace_back(token);
    if (token.type == ProtoWireType::kLengthDelimited) {
      ASSERT_EQ(eat_bytes, 0u);
      eat_bytes = static_cast<size_t>(token.value);
    }
  }
  EXPECT_TRUE(tokenizer.idle());
  EXPECT_THAT(
      tokens,
      ElementsAre(
          Token{1, ProtoWireType::kVarInt, 42u},
          Token{1, ProtoWireType::kVarInt, 1000u},
          Token{2, ProtoWireType::kVarInt, 1000000000ull},
          Token{3, ProtoWireType::kVarInt, 0xFF001234DEADBEEFull},
          Token{4, ProtoWireType::kLengthDelimited, 3},
          Token{5, ProtoWireType::kFixed32, 0xFFAAFFFFu},
          Token{4, ProtoWireType::kLengthDelimited, 6},
          Token{6, ProtoWireType::kFixed64, uint64_t(1ull << 63)},
          Token{1000, ProtoWireType::kVarInt, 1001ull},
          Token{1000000, ProtoWireType::kVarInt, 1000001ull},
          Token{1 << 28, ProtoWireType::kVarInt, uint64_t(1ull << 63)}));
}

TEST(MessageTokenizerTest, NestedMessage) {
  HeapBuffered<Message> msg;
  msg->AppendVarInt(/*field_id*/ 1, 101u);
  {
    auto* nested = msg->BeginNestedMessage<Message>(2);
    nested->AppendVarInt(/*field_id*/ 3, 103u);
    nested->AppendFixed(/*field_id*/ 4, 104u);
    {
      auto* nested2 = nested->BeginNestedMessage<Message>(5);
      nested2->AppendVarInt(/*field_id*/ 6, 106u);
      nested2->AppendFixed(/*field_id*/ 7, 107u);
      nested2->Finalize();
    }
    nested->AppendFixed(/*field_id*/ 8, 0x42420000u);
    nested->Finalize();
  }
  msg->AppendFixed(/*field_id*/ 9, uint64_t(1ull << 63));

  // Tokenize the message. This treat all len delimited fields as submessage
  // and test the recursion logic.
  MessageTokenizer tokenizer;
  std::vector<Token> tokens;
  for (uint8_t octet : msg.SerializeAsArray()) {
    auto token = tokenizer.Push(octet);
    if (token.valid())
      tokens.emplace_back(token);
  }
  EXPECT_TRUE(tokenizer.idle());
  EXPECT_THAT(
      tokens,
      ElementsAre(Token{1, ProtoWireType::kVarInt, 101u},
                  Token{2, ProtoWireType::kLengthDelimited, 24u},
                  Token{3, ProtoWireType::kVarInt, 103u},
                  Token{4, ProtoWireType::kFixed32, 104u},
                  Token{5, ProtoWireType::kLengthDelimited, 7},
                  Token{6, ProtoWireType::kVarInt, 106u},
                  Token{7, ProtoWireType::kFixed32, 107u},
                  Token{8, ProtoWireType::kFixed32, 0x42420000u},
                  Token{9, ProtoWireType::kFixed64, uint64_t(1ull << 63)}));
}

TEST(MessageTokenizerTest, InvlidCases) {
  {
    // A very large varint.
    MessageTokenizer tokenizer;
    EXPECT_FALSE(tokenizer.Push(0x08).valid());
    for (int i = 0; i < 14; ++i)
      EXPECT_FALSE(tokenizer.Push(0xff).valid());
    EXPECT_FALSE(tokenizer.Push(0x0).valid());
    EXPECT_FALSE(tokenizer.idle());
    EXPECT_EQ(tokenizer.state(), 6u);
  }
  {
    // A very large string.
    MessageTokenizer tokenizer;
    EXPECT_FALSE(tokenizer.Push(0x0A).valid());
    EXPECT_FALSE(tokenizer.Push(0xFF).valid());
    EXPECT_FALSE(tokenizer.Push(0xFF).valid());
    EXPECT_FALSE(tokenizer.Push(0xFF).valid());
    EXPECT_FALSE(tokenizer.Push(0xFF).valid());
    EXPECT_FALSE(tokenizer.Push(0x20).valid());
    EXPECT_FALSE(tokenizer.idle());
    EXPECT_EQ(tokenizer.state(), 5u);
  }
  {
    // A field of unknown type (wire type = 0x3).
    MessageTokenizer tokenizer;
    EXPECT_FALSE(tokenizer.Push(0x0B).valid());
    EXPECT_FALSE(tokenizer.Push(0).valid());
    EXPECT_FALSE(tokenizer.Push(0).valid());
    EXPECT_FALSE(tokenizer.idle());
    EXPECT_EQ(tokenizer.state(), 4u);
  }
}

}  // namespace
}  // namespace protozero