// Copyright 2020 The Fuchsia 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 "phase_2_secure_connections.h"

#include <cstdint>
#include <memory>

#include <gtest/gtest.h>

#include "src/connectivity/bluetooth/core/bt-host/common/byte_buffer.h"
#include "src/connectivity/bluetooth/core/bt-host/common/device_address.h"
#include "src/connectivity/bluetooth/core/bt-host/common/random.h"
#include "src/connectivity/bluetooth/core/bt-host/common/test_helpers.h"
#include "src/connectivity/bluetooth/core/bt-host/common/uint128.h"
#include "src/connectivity/bluetooth/core/bt-host/hci-spec/constants.h"
#include "src/connectivity/bluetooth/core/bt-host/hci/connection.h"
#include "src/connectivity/bluetooth/core/bt-host/l2cap/fake_channel_test.h"
#include "src/connectivity/bluetooth/core/bt-host/l2cap/l2cap_defs.h"
#include "src/connectivity/bluetooth/core/bt-host/sm/ecdh_key.h"
#include "src/connectivity/bluetooth/core/bt-host/sm/fake_phase_listener.h"
#include "src/connectivity/bluetooth/core/bt-host/sm/packet.h"
#include "src/connectivity/bluetooth/core/bt-host/sm/smp.h"
#include "src/connectivity/bluetooth/core/bt-host/sm/types.h"
#include "src/connectivity/bluetooth/core/bt-host/sm/util.h"
#include "src/lib/fxl/memory/weak_ptr.h"

namespace bt::sm {
namespace {
using ConfirmCallback = FakeListener::ConfirmCallback;
using PasskeyResponseCallback = FakeListener::PasskeyResponseCallback;

// clang-format off
const PairingFeatures kDefaultFeatures(
    true,                                     // initiator
    true,                                     // secure_connections
    true,                                     // will_bond
    std::optional<CrossTransportKeyAlgo>{std::nullopt},
    PairingMethod::kJustWorks,
    kMaxEncryptionKeySize,                    // encryption_key_size
    KeyDistGen::kIdKey,                       // local_key_distribution
    KeyDistGen::kIdKey | KeyDistGen::kEncKey  // remote_key_distribution
);

const PairingRequestParams kDefaultPreq{
    .io_capability = IOCapability::kNoInputNoOutput,
    .oob_data_flag = OOBDataFlag::kNotPresent,
    .auth_req = AuthReq::kSC | AuthReq::kBondingFlag,
    .max_encryption_key_size = kMaxEncryptionKeySize,
    .initiator_key_dist_gen = KeyDistGen::kIdKey,
    .responder_key_dist_gen = KeyDistGen::kIdKey | KeyDistGen::kEncKey
};

const PairingResponseParams kDefaultPres{
    .io_capability = IOCapability::kNoInputNoOutput,
    .oob_data_flag = OOBDataFlag::kNotPresent,
    .auth_req = AuthReq::kSC | AuthReq::kBondingFlag,
    .max_encryption_key_size = kMaxEncryptionKeySize,
    .initiator_key_dist_gen = KeyDistGen::kIdKey,
    .responder_key_dist_gen = KeyDistGen::kIdKey | KeyDistGen::kEncKey
};
// clang-format on
const DeviceAddress kAddr1(DeviceAddress::Type::kLEPublic, {0x00, 0x00, 0x00, 0x00, 0x00, 0x01});
const DeviceAddress kAddr2(DeviceAddress::Type::kLEPublic, {0x00, 0x00, 0x00, 0x00, 0x00, 0x02});
const LocalEcdhKey kDefaultEcdhKey = LocalEcdhKey::Create().value();

using util::PacketSize;

class SMP_Phase2SecureConnectionsTest : public l2cap::testing::FakeChannelTest {
 public:
  SMP_Phase2SecureConnectionsTest() = default;
  ~SMP_Phase2SecureConnectionsTest() override = default;

 protected:
  void SetUp() override { NewPhase2SecureConnections(); }

  void TearDown() override { phase_2_sc_ = nullptr; }

  void NewPhase2SecureConnections(Role local_role = Role::kInitiator,
                                  PairingMethod method = PairingMethod::kJustWorks,
                                  uint16_t mtu = kLeSecureConnectionsMtu) {
    features_.initiator = (local_role == Role::kInitiator);
    features_.method = method;
    ChannelOptions options(l2cap::kLESMPChannelId, mtu);
    options.link_type = hci::Connection::LinkType::kLE;
    sm_chan_ = std::make_unique<PairingChannel>(CreateFakeChannel(options));

    listener_ = std::make_unique<FakeListener>();
    phase_2_sc_ = std::make_unique<Phase2SecureConnections>(
        sm_chan_->GetWeakPtr(), listener_->as_weak_ptr(), local_role, features_, preq_, pres_,
        initiator_addr_, responder_addr_, [this](const UInt128& ltk) {
          phase_2_complete_count_++;
          ltk_ = ltk;
        });
  }

  template <typename T>
  void ReceiveCmd(Code cmd_code, const T& value) {
    fake_chan()->Receive(MakeCmd(cmd_code, value));
  }

  template <typename T>
  StaticByteBuffer<PacketSize<T>()> MakeCmd(Code cmd_code, const T& value) {
    StaticByteBuffer<PacketSize<T>()> buffer;
    PacketWriter writer(cmd_code, &buffer);
    *writer.mutable_payload<T>() = value;
    return buffer;
  }

  static std::pair<Code, UInt128> ExtractCodeAnd128BitCmd(ByteBufferPtr sdu) {
    ZX_ASSERT_MSG(sdu, "Tried to ExtractCodeAnd128BitCmd from nullptr in test");
    auto maybe_reader = ValidPacketReader::ParseSdu(sdu);
    ZX_ASSERT_MSG(maybe_reader.is_ok(), "Tried to ExtractCodeAnd128BitCmd from invalid SMP packet");
    return {maybe_reader.value().code(), maybe_reader.value().payload<UInt128>()};
  }

  std::optional<PairingConfirmValue> FastForwardPublicKeyExchange() {
    std::optional<PairingConfirmValue> responder_jw_nc_confirm = std::nullopt;
    fake_chan()->SetSendCallback(
        [&responder_jw_nc_confirm, this](ByteBufferPtr sdu) {
          auto reader = ValidPacketReader::ParseSdu(sdu).value();
          if (reader.code() == kPairingPublicKey) {
            local_key_ = EcdhKey::ParseFromPublicKey(reader.payload<PairingPublicKeyParams>());
          } else if (reader.code() == kPairingConfirm) {
            // Confirm must come after ECDH Pub Key if it comes now
            ASSERT_TRUE(local_key_.has_value());
            responder_jw_nc_confirm = reader.payload<PairingConfirmValue>();
          } else {
            ADD_FAILURE() << "unexpected packet code " << reader.code();
          }
        },
        dispatcher());

    if (phase_2_sc_->role() == Role::kInitiator) {
      phase_2_sc_->Start();
      RunLoopUntilIdle();
      ZX_ASSERT_MSG(local_key_.has_value(), "initiator did not send ecdh key upon starting");
      ReceiveCmd(kPairingPublicKey, peer_key_.GetSerializedPublicKey());
      RunLoopUntilIdle();
    } else {
      phase_2_sc_->Start();
      ReceiveCmd(kPairingPublicKey, peer_key_.GetSerializedPublicKey());
      RunLoopUntilIdle();
      ZX_ASSERT_MSG(local_key_.has_value(), "responder did not send ecdh key upon peer key");
      if (features_.method == PairingMethod::kJustWorks ||
          features_.method == PairingMethod::kNumericComparison) {
        ZX_ASSERT(responder_jw_nc_confirm.has_value());
        return responder_jw_nc_confirm;
      }
    }
    // We should only send confirm value immediately after ECDH key as responder in Numeric
    // Comparison/Just Works pairing, in which case we would've already returned.
    ZX_ASSERT(!responder_jw_nc_confirm.has_value());
    return responder_jw_nc_confirm;
  }

  UInt128 GenerateConfirmValue(const UInt128& random, bool gen_initiator_confirm,
                               uint8_t r = 0) const {
    ZX_ASSERT_MSG(local_key_.has_value(), "cannot compute confirm, missing key!");
    UInt256 pka = local_key_->GetPublicKeyX(), pkb = peer_key_.GetPublicKeyX();
    if (phase_2_sc_->role() == Role::kResponder) {
      std::swap(pka, pkb);
    }
    return gen_initiator_confirm ? util::F4(pka, pkb, random, r).value()
                                 : util::F4(pkb, pka, random, r).value();
  }

  struct MatchingPair {
    UInt128 confirm;
    UInt128 random;
  };
  MatchingPair GenerateMatchingConfirmAndRandom(uint8_t r = 0) const {
    MatchingPair pair{.random = {1, 2, 3, 4}};
    pair.confirm = GenerateConfirmValue(pair.random, phase_2_sc_->role() == Role::kResponder, r);
    return pair;
  }

  struct LtkAndChecks {
    UInt128 ltk;
    UInt128 dhkey_check_a;
    UInt128 dhkey_check_b;
  };
  LtkAndChecks GenerateLtkAndChecks(const UInt128& initiator_rand, const UInt128& responder_rand,
                                    uint64_t r = 0) {
    LtkAndChecks vals;
    util::F5Results f5 = *util::F5(peer_key_.CalculateDhKey(*local_key_), initiator_rand,
                                   responder_rand, initiator_addr_, responder_addr_);
    vals.ltk = f5.ltk;

    UInt128 r_array{0};
    // Copy little-endian uint64 r to the UInt128 array needed for Stage 2
    std::memcpy(r_array.data(), &r, sizeof(uint64_t));
    vals.dhkey_check_a =
        *util::F6(f5.mac_key, initiator_rand, responder_rand, r_array, preq_.auth_req,
                  preq_.oob_data_flag, preq_.io_capability, initiator_addr_, responder_addr_);
    vals.dhkey_check_b =
        *util::F6(f5.mac_key, responder_rand, initiator_rand, r_array, pres_.auth_req,
                  pres_.oob_data_flag, pres_.io_capability, responder_addr_, initiator_addr_);
    return vals;
  }

  LtkAndChecks FastForwardToDhKeyCheck() {
    ZX_ASSERT_MSG(features_.method == PairingMethod::kJustWorks,
                  "Fast forward to DHKey check only implemented for JustWorks method");
    return phase_2_sc_->role() == Role::kInitiator ? FastForwardToDhKeyCheckInitiatorJustWorks()
                                                   : FastForwardToDhKeyCheckResponderJustWorks();
  }

  void DestroyPhase2() { phase_2_sc_.reset(nullptr); }
  Phase2SecureConnections* phase_2_sc() { return phase_2_sc_.get(); }
  FakeListener* listener() { return listener_.get(); }
  const LocalEcdhKey& peer_key() { return peer_key_; }
  const std::optional<EcdhKey>& local_key() { return local_key_; }
  std::optional<EcdhKey>& mut_local_key() { return local_key_; }
  int phase_2_complete_count() const { return phase_2_complete_count_; }
  UInt128 ltk() const { return ltk_; }

 private:
  LtkAndChecks FastForwardToDhKeyCheckInitiatorJustWorks() {
    FastForwardPublicKeyExchange();
    Code sent_code = kPairingFailed;
    PairingRandomValue initiator_rand;
    fake_chan()->SetSendCallback(
        [&](ByteBufferPtr sdu) {
          std::tie(sent_code, initiator_rand) = ExtractCodeAnd128BitCmd(std::move(sdu));
        },
        dispatcher());
    MatchingPair stage1_vals = GenerateMatchingConfirmAndRandom();
    ReceiveCmd<PairingConfirmValue>(kPairingConfirm, stage1_vals.confirm);
    RunLoopUntilIdle();
    ZX_ASSERT_MSG(kPairingRandom == sent_code, "did not send pairing random when expected!");

    ReceiveCmd(kPairingRandom, stage1_vals.random);
    RunLoopUntilIdle();
    return GenerateLtkAndChecks(initiator_rand, stage1_vals.random);
  }

  LtkAndChecks FastForwardToDhKeyCheckResponderJustWorks() {
    UInt128 rsp_confirm = *FastForwardPublicKeyExchange();
    Code sent_code = kPairingFailed;
    UInt128 rsp_rand;
    fake_chan()->SetSendCallback(
        [&](ByteBufferPtr sdu) {
          std::tie(sent_code, rsp_rand) = ExtractCodeAnd128BitCmd(std::move(sdu));
        },
        dispatcher());
    PairingRandomValue initiator_rand{1};
    ReceiveCmd(kPairingRandom, initiator_rand);
    RunLoopUntilIdle();

    ZX_ASSERT_MSG(kPairingRandom == sent_code, "did not send pairing random when expected!");
    ZX_ASSERT_MSG(GenerateConfirmValue(rsp_rand, false /*gen_init_confirm*/) == rsp_confirm,
                  "send invalid confirm value as JustWorks responder");
    return GenerateLtkAndChecks(initiator_rand, rsp_rand);
  }

  std::unique_ptr<FakeListener> listener_;
  std::unique_ptr<PairingChannel> sm_chan_;
  std::unique_ptr<Phase2SecureConnections> phase_2_sc_;
  // Key of the internal bt-host SMP stack
  std::optional<EcdhKey> local_key_;
  // Key of the unit test peer we're mocking
  const LocalEcdhKey& peer_key_ = kDefaultEcdhKey;

  PairingFeatures features_ = kDefaultFeatures;
  PairingRequestParams preq_ = kDefaultPreq;
  PairingResponseParams pres_ = kDefaultPres;
  DeviceAddress initiator_addr_ = kAddr1;
  DeviceAddress responder_addr_ = kAddr2;
  int phase_2_complete_count_ = 0;
  UInt128 ltk_;

  DISALLOW_COPY_AND_ASSIGN_ALLOW_MOVE(SMP_Phase2SecureConnectionsTest);
};

using SMP_Phase2SecureConnectionsDeathTest = SMP_Phase2SecureConnectionsTest;

TEST_F(SMP_Phase2SecureConnectionsDeathTest, MtuTooSmallDies) {
  ASSERT_DEATH_IF_SUPPORTED(NewPhase2SecureConnections(Role::kInitiator, PairingMethod::kJustWorks,
                                                       kNoSecureConnectionsMtu),
                            ".*SecureConnections.*");
}

TEST_F(SMP_Phase2SecureConnectionsTest, ReceivePairingFailed) {
  phase_2_sc()->Start();
  fake_chan()->Receive(
      StaticByteBuffer<PacketSize<ErrorCode>()>{kPairingFailed, ErrorCode::kPairingNotSupported});
  RunLoopUntilIdle();

  ASSERT_TRUE(listener()->last_error().is_protocol_error());
  EXPECT_EQ(ErrorCode::kPairingNotSupported, listener()->last_error().protocol_error());
}

TEST_F(SMP_Phase2SecureConnectionsTest, ReceiveMalformedPacket) {
  phase_2_sc()->Start();
  // PairingPublicKeyParams is expected to have both an X and Y value, not just an X.
  const UInt256 kX = peer_key().GetPublicKeyX();
  const auto kPairingPublicKeyCmd = MakeCmd(kPairingPublicKey, kX);
  const StaticByteBuffer<PacketSize<ErrorCode>()> kExpectedFailure{kPairingFailed,
                                                                   ErrorCode::kInvalidParameters};

  EXPECT_TRUE(ReceiveAndExpect(kPairingPublicKeyCmd, kExpectedFailure));
}

TEST_F(SMP_Phase2SecureConnectionsTest, ReceiveUnexpectedPacket) {
  phase_2_sc()->Start();
  // Pairing Responses should only be sent during Phase 1 of pairing.
  const auto kPairingResponseCmd = MakeCmd(kPairingResponse, PairingResponseParams());
  const StaticByteBuffer<PacketSize<ErrorCode>()> kExpectedFailure{kPairingFailed,
                                                                   ErrorCode::kUnspecifiedReason};

  EXPECT_TRUE(ReceiveAndExpect(kPairingResponseCmd, kExpectedFailure));
}

TEST_F(SMP_Phase2SecureConnectionsTest, InitiatorPubKeyOutOfOrder) {
  NewPhase2SecureConnections(Role::kInitiator);

  const auto kPairingPublicKeyCmd = MakeCmd(kPairingPublicKey, peer_key().GetSerializedPublicKey());
  const StaticByteBuffer<PacketSize<ErrorCode>()> kExpectedFailure{kPairingFailed,
                                                                   ErrorCode::kUnspecifiedReason};
  ASSERT_TRUE(ReceiveAndExpect(kPairingPublicKeyCmd, kExpectedFailure));
  ASSERT_EQ(1, listener()->pairing_error_count());
}

TEST_F(SMP_Phase2SecureConnectionsTest, RejectsPublicKeyOffCurve) {
  phase_2_sc()->Start();
  const auto kPairingPublicKeyCmd =
      MakeCmd(kPairingPublicKey, PairingPublicKeyParams{.x = {0x01}, .y = {0x02}});
  const StaticByteBuffer<PacketSize<ErrorCode>()> kExpectedFailure{kPairingFailed,
                                                                   ErrorCode::kInvalidParameters};
  ASSERT_TRUE(ReceiveAndExpect(kPairingPublicKeyCmd, kExpectedFailure));
  ASSERT_EQ(1, listener()->pairing_error_count());
}

TEST_F(SMP_Phase2SecureConnectionsTest, ReceivePeerPublicKeyTwice) {
  phase_2_sc()->Start();
  const auto kPairingPublicKeyCmd = MakeCmd(kPairingPublicKey, peer_key().GetSerializedPublicKey());
  fake_chan()->Receive(kPairingPublicKeyCmd);
  const StaticByteBuffer<PacketSize<ErrorCode>()> kExpectedFailure{kPairingFailed,
                                                                   ErrorCode::kUnspecifiedReason};
  ASSERT_TRUE(ReceiveAndExpect(kPairingPublicKeyCmd, kExpectedFailure));
  ASSERT_EQ(1, listener()->pairing_error_count());
}

TEST_F(SMP_Phase2SecureConnectionsTest, ReceiveConfirmValueBeforeStage1Fails) {
  phase_2_sc()->Start();
  const auto kPairingConfirmCmd = MakeCmd(kPairingConfirm, PairingConfirmValue{1});
  const StaticByteBuffer<PacketSize<ErrorCode>()> kExpectedFailure{kPairingFailed,
                                                                   ErrorCode::kUnspecifiedReason};
  ASSERT_TRUE(ReceiveAndExpect(kPairingConfirmCmd, kExpectedFailure));
  ASSERT_EQ(1, listener()->pairing_error_count());
}

TEST_F(SMP_Phase2SecureConnectionsTest, ReceiveRandomValueBeforeStage1Fails) {
  phase_2_sc()->Start();
  const auto kPairingRandomCmd = MakeCmd(kPairingRandom, PairingRandomValue{1});
  const StaticByteBuffer<PacketSize<ErrorCode>()> kExpectedFailure{kPairingFailed,
                                                                   ErrorCode::kUnspecifiedReason};
  ASSERT_TRUE(ReceiveAndExpect(kPairingRandomCmd, kExpectedFailure));
  ASSERT_EQ(1, listener()->pairing_error_count());
}

TEST_F(SMP_Phase2SecureConnectionsTest, ReceiveDhKeyCheckValueBeforeStage1Fails) {
  phase_2_sc()->Start();
  const auto kPairingDHKeyCheckCmd = MakeCmd(kPairingDHKeyCheck, PairingDHKeyCheckValueE{1});
  const StaticByteBuffer<PacketSize<ErrorCode>()> kExpectedFailure{kPairingFailed,
                                                                   ErrorCode::kUnspecifiedReason};
  ASSERT_TRUE(ReceiveAndExpect(kPairingDHKeyCheckCmd, kExpectedFailure));
  ASSERT_EQ(1, listener()->pairing_error_count());
}

TEST_F(SMP_Phase2SecureConnectionsTest, ReceiveConfirmValueAfterStage1Fails) {
  FastForwardToDhKeyCheck();
  const auto kPairingConfirmCmd = MakeCmd(kPairingConfirm, PairingConfirmValue{1});
  const StaticByteBuffer<PacketSize<ErrorCode>()> kExpectedFailure{kPairingFailed,
                                                                   ErrorCode::kUnspecifiedReason};
  ASSERT_TRUE(ReceiveAndExpect(kPairingConfirmCmd, kExpectedFailure));
  ASSERT_EQ(1, listener()->pairing_error_count());
}

TEST_F(SMP_Phase2SecureConnectionsTest, ReceiveRandomValueAfterStage1Fails) {
  FastForwardToDhKeyCheck();
  const auto kPairingRandomCmd = MakeCmd(kPairingRandom, PairingRandomValue{1});
  const StaticByteBuffer<PacketSize<ErrorCode>()> kExpectedFailure{kPairingFailed,
                                                                   ErrorCode::kUnspecifiedReason};
  ASSERT_TRUE(ReceiveAndExpect(kPairingRandomCmd, kExpectedFailure));
  ASSERT_EQ(1, listener()->pairing_error_count());
}

TEST_F(SMP_Phase2SecureConnectionsTest, InitiatorReceiveDhKeyCheckWhileWaitingForConfirmFails) {
  NewPhase2SecureConnections(Role::kInitiator);
  ConfirmCallback confirm_cb = nullptr;
  listener()->set_confirm_delegate([&](ConfirmCallback cb) { confirm_cb = std::move(cb); });
  LtkAndChecks expected_stage2_vals = FastForwardToDhKeyCheck();
  ASSERT_TRUE(confirm_cb);
  // Receiving the peer DHKey check before user confirmation should fail as initiator.
  const auto kPairingDHKeyCheckCmd =
      MakeCmd(kPairingDHKeyCheck, expected_stage2_vals.dhkey_check_b);
  const StaticByteBuffer<PacketSize<ErrorCode>()> kExpectedFailure{kPairingFailed,
                                                                   ErrorCode::kUnspecifiedReason};
  ASSERT_TRUE(ReceiveAndExpect(kPairingDHKeyCheckCmd, kExpectedFailure));
  ASSERT_EQ(1, listener()->pairing_error_count());
}

TEST_F(SMP_Phase2SecureConnectionsTest, Stage1JustWorksErrorPropagates) {
  NewPhase2SecureConnections(Role::kInitiator, PairingMethod::kJustWorks);
  ConfirmCallback confirm_cb = nullptr;
  listener()->set_confirm_delegate([&](ConfirmCallback cb) { confirm_cb = std::move(cb); });
  FastForwardPublicKeyExchange();

  Code sent_code = kPairingFailed;
  fake_chan()->SetSendCallback(
      [&](ByteBufferPtr sdu) {
        std::tie(sent_code, std::ignore) = ExtractCodeAnd128BitCmd(std::move(sdu));
      },
      dispatcher());
  MatchingPair stage1_vals = GenerateMatchingConfirmAndRandom();
  ReceiveCmd<PairingConfirmValue>(kPairingConfirm, stage1_vals.confirm);
  RunLoopUntilIdle();
  ASSERT_EQ(kPairingRandom, sent_code);
  UInt128 mismatched_random = stage1_vals.random;
  mismatched_random[0] += 1;
  const auto kMismatchedPairingRandomCmd = MakeCmd(kPairingRandom, mismatched_random);
  const StaticByteBuffer<PacketSize<ErrorCode>()> kExpectedFailure{kPairingFailed,
                                                                   ErrorCode::kConfirmValueFailed};
  ASSERT_TRUE(ReceiveAndExpect(kMismatchedPairingRandomCmd, kExpectedFailure));
  ASSERT_EQ(1, listener()->pairing_error_count());
}

TEST_F(SMP_Phase2SecureConnectionsTest, Stage1PasskeyErrorPropagates) {
  NewPhase2SecureConnections(Role::kInitiator, PairingMethod::kPasskeyEntryDisplay);
  uint32_t passkey;
  ConfirmCallback confirm_cb = nullptr;
  listener()->set_display_delegate(
      [&](uint32_t disp_passkey, Delegate::DisplayMethod method, ConfirmCallback cb) {
        ASSERT_EQ(Delegate::DisplayMethod::kPeerEntry, method);
        confirm_cb = std::move(cb);
        passkey = disp_passkey;
      });
  FastForwardPublicKeyExchange();

  ASSERT_TRUE(confirm_cb);
  const StaticByteBuffer<PacketSize<ErrorCode>()> kExpectedFailure{kPairingFailed,
                                                                   ErrorCode::kPasskeyEntryFailed};
  bool failure_sent = false;
  fake_chan()->SetSendCallback(
      [&](ByteBufferPtr sdu) {
        ASSERT_TRUE(ContainersEqual(kExpectedFailure, *sdu));
        failure_sent = true;
      },
      dispatcher());
  confirm_cb(false);
  RunLoopUntilIdle();
  ASSERT_EQ(1, listener()->pairing_error_count());
  ASSERT_TRUE(failure_sent);
}

TEST_F(SMP_Phase2SecureConnectionsTest, InitiatorReceiveWrongDhKeyCheckFails) {
  NewPhase2SecureConnections(Role::kInitiator);
  ConfirmCallback confirm_cb = nullptr;
  listener()->set_confirm_delegate([&](ConfirmCallback cb) { confirm_cb = std::move(cb); });
  LtkAndChecks expected_stage2_vals = FastForwardToDhKeyCheck();
  ASSERT_TRUE(confirm_cb);
  Code sent_code = kPairingFailed;
  UInt128 sent_payload;
  fake_chan()->SetSendCallback(
      [&](ByteBufferPtr sdu) {
        std::tie(sent_code, sent_payload) = ExtractCodeAnd128BitCmd(std::move(sdu));
      },
      dispatcher());
  confirm_cb(true);
  RunLoopUntilIdle();
  ASSERT_EQ(kPairingDHKeyCheck, sent_code);
  // As initiator, we expect the dhkey_check_b value, not the a.
  const auto kPairingDHKeyCheckCmd =
      MakeCmd(kPairingDHKeyCheck, expected_stage2_vals.dhkey_check_a);
  const StaticByteBuffer<PacketSize<ErrorCode>()> kExpectedFailure{kPairingFailed,
                                                                   ErrorCode::kDHKeyCheckFailed};
  ASSERT_TRUE(ReceiveAndExpect(kPairingDHKeyCheckCmd, kExpectedFailure));
  ASSERT_EQ(1, listener()->pairing_error_count());
  ASSERT_EQ(0, phase_2_complete_count());
}

TEST_F(SMP_Phase2SecureConnectionsTest, InitiatorFlowSuccessJustWorks) {
  NewPhase2SecureConnections(Role::kInitiator, PairingMethod::kJustWorks);
  ConfirmCallback confirm_cb = nullptr;
  listener()->set_confirm_delegate([&](ConfirmCallback cb) { confirm_cb = std::move(cb); });
  FastForwardPublicKeyExchange();

  Code sent_code = kPairingFailed;
  UInt128 sent_payload;
  fake_chan()->SetSendCallback(
      [&](ByteBufferPtr sdu) {
        std::tie(sent_code, sent_payload) = ExtractCodeAnd128BitCmd(std::move(sdu));
      },
      dispatcher());
  MatchingPair stage1_vals = GenerateMatchingConfirmAndRandom();
  ReceiveCmd<PairingConfirmValue>(kPairingConfirm, stage1_vals.confirm);
  RunLoopUntilIdle();
  ASSERT_EQ(kPairingRandom, sent_code);
  UInt128 local_rand = sent_payload;
  ASSERT_FALSE(confirm_cb);
  ReceiveCmd(kPairingRandom, stage1_vals.random);
  RunLoopUntilIdle();
  ASSERT_TRUE(confirm_cb);

  LtkAndChecks expected_stage2_vals = GenerateLtkAndChecks(local_rand, stage1_vals.random);
  confirm_cb(true);
  RunLoopUntilIdle();
  // After receiving user confirmation, we should send (the correct) DHKey Check Ea
  ASSERT_EQ(kPairingDHKeyCheck, sent_code);
  ASSERT_EQ(expected_stage2_vals.dhkey_check_a, sent_payload);
  ReceiveCmd(kPairingDHKeyCheck, expected_stage2_vals.dhkey_check_b);
  RunLoopUntilIdle();
  ASSERT_EQ(1, phase_2_complete_count());
  // We should generate the same LTK on "both sides"
  ASSERT_EQ(expected_stage2_vals.ltk, ltk());
}

TEST_F(SMP_Phase2SecureConnectionsTest, InitiatorFlowSuccessNumericComparison) {
  NewPhase2SecureConnections(Role::kInitiator, PairingMethod::kNumericComparison);
  ConfirmCallback confirm_cb = nullptr;
  listener()->set_display_delegate(
      [&](uint32_t, Delegate::DisplayMethod method, ConfirmCallback cb) {
        ASSERT_EQ(Delegate::DisplayMethod::kComparison, method);
        confirm_cb = std::move(cb);
      });
  FastForwardPublicKeyExchange();

  Code sent_code = kPairingFailed;
  UInt128 sent_payload;
  fake_chan()->SetSendCallback(
      [&](ByteBufferPtr sdu) {
        std::tie(sent_code, sent_payload) = ExtractCodeAnd128BitCmd(std::move(sdu));
      },
      dispatcher());
  MatchingPair stage1_vals = GenerateMatchingConfirmAndRandom();
  ReceiveCmd(kPairingConfirm, stage1_vals.confirm);
  RunLoopUntilIdle();
  EXPECT_EQ(kPairingRandom, sent_code);
  UInt128 initiator_random = sent_payload;
  ASSERT_FALSE(confirm_cb);

  ReceiveCmd(kPairingRandom, stage1_vals.random);
  RunLoopUntilIdle();
  ASSERT_TRUE(confirm_cb);
  confirm_cb(true);
  RunLoopUntilIdle();
  LtkAndChecks vals = GenerateLtkAndChecks(initiator_random, stage1_vals.random);
  EXPECT_EQ(kPairingDHKeyCheck, sent_code);
  EXPECT_EQ(vals.dhkey_check_a, sent_payload);
  ReceiveCmd(kPairingDHKeyCheck, vals.dhkey_check_b);
  RunLoopUntilIdle();
  ASSERT_EQ(vals.ltk, ltk());
}

TEST_F(SMP_Phase2SecureConnectionsTest, InitiatorFlowSuccessPasskeyEntryDisplay) {
  NewPhase2SecureConnections(Role::kInitiator, PairingMethod::kPasskeyEntryDisplay);
  uint32_t passkey;
  ConfirmCallback confirm_cb = nullptr;
  listener()->set_display_delegate(
      [&](uint32_t disp_passkey, Delegate::DisplayMethod method, ConfirmCallback cb) {
        ASSERT_EQ(Delegate::DisplayMethod::kPeerEntry, method);
        confirm_cb = std::move(cb);
        passkey = disp_passkey;
      });
  FastForwardPublicKeyExchange();
  ASSERT_TRUE(confirm_cb);
  Code sent_code = kPairingFailed;
  UInt128 sent_payload;
  fake_chan()->SetSendCallback(
      [&](ByteBufferPtr sdu) {
        std::tie(sent_code, sent_payload) = ExtractCodeAnd128BitCmd(std::move(sdu));
      },
      dispatcher());
  confirm_cb(true);
  MatchingPair stage1_vals;
  UInt128 last_sent_rand;
  for (size_t i = 0; i < 20; ++i) {
    const uint8_t r = (passkey & (1 << i)) ? 0x81 : 0x80;
    stage1_vals = GenerateMatchingConfirmAndRandom(r);
    RunLoopUntilIdle();

    ASSERT_EQ(kPairingConfirm, sent_code);
    PairingConfirmValue init_confirm = sent_payload;
    ReceiveCmd(kPairingConfirm, stage1_vals.confirm);
    RunLoopUntilIdle();

    ASSERT_EQ(kPairingRandom, sent_code);
    last_sent_rand = sent_payload;
    EXPECT_EQ(GenerateConfirmValue(sent_payload, true /*gen_initiator_confirm*/, r), init_confirm);
    ReceiveCmd(kPairingRandom, stage1_vals.random);
  }
  LtkAndChecks vals = GenerateLtkAndChecks(last_sent_rand, stage1_vals.random, uint64_t{passkey});
  RunLoopUntilIdle();
  EXPECT_EQ(kPairingDHKeyCheck, sent_code);
  EXPECT_EQ(vals.dhkey_check_a, sent_payload);
  ReceiveCmd(kPairingDHKeyCheck, vals.dhkey_check_b);
  RunLoopUntilIdle();
  ASSERT_EQ(vals.ltk, ltk());
}

TEST_F(SMP_Phase2SecureConnectionsTest, InitiatorFlowSuccessPasskeyEntryInput) {
  NewPhase2SecureConnections(Role::kInitiator, PairingMethod::kPasskeyEntryInput);
  PasskeyResponseCallback passkey_cb = nullptr;
  listener()->set_request_passkey_delegate(
      [&](PasskeyResponseCallback cb) { passkey_cb = std::move(cb); });
  FastForwardPublicKeyExchange();
  ASSERT_TRUE(passkey_cb);
  Code sent_code = kPairingFailed;
  UInt128 sent_payload;
  fake_chan()->SetSendCallback(
      [&](ByteBufferPtr sdu) {
        std::tie(sent_code, sent_payload) = ExtractCodeAnd128BitCmd(std::move(sdu));
      },
      dispatcher());
  const int64_t passkey = 123456;
  passkey_cb(passkey);
  MatchingPair stage1_vals;
  UInt128 last_sent_rand;
  for (size_t i = 0; i < 20; ++i) {
    const uint8_t r = (passkey & (1 << i)) ? 0x81 : 0x80;
    stage1_vals = GenerateMatchingConfirmAndRandom(r);
    RunLoopUntilIdle();

    ASSERT_EQ(kPairingConfirm, sent_code);
    PairingConfirmValue init_confirm = sent_payload;
    ReceiveCmd(kPairingConfirm, stage1_vals.confirm);
    RunLoopUntilIdle();

    ASSERT_EQ(kPairingRandom, sent_code);
    last_sent_rand = sent_payload;
    EXPECT_EQ(GenerateConfirmValue(sent_payload, true /*gen_initiator_confirm*/, r), init_confirm);
    ReceiveCmd(kPairingRandom, stage1_vals.random);
  }
  LtkAndChecks vals = GenerateLtkAndChecks(last_sent_rand, stage1_vals.random, uint64_t{passkey});
  RunLoopUntilIdle();
  EXPECT_EQ(kPairingDHKeyCheck, sent_code);
  EXPECT_EQ(vals.dhkey_check_a, sent_payload);
  ReceiveCmd(kPairingDHKeyCheck, vals.dhkey_check_b);
  RunLoopUntilIdle();
  ASSERT_EQ(vals.ltk, ltk());
}

TEST_F(SMP_Phase2SecureConnectionsTest, ResponderReceiveWrongDhKeyCheckFails) {
  NewPhase2SecureConnections(Role::kResponder);
  ConfirmCallback confirm_cb = nullptr;
  listener()->set_confirm_delegate([&](ConfirmCallback cb) { confirm_cb = std::move(cb); });
  LtkAndChecks expected_stage2_vals = FastForwardToDhKeyCheck();
  ASSERT_TRUE(confirm_cb);
  Code sent_code = kPairingFailed;
  UInt128 sent_payload;
  fake_chan()->SetSendCallback(
      [&](ByteBufferPtr sdu) {
        std::tie(sent_code, sent_payload) = ExtractCodeAnd128BitCmd(std::move(sdu));
      },
      dispatcher());
  confirm_cb(true);
  RunLoopUntilIdle();
  // As responder, we expect the dhkey_check_a value, not the b.
  const auto kPairingDHKeyCheckCmdWithBValue =
      MakeCmd(kPairingDHKeyCheck, expected_stage2_vals.dhkey_check_b);
  const StaticByteBuffer<PacketSize<ErrorCode>()> kExpectedFailure{kPairingFailed,
                                                                   ErrorCode::kDHKeyCheckFailed};
  ASSERT_TRUE(ReceiveAndExpect(kPairingDHKeyCheckCmdWithBValue, kExpectedFailure));
  ASSERT_EQ(1, listener()->pairing_error_count());
  ASSERT_EQ(0, phase_2_complete_count());
}

TEST_F(SMP_Phase2SecureConnectionsTest, ResponderFlowSuccessJustWorks) {
  NewPhase2SecureConnections(Role::kResponder, PairingMethod::kJustWorks);
  ConfirmCallback confirm_cb = nullptr;
  listener()->set_confirm_delegate([&](ConfirmCallback cb) { confirm_cb = std::move(cb); });
  PairingConfirmValue responder_confirm = *FastForwardPublicKeyExchange();

  Code sent_code = kPairingFailed;
  UInt128 sent_payload;
  fake_chan()->SetSendCallback(
      [&](ByteBufferPtr sdu) {
        std::tie(sent_code, sent_payload) = ExtractCodeAnd128BitCmd(std::move(sdu));
      },
      dispatcher());
  ASSERT_FALSE(confirm_cb);
  const PairingRandomValue kInitiatorRand{1};
  ReceiveCmd(kPairingRandom, kInitiatorRand);
  RunLoopUntilIdle();

  ASSERT_EQ(kPairingRandom, sent_code);
  UInt128 responder_rand = sent_payload;
  ASSERT_EQ(GenerateConfirmValue(responder_rand, false /*gen_init_confirm*/), responder_confirm);
  ASSERT_TRUE(confirm_cb);
  LtkAndChecks expected_stage2_vals = GenerateLtkAndChecks(kInitiatorRand, responder_rand);
  // After receiving user confirmation & the peer dhkey check, we should send the DHKey Check Eb.
  confirm_cb(true);
  ReceiveCmd(kPairingDHKeyCheck, expected_stage2_vals.dhkey_check_a);
  RunLoopUntilIdle();
  ASSERT_EQ(kPairingDHKeyCheck, sent_code);
  ASSERT_EQ(expected_stage2_vals.dhkey_check_b, sent_payload);
  ASSERT_EQ(1, phase_2_complete_count());
  // We should generate the same LTK on "both sides"
  ASSERT_EQ(expected_stage2_vals.ltk, ltk());
}

TEST_F(SMP_Phase2SecureConnectionsTest, ResponderReceiveDhKeyCheckWhileWaitingForConfirmSuccess) {
  NewPhase2SecureConnections(Role::kResponder, PairingMethod::kJustWorks);
  ConfirmCallback confirm_cb = nullptr;
  listener()->set_confirm_delegate([&](ConfirmCallback cb) { confirm_cb = std::move(cb); });
  LtkAndChecks expected_stage2_vals = FastForwardToDhKeyCheck();

  Code sent_code;
  PairingDHKeyCheckValueE sent_dhkey_check;
  fake_chan()->SetSendCallback(
      [&](ByteBufferPtr sdu) {
        std::tie(sent_code, sent_dhkey_check) = ExtractCodeAnd128BitCmd(std::move(sdu));
      },
      dispatcher());
  ASSERT_TRUE(confirm_cb);
  // Receiving the peer DHKey check before user confirmation should work as responder.
  ReceiveCmd(kPairingDHKeyCheck, expected_stage2_vals.dhkey_check_a);
  RunLoopUntilIdle();
  confirm_cb(true);
  RunLoopUntilIdle();

  ASSERT_EQ(kPairingDHKeyCheck, sent_code);
  ASSERT_EQ(expected_stage2_vals.dhkey_check_b, sent_dhkey_check);
  ASSERT_EQ(1, phase_2_complete_count());
}

}  // namespace
}  // namespace bt::sm