// Copyright 2018 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 "device/fido/u2f_sign_operation.h"

#include <string>
#include <utility>

#include "base/test/scoped_task_environment.h"
#include "crypto/ec_private_key.h"
#include "device/fido/authenticator_get_assertion_response.h"
#include "device/fido/ctap_get_assertion_request.h"
#include "device/fido/fido_constants.h"
#include "device/fido/fido_parsing_utils.h"
#include "device/fido/fido_test_data.h"
#include "device/fido/mock_fido_device.h"
#include "device/fido/test_callback_receiver.h"
#include "device/fido/virtual_u2f_device.h"
#include "testing/gmock/include/gmock/gmock.h"
#include "testing/gtest/include/gtest/gtest.h"

namespace device {

using ::testing::_;
using ::testing::InSequence;

namespace {

using TestSignCallback = ::device::test::StatusAndValueCallbackReceiver<
    CtapDeviceResponseCode,
    base::Optional<AuthenticatorGetAssertionResponse>>;

}  // namespace

class U2fSignOperationTest : public ::testing::Test {
 public:
  CtapGetAssertionRequest CreateSignRequest(
      std::vector<std::vector<uint8_t>> key_handles) {
    CtapGetAssertionRequest request(test_data::kRelyingPartyId,
                                    test_data::kClientDataHash);

    std::vector<PublicKeyCredentialDescriptor> allowed_list;
    for (auto& key_handle : key_handles) {
      allowed_list.emplace_back(CredentialType::kPublicKey,
                                std::move(key_handle));
    }

    request.SetAllowList(std::move(allowed_list));
    return request;
  }

  TestSignCallback& sign_callback_receiver() { return sign_callback_receiver_; }

 protected:
  base::test::ScopedTaskEnvironment scoped_task_environment_;
  TestSignCallback sign_callback_receiver_;
};

TEST_F(U2fSignOperationTest, SignSuccess) {
  auto request = CreateSignRequest(
      {fido_parsing_utils::Materialize(test_data::kU2fSignKeyHandle)});

  auto device = std::make_unique<MockFidoDevice>();
  EXPECT_CALL(*device, GetId()).WillRepeatedly(testing::Return("device"));
  InSequence s;
  device->ExpectRequestAndRespondWith(
      test_data::kU2fCheckOnlySignCommandApdu,
      test_data::kApduEncodedNoErrorSignResponse);
  device->ExpectRequestAndRespondWith(
      test_data::kU2fSignCommandApdu,
      test_data::kApduEncodedNoErrorSignResponse);

  auto u2f_sign = std::make_unique<U2fSignOperation>(
      device.get(), std::move(request), sign_callback_receiver().callback());
  u2f_sign->Start();

  sign_callback_receiver().WaitForCallback();
  EXPECT_EQ(CtapDeviceResponseCode::kSuccess,
            sign_callback_receiver().status());
  EXPECT_THAT(sign_callback_receiver().value()->signature(),
              ::testing::ElementsAreArray(test_data::kU2fSignature));
  EXPECT_THAT(sign_callback_receiver().value()->raw_credential_id(),
              ::testing::ElementsAreArray(test_data::kU2fSignKeyHandle));
}

TEST_F(U2fSignOperationTest, SignSuccessWithFakeDevice) {
  auto private_key = crypto::ECPrivateKey::Create();
  std::string public_key;
  private_key->ExportRawPublicKey(&public_key);
  auto hash = fido_parsing_utils::CreateSHA256Hash(public_key);
  std::vector<uint8_t> key_handle(hash.begin(), hash.end());
  auto request = CreateSignRequest({key_handle});

  auto device = std::make_unique<VirtualU2fDevice>();
  device->mutable_state()->registrations.emplace(
      key_handle, VirtualFidoDevice::RegistrationData(
                      std::move(private_key), test_data::kApplicationParameter,
                      42 /* counter */));

  auto u2f_sign = std::make_unique<U2fSignOperation>(
      device.get(), std::move(request), sign_callback_receiver().callback());
  u2f_sign->Start();

  sign_callback_receiver().WaitForCallback();
  EXPECT_EQ(CtapDeviceResponseCode::kSuccess,
            sign_callback_receiver().status());

  // Just a sanity check, we don't verify the actual signature.
  ASSERT_GE(32u + 1u + 4u + 8u,  // Minimal ECDSA signature is 8 bytes
            sign_callback_receiver()
                .value()
                ->auth_data()
                .SerializeToByteArray()
                .size());
  EXPECT_EQ(0x01,
            sign_callback_receiver()
                .value()
                ->auth_data()
                .SerializeToByteArray()[32]);  // UP flag
  // Counter is incremented for every sign request. Since we conduct 2 sign
  // requests (check only sign followed by a regular sign), counter should be
  // incremented twice.
  EXPECT_EQ(44, sign_callback_receiver()
                    .value()
                    ->auth_data()
                    .SerializeToByteArray()[36]);  // counter
}

TEST_F(U2fSignOperationTest, DelayedSuccess) {
  auto request = CreateSignRequest(
      {fido_parsing_utils::Materialize(test_data::kU2fSignKeyHandle)});

  // Simulates a device that times out waiting for user touch once before
  // responding successfully.
  auto device = std::make_unique<MockFidoDevice>();
  EXPECT_CALL(*device, GetId()).WillRepeatedly(::testing::Return("device"));

  InSequence s;
  device->ExpectRequestAndRespondWith(
      test_data::kU2fCheckOnlySignCommandApdu,
      test_data::kApduEncodedNoErrorSignResponse);
  device->ExpectRequestAndRespondWith(
      test_data::kU2fSignCommandApdu,
      test_data::kU2fConditionNotSatisfiedApduResponse);
  device->ExpectRequestAndRespondWith(
      test_data::kU2fSignCommandApdu,
      test_data::kApduEncodedNoErrorSignResponse);

  auto u2f_sign = std::make_unique<U2fSignOperation>(
      device.get(), std::move(request), sign_callback_receiver().callback());
  u2f_sign->Start();

  sign_callback_receiver().WaitForCallback();
  EXPECT_EQ(CtapDeviceResponseCode::kSuccess,
            sign_callback_receiver().status());
  EXPECT_THAT(sign_callback_receiver().value()->signature(),
              ::testing::ElementsAreArray(test_data::kU2fSignature));
  EXPECT_THAT(sign_callback_receiver().value()->raw_credential_id(),
              ::testing::ElementsAreArray(test_data::kU2fSignKeyHandle));
}

TEST_F(U2fSignOperationTest, MultipleHandles) {
  // Two wrong keys followed by a correct key ensuring the wrong keys will be
  // tested first.
  auto request = CreateSignRequest(
      {fido_parsing_utils::Materialize(test_data::kKeyHandleAlpha),
       fido_parsing_utils::Materialize(test_data::kKeyHandleBeta),
       fido_parsing_utils::Materialize(test_data::kU2fSignKeyHandle)});

  auto device = std::make_unique<MockFidoDevice>();
  EXPECT_CALL(*device, GetId()).WillRepeatedly(::testing::Return("device"));
  InSequence s;
  // Wrong key would respond with SW_WRONG_DATA.
  device->ExpectRequestAndRespondWith(
      test_data::kU2fCheckOnlySignCommandApduWithKeyAlpha,
      test_data::kU2fWrongDataApduResponse);
  device->ExpectRequestAndRespondWith(
      test_data::kU2fCheckOnlySignCommandApduWithKeyBeta,
      test_data::kU2fWrongDataApduResponse);
  device->ExpectRequestAndRespondWith(
      test_data::kU2fCheckOnlySignCommandApdu,
      test_data::kApduEncodedNoErrorSignResponse);
  device->ExpectRequestAndRespondWith(
      test_data::kU2fSignCommandApdu,
      test_data::kApduEncodedNoErrorSignResponse);

  auto u2f_sign = std::make_unique<U2fSignOperation>(
      device.get(), std::move(request), sign_callback_receiver().callback());
  u2f_sign->Start();

  sign_callback_receiver().WaitForCallback();
  EXPECT_EQ(CtapDeviceResponseCode::kSuccess,
            sign_callback_receiver().status());
  EXPECT_THAT(sign_callback_receiver().value()->signature(),
              ::testing::ElementsAreArray(test_data::kU2fSignature));
  EXPECT_THAT(sign_callback_receiver().value()->raw_credential_id(),
              ::testing::ElementsAreArray(test_data::kU2fSignKeyHandle));
}

// Test that Fake U2F registration is invoked when no credentials in the allowed
// list are recognized by the device.
TEST_F(U2fSignOperationTest, FakeEnroll) {
  auto request = CreateSignRequest(
      {fido_parsing_utils::Materialize(test_data::kKeyHandleAlpha),
       fido_parsing_utils::Materialize(test_data::kKeyHandleBeta)});

  auto device = std::make_unique<MockFidoDevice>();
  InSequence s;
  device->ExpectRequestAndRespondWith(
      test_data::kU2fCheckOnlySignCommandApduWithKeyAlpha,
      test_data::kU2fWrongDataApduResponse);
  device->ExpectRequestAndRespondWith(
      test_data::kU2fCheckOnlySignCommandApduWithKeyBeta,
      test_data::kU2fWrongDataApduResponse);
  device->ExpectRequestAndRespondWith(
      test_data::kU2fFakeRegisterCommand,
      test_data::kApduEncodedNoErrorRegisterResponse);

  auto u2f_sign = std::make_unique<U2fSignOperation>(
      device.get(), std::move(request), sign_callback_receiver().callback());
  u2f_sign->Start();
  sign_callback_receiver().WaitForCallback();

  EXPECT_EQ(CtapDeviceResponseCode::kCtap2ErrCredentialNotValid,
            sign_callback_receiver().status());
  EXPECT_FALSE(sign_callback_receiver().value());
}

// Tests that U2F fake enrollment should be re-tried repeatedly if no
// credentials are valid for the authenticator and user presence is not
// obtained.
TEST_F(U2fSignOperationTest, DelayedFakeEnrollment) {
  auto request = CreateSignRequest(
      {fido_parsing_utils::Materialize(test_data::kU2fSignKeyHandle)});

  // Simulates a device that times out waiting for user presence during fake
  // enrollment.
  auto device = std::make_unique<MockFidoDevice>();
  EXPECT_CALL(*device, GetId()).WillRepeatedly(::testing::Return("device0"));
  InSequence s;
  device->ExpectRequestAndRespondWith(test_data::kU2fCheckOnlySignCommandApdu,
                                      test_data::kU2fWrongDataApduResponse);
  device->ExpectRequestAndRespondWith(
      test_data::kU2fFakeRegisterCommand,
      test_data::kU2fConditionNotSatisfiedApduResponse);
  device->ExpectRequestAndRespondWith(
      test_data::kU2fFakeRegisterCommand,
      test_data::kApduEncodedNoErrorRegisterResponse);

  auto u2f_sign = std::make_unique<U2fSignOperation>(
      device.get(), std::move(request), sign_callback_receiver().callback());
  u2f_sign->Start();
  sign_callback_receiver().WaitForCallback();

  EXPECT_EQ(CtapDeviceResponseCode::kCtap2ErrCredentialNotValid,
            sign_callback_receiver().status());
  EXPECT_FALSE(sign_callback_receiver().value());
}

// Tests that request is dropped gracefully if device returns error on all
// requests (including fake enrollment).
TEST_F(U2fSignOperationTest, FakeEnrollErroringOut) {
  auto request = CreateSignRequest(
      {fido_parsing_utils::Materialize(test_data::kU2fSignKeyHandle)});

  // Simulates a device that errors out on all requests (including the sign
  // request and fake registration attempt). The device should then be abandoned
  // to prevent the test from crashing or timing out.
  auto device = std::make_unique<MockFidoDevice>();
  EXPECT_CALL(*device, GetId()).WillRepeatedly(::testing::Return("device0"));
  InSequence s;
  device->ExpectRequestAndRespondWith(test_data::kU2fCheckOnlySignCommandApdu,
                                      test_data::kU2fWrongDataApduResponse);
  device->ExpectRequestAndRespondWith(test_data::kU2fFakeRegisterCommand,
                                      test_data::kU2fWrongDataApduResponse);

  auto u2f_sign = std::make_unique<U2fSignOperation>(
      device.get(), std::move(request), sign_callback_receiver().callback());
  u2f_sign->Start();
  sign_callback_receiver().WaitForCallback();

  EXPECT_EQ(CtapDeviceResponseCode::kCtap2ErrOther,
            sign_callback_receiver().status());
  EXPECT_FALSE(sign_callback_receiver().value());
}

// Tests the scenario where device returns success response, but the response is
// unparse-able.
TEST_F(U2fSignOperationTest, SignWithCorruptedResponse) {
  auto request = CreateSignRequest(
      {fido_parsing_utils::Materialize(test_data::kU2fSignKeyHandle)});

  auto device = std::make_unique<MockFidoDevice>();
  EXPECT_CALL(*device, GetId()).WillRepeatedly(::testing::Return("device"));
  InSequence s;
  device->ExpectRequestAndRespondWith(test_data::kU2fCheckOnlySignCommandApdu,
                                      test_data::kTestCorruptedU2fSignResponse);
  device->ExpectRequestAndRespondWith(test_data::kU2fSignCommandApdu,
                                      test_data::kTestCorruptedU2fSignResponse);

  auto u2f_sign = std::make_unique<U2fSignOperation>(
      device.get(), std::move(request), sign_callback_receiver().callback());
  u2f_sign->Start();
  sign_callback_receiver().WaitForCallback();

  EXPECT_EQ(CtapDeviceResponseCode::kCtap2ErrOther,
            sign_callback_receiver().status());
  EXPECT_FALSE(sign_callback_receiver().value());
}

TEST_F(U2fSignOperationTest, AlternativeApplicationParameter) {
  auto request = CreateSignRequest(
      {fido_parsing_utils::Materialize(test_data::kU2fSignKeyHandle)});
  request.SetAlternativeApplicationParameter(
      test_data::kAlternativeApplicationParameter);

  auto device = std::make_unique<MockFidoDevice>();
  EXPECT_CALL(*device, GetId()).WillRepeatedly(::testing::Return("device"));
  InSequence s;
  // The first request will use the primary app_param, which will be rejected.
  device->ExpectRequestAndRespondWith(test_data::kU2fCheckOnlySignCommandApdu,
                                      test_data::kU2fWrongDataApduResponse);
  // After the rejection, the U2F sign request with alternative application
  // parameter should be tried.
  device->ExpectRequestAndRespondWith(
      test_data::
          kU2fCheckOnlySignCommandApduWithAlternativeApplicationParameter,
      test_data::kApduEncodedNoErrorSignResponse);
  device->ExpectRequestAndRespondWith(
      test_data::kU2fSignCommandApduWithAlternativeApplicationParameter,
      test_data::kApduEncodedNoErrorSignResponse);

  auto u2f_sign = std::make_unique<U2fSignOperation>(
      device.get(), std::move(request), sign_callback_receiver().callback());
  u2f_sign->Start();
  sign_callback_receiver().WaitForCallback();

  EXPECT_EQ(CtapDeviceResponseCode::kSuccess,
            sign_callback_receiver().status());
  const auto& response_value = sign_callback_receiver().value();
  EXPECT_THAT(response_value->signature(),
              ::testing::ElementsAreArray(test_data::kU2fSignature));
  EXPECT_THAT(response_value->raw_credential_id(),
              ::testing::ElementsAreArray(test_data::kU2fSignKeyHandle));
  EXPECT_THAT(
      response_value->GetRpIdHash(),
      ::testing::ElementsAreArray(test_data::kAlternativeApplicationParameter));
}

// This is a regression test in response to https://crbug.com/833398.
TEST_F(U2fSignOperationTest, AlternativeApplicationParameterRejection) {
  auto request = CreateSignRequest(
      {fido_parsing_utils::Materialize(test_data::kU2fSignKeyHandle)});
  request.SetAlternativeApplicationParameter(
      test_data::kAlternativeApplicationParameter);

  auto device = std::make_unique<MockFidoDevice>();
  EXPECT_CALL(*device, GetId()).WillRepeatedly(::testing::Return("device"));
  InSequence s;
  // The first request will use the primary app_param, which will be rejected.
  device->ExpectRequestAndRespondWith(test_data::kU2fCheckOnlySignCommandApdu,
                                      test_data::kU2fWrongDataApduResponse);
  // After the rejection, request with alternative application parameter should
  // be tried, which will also be rejected.
  device->ExpectRequestAndRespondWith(
      test_data::
          kU2fCheckOnlySignCommandApduWithAlternativeApplicationParameter,
      test_data::kU2fWrongDataApduResponse);
  // The second rejection will trigger a bogus register command. This will be
  // rejected as well, triggering the device to be abandoned.
  device->ExpectRequestAndRespondWith(test_data::kU2fFakeRegisterCommand,
                                      test_data::kU2fWrongDataApduResponse);

  auto u2f_sign = std::make_unique<U2fSignOperation>(
      device.get(), std::move(request), sign_callback_receiver().callback());
  u2f_sign->Start();
  sign_callback_receiver().WaitForCallback();

  EXPECT_EQ(CtapDeviceResponseCode::kCtap2ErrOther,
            sign_callback_receiver().status());
  EXPECT_FALSE(sign_callback_receiver().value());
}

}  // namespace device