// Copyright 2021 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "net/quic/dedicated_web_transport_http3_client.h"

#include <memory>
#include <string_view>

#include "base/memory/raw_ptr.h"
#include "base/strings/strcat.h"
#include "base/strings/string_number_conversions.h"
#include "base/test/gmock_callback_support.h"
#include "build/build_config.h"
#include "net/base/proxy_chain.h"
#include "net/base/proxy_server.h"
#include "net/base/schemeful_site.h"
#include "net/cert/mock_cert_verifier.h"
#include "net/dns/mock_host_resolver.h"
#include "net/proxy_resolution/configured_proxy_resolution_service.h"
#include "net/quic/crypto/proof_source_chromium.h"
#include "net/quic/quic_context.h"
#include "net/test/test_data_directory.h"
#include "net/test/test_with_task_environment.h"
#include "net/third_party/quiche/src/quiche/quic/test_tools/crypto_test_utils.h"
#include "net/third_party/quiche/src/quiche/quic/test_tools/quic_test_backend.h"
#include "net/tools/quic/quic_simple_server.h"
#include "net/tools/quic/quic_simple_server_socket.h"
#include "net/traffic_annotation/network_traffic_annotation_test_helper.h"
#include "net/url_request/url_request_context.h"
#include "net/url_request/url_request_context_builder.h"
#include "testing/gmock/include/gmock/gmock.h"
#include "testing/gtest/include/gtest/gtest.h"
#include "url/gurl.h"
#include "url/origin.h"

namespace net::test {
namespace {

using ::quic::test::MemSliceFromString;
using ::testing::_;
using ::testing::DoAll;
using ::testing::Optional;
using ::testing::SaveArg;

class MockVisitor : public WebTransportClientVisitor {
 public:
  MOCK_METHOD(void,
              OnConnected,
              (scoped_refptr<HttpResponseHeaders>),
              (override));
  MOCK_METHOD(void, OnConnectionFailed, (const WebTransportError&), (override));
  MOCK_METHOD(void,
              OnLocalNetworkAccessCheck,
              (const IPEndPoint&,
               const NetLogWithSource&,
               CompletionOnceCallback callback),
              (override));
  MOCK_METHOD(void, OnBeforeConnect, (const IPEndPoint&), (override));
  MOCK_METHOD(void,
              OnClosed,
              (const std::optional<WebTransportCloseInfo>&),
              (override));
  MOCK_METHOD(void, OnError, (const WebTransportError&), (override));
  MOCK_METHOD(void, OnDraining, (), (override));

  MOCK_METHOD0(OnIncomingBidirectionalStreamAvailable, void());
  MOCK_METHOD0(OnIncomingUnidirectionalStreamAvailable, void());
  MOCK_METHOD1(OnDatagramReceived, void(std::string_view));
  MOCK_METHOD0(OnCanCreateNewOutgoingBidirectionalStream, void());
  MOCK_METHOD0(OnCanCreateNewOutgoingUnidirectionalStream, void());
  MOCK_METHOD1(OnDatagramProcessed, void(std::optional<quic::DatagramStatus>));
};

// A clock that only mocks out WallNow(), but uses real Now() and
// ApproximateNow().  Useful for certificate verification.
class TestWallClock : public quic::QuicClock {
 public:
  quic::QuicTime Now() const override {
    return quic::QuicChromiumClock::GetInstance()->Now();
  }
  quic::QuicTime ApproximateNow() const override {
    return quic::QuicChromiumClock::GetInstance()->ApproximateNow();
  }
  quic::QuicWallTime WallNow() const override { return wall_now_; }

  void set_wall_now(quic::QuicWallTime now) { wall_now_ = now; }

 private:
  quic::QuicWallTime wall_now_ = quic::QuicWallTime::Zero();
};

class TestConnectionHelper : public quic::QuicConnectionHelperInterface {
 public:
  const quic::QuicClock* GetClock() const override { return &clock_; }
  quic::QuicRandom* GetRandomGenerator() override {
    return quic::QuicRandom::GetInstance();
  }
  quiche::QuicheBufferAllocator* GetStreamSendBufferAllocator() override {
    return &allocator_;
  }

  TestWallClock& clock() { return clock_; }

 private:
  TestWallClock clock_;
  quiche::SimpleBufferAllocator allocator_;
};

class DedicatedWebTransportHttp3Test : public TestWithTaskEnvironment {
 public:
  ~DedicatedWebTransportHttp3Test() override {
    if (server_ != nullptr) {
      server_->Shutdown();
    }
  }

  void SetUp() override {
    BuildContext(ConfiguredProxyResolutionService::CreateDirect());
    quic::QuicEnableVersion(quic::ParsedQuicVersion::RFCv1());
    origin_ = url::Origin::Create(GURL{"https://example.org"});
    anonymization_key_ =
        NetworkAnonymizationKey::CreateSameSite(SchemefulSite(origin_));

    // By default, quit on error instead of waiting for RunLoop() to time out.
    ON_CALL(visitor_, OnConnectionFailed(_))
        .WillByDefault([this](const WebTransportError& error) {
          LOG(ERROR) << "Connection failed: " << error;
          if (run_loop_) {
            run_loop_->Quit();
          }
        });
    ON_CALL(visitor_, OnError(_))
        .WillByDefault([this](const WebTransportError& error) {
          LOG(ERROR) << "Connection error: " << error;
          if (run_loop_) {
            run_loop_->Quit();
          }
        });
    ON_CALL(visitor_, OnLocalNetworkAccessCheck(_, _, _))
        .WillByDefault(base::test::RunOnceCallback<2>(OK));
  }

  // Use a URLRequestContextBuilder to set `context_`.
  void BuildContext(
      std::unique_ptr<ProxyResolutionService> proxy_resolution_service) {
    URLRequestContextBuilder builder;
    builder.set_proxy_resolution_service(std::move(proxy_resolution_service));

    auto cert_verifier = std::make_unique<MockCertVerifier>();
    cert_verifier->set_default_result(OK);
    builder.SetCertVerifier(std::move(cert_verifier));

    auto host_resolver = std::make_unique<MockHostResolver>();
    host_resolver->rules()->AddRule("test.example.com", "127.0.0.1");
    builder.set_host_resolver(std::move(host_resolver));

    auto helper = std::make_unique<TestConnectionHelper>();
    helper_ = helper.get();
    auto quic_context = std::make_unique<QuicContext>(std::move(helper));
    quic_context->params()->supported_versions.clear();
    // This is required to bypass the check that only allows known certificate
    // roots in QUIC.
    quic_context->params()->origins_to_force_quic_on.insert(
        url::SchemeHostPort("https", "test.example.com", 443));
    builder.set_quic_context(std::move(quic_context));

    builder.set_net_log(NetLog::Get());
    context_ = builder.Build();
  }

  GURL GetURL(const std::string& suffix) {
    return GURL{base::StrCat(
        {"https://test.example.com:", base::NumberToString(port_), suffix})};
  }

  void StartServer(std::unique_ptr<quic::ProofSource> proof_source = nullptr) {
    if (proof_source == nullptr) {
      proof_source = quic::test::crypto_test_utils::ProofSourceForTesting();
    }
    backend_.set_enable_webtransport(true);
    server_ = std::make_unique<QuicSimpleServer>(
        std::move(proof_source), quic::QuicConfig(),
        quic::QuicCryptoServerConfig::ConfigOptions(),
        AllSupportedQuicVersions(), &backend_);
    ASSERT_TRUE(server_->CreateUDPSocketAndListen(
        quic::QuicSocketAddress(quiche::QuicheIpAddress::Any6(), /*port=*/0)));
    port_ = server_->server_address().port();
  }

  void Run() {
    run_loop_ = std::make_unique<base::RunLoop>();
    run_loop_->Run();
  }

  auto StopRunning() {
    return [this]() {
      if (run_loop_) {
        run_loop_->Quit();
      }
    };
  }

 protected:
  quic::test::QuicFlagSaver flags_;  // Save/restore all QUIC flag values.
  std::unique_ptr<URLRequestContext> context_;
  std::unique_ptr<DedicatedWebTransportHttp3Client> client_;
  raw_ptr<TestConnectionHelper> helper_;  // Owned by |context_|.
  ::testing::NiceMock<MockVisitor> visitor_;
  std::unique_ptr<QuicSimpleServer> server_;
  std::unique_ptr<base::RunLoop> run_loop_;
  quic::test::QuicTestBackend backend_;

  int port_ = 0;
  url::Origin origin_;
  NetworkAnonymizationKey anonymization_key_;
};

TEST_F(DedicatedWebTransportHttp3Test, Connect) {
  StartServer();
  client_ = std::make_unique<DedicatedWebTransportHttp3Client>(
      GetURL("/echo"), origin_, &visitor_, anonymization_key_,
      handles::kInvalidNetworkHandle, context_.get(), WebTransportParameters());

  EXPECT_CALL(visitor_, OnBeforeConnect);
  EXPECT_CALL(visitor_, OnConnected).WillOnce(StopRunning());
  client_->Connect();
  Run();
  ASSERT_TRUE(client_->session() != nullptr);

  client_->Close(std::nullopt);
  EXPECT_CALL(visitor_, OnClosed(_)).WillOnce(StopRunning());
  Run();
}

// Check that the Local Network Access check returning an error correctly fails
// the connection before attempting the connection.
TEST_F(DedicatedWebTransportHttp3Test, ConnectLocalNetworkAccessCheckFail) {
  StartServer();
  client_ = std::make_unique<DedicatedWebTransportHttp3Client>(
      GetURL("/echo"), origin_, &visitor_, anonymization_key_,
      handles::kInvalidNetworkHandle, context_.get(), WebTransportParameters());

  EXPECT_CALL(visitor_, OnLocalNetworkAccessCheck)
      .WillOnce(base::test::RunOnceCallback<2>(
          ERR_BLOCKED_BY_LOCAL_NETWORK_ACCESS_CHECKS));

  WebTransportError error;
  EXPECT_CALL(visitor_, OnConnectionFailed)
      .WillOnce(DoAll(StopRunning(), SaveArg<0>(&error)));
  client_->Connect();
  Run();
  ASSERT_TRUE(client_->session() == nullptr);
  EXPECT_EQ(error.net_error, ERR_BLOCKED_BY_LOCAL_NETWORK_ACCESS_CHECKS);
}

// Check that connecting via a proxy fails. This is currently not implemented,
// but it's important that WebTransport not be usable to _bypass_ a proxy -- if
// a proxy is configured, it must be used.
TEST_F(DedicatedWebTransportHttp3Test, ConnectViaProxy) {
  BuildContext(
      ConfiguredProxyResolutionService::CreateFixedFromProxyChainsForTest(
          {ProxyChain::FromSchemeHostAndPort(ProxyServer::SCHEME_HTTPS, "test",
                                             80)},
          TRAFFIC_ANNOTATION_FOR_TESTS));
  StartServer();
  client_ = std::make_unique<DedicatedWebTransportHttp3Client>(
      GetURL("/echo"), origin_, &visitor_, anonymization_key_,
      handles::kInvalidNetworkHandle, context_.get(), WebTransportParameters());

  client_->Connect();
}

// TODO(crbug.com/40816637): The test is flaky on Mac and iOS.
#if BUILDFLAG(IS_IOS) || BUILDFLAG(IS_MAC)
#define MAYBE_CloseTimeout DISABLED_CloseTimeout
#else
#define MAYBE_CloseTimeout CloseTimeout
#endif
TEST_F(DedicatedWebTransportHttp3Test, MAYBE_CloseTimeout) {
  StartServer();
  client_ = std::make_unique<DedicatedWebTransportHttp3Client>(
      GetURL("/echo"), origin_, &visitor_, anonymization_key_,
      handles::kInvalidNetworkHandle, context_.get(), WebTransportParameters());

  EXPECT_CALL(visitor_, OnBeforeConnect);
  EXPECT_CALL(visitor_, OnConnected).WillOnce(StopRunning());
  client_->Connect();
  Run();
  ASSERT_TRUE(client_->session() != nullptr);

  // Delete the server and put up a no-op socket in its place to simulate the
  // traffic being dropped.  Note that this is normally not a supported way of
  // shutting down a QuicServer, and will generate a lot of errors in the logs.
  server_.reset();
  IPEndPoint bind_address(IPAddress::IPv6AllZeros(), port_);
  auto noop_socket =
      std::make_unique<UDPServerSocket>(/*net_log=*/nullptr, NetLogSource());
  noop_socket->AllowAddressReuse();
  ASSERT_GE(noop_socket->Listen(bind_address), 0);

  client_->Close(std::nullopt);
  EXPECT_CALL(visitor_, OnError(_)).WillOnce(StopRunning());
  Run();
}

TEST_F(DedicatedWebTransportHttp3Test, CloseReason) {
  StartServer();
  client_ = std::make_unique<DedicatedWebTransportHttp3Client>(
      GetURL("/session-close"), origin_, &visitor_, anonymization_key_,
      handles::kInvalidNetworkHandle, context_.get(), WebTransportParameters());

  EXPECT_CALL(visitor_, OnBeforeConnect);
  EXPECT_CALL(visitor_, OnConnected).WillOnce(StopRunning());
  client_->Connect();
  Run();
  ASSERT_TRUE(client_->session() != nullptr);

  quic::WebTransportStream* stream =
      client_->session()->OpenOutgoingUnidirectionalStream();
  ASSERT_TRUE(stream != nullptr);
  EXPECT_TRUE(stream->Write("42 test error"));
  EXPECT_TRUE(stream->SendFin());

  WebTransportCloseInfo close_info(42, "test error");
  std::optional<WebTransportCloseInfo> received_close_info;
  EXPECT_CALL(visitor_, OnClosed(_))
      .WillOnce(DoAll(StopRunning(), SaveArg<0>(&received_close_info)));
  Run();
  EXPECT_THAT(received_close_info, Optional(close_info));
}

TEST_F(DedicatedWebTransportHttp3Test, SessionDraining) {
  StartServer();
  client_ = std::make_unique<DedicatedWebTransportHttp3Client>(
      GetURL("/session-close"), origin_, &visitor_, anonymization_key_,
      handles::kInvalidNetworkHandle, context_.get(), WebTransportParameters());

  EXPECT_CALL(visitor_, OnBeforeConnect);
  EXPECT_CALL(visitor_, OnConnected).WillOnce(StopRunning());
  EXPECT_CALL(visitor_, OnClosed(_)).Times(0);
  EXPECT_CALL(visitor_, OnDraining()).WillOnce(StopRunning());
  client_->Connect();
  Run();
  ASSERT_TRUE(client_->session() != nullptr);

  // The "/session-close" endpoint sends a DRAIN_WEBTRANSPORT_SESSION capsule
  // when it receives the string "DRAIN" on a stream.
  quic::WebTransportStream* stream =
      client_->session()->OpenOutgoingUnidirectionalStream();
  ASSERT_TRUE(stream != nullptr);
  EXPECT_TRUE(stream->Write("DRAIN"));
  EXPECT_TRUE(stream->SendFin());

  Run();
}

TEST_F(DedicatedWebTransportHttp3Test,
       DrainingAfterConnectionFailureIsIgnored) {
  StartServer();
  client_ = std::make_unique<DedicatedWebTransportHttp3Client>(
      GetURL("/not-found"), origin_, &visitor_, anonymization_key_,
      handles::kInvalidNetworkHandle, context_.get(), WebTransportParameters());

  EXPECT_CALL(visitor_, OnBeforeConnect);
  EXPECT_CALL(visitor_, OnConnectionFailed).WillOnce(StopRunning());
  EXPECT_CALL(visitor_, OnDraining()).Times(0);
  client_->Connect();
  Run();
  ASSERT_EQ(client_->state(), WebTransportState::FAILED);

  client_->OnSessionDraining();
}

// Test negotiation of the application protocol via
// https://www.ietf.org/archive/id/draft-ietf-webtrans-http3-12.html#name-application-protocol-negoti
TEST_F(DedicatedWebTransportHttp3Test, SubprotocolHeader) {
  StartServer();
  WebTransportParameters parameters;
  parameters.application_protocols = {"first", "second", "third"};
  // The selected-subprotocol endpoint selects the first of the offered
  // protocols by default, and echoes it on a unidirectional stream.
  client_ = std::make_unique<DedicatedWebTransportHttp3Client>(
      GetURL("/selected-subprotocol"), origin_, &visitor_, anonymization_key_,
      handles::kInvalidNetworkHandle, context_.get(), parameters);

  bool stream_received = false;
  EXPECT_CALL(visitor_, OnConnected).WillOnce(StopRunning());
  EXPECT_CALL(visitor_, OnIncomingUnidirectionalStreamAvailable).WillOnce([&] {
    stream_received = true;
    StopRunning();
  });
  client_->Connect();
  Run();
  ASSERT_TRUE(client_->session() != nullptr);

  EXPECT_EQ(client_->session()->GetNegotiatedSubprotocol(), "first");

  if (!stream_received) {
    Run();
  }

  quic::WebTransportStream* stream =
      client_->session()->AcceptIncomingUnidirectionalStream();
  ASSERT_TRUE(stream != nullptr);
  std::string read_buffer;
  webtransport::Stream::ReadResult read_result = stream->Read(&read_buffer);
  ASSERT_TRUE(read_result.fin);
  EXPECT_EQ(read_buffer, "first");
}

// Test backend that captures request headers for inspection.
class HeaderCapturingBackend : public quic::test::QuicTestBackend {
 public:
  quic::QuicSimpleServerBackend::WebTransportResponse
  ProcessWebTransportRequest(const quiche::HttpHeaderBlock& request_headers,
                             quic::WebTransportSession* session) override {
    for (const auto& [name, value] : request_headers) {
      captured_headers_.emplace_back(std::string(name), std::string(value));
    }
    return QuicTestBackend::ProcessWebTransportRequest(request_headers,
                                                       session);
  }

  const std::vector<std::pair<std::string, std::string>>& captured_headers()
      const {
    return captured_headers_;
  }

 private:
  std::vector<std::pair<std::string, std::string>> captured_headers_;
};

class DedicatedWebTransportHttp3HeadersTest
    : public DedicatedWebTransportHttp3Test {
 public:
  ~DedicatedWebTransportHttp3HeadersTest() override {
    if (server_ != nullptr) {
      server_->Shutdown();
      server_.reset();
    }
  }

  void StartServerWithCapture() {
    capturing_backend_.set_enable_webtransport(true);
    server_ = std::make_unique<QuicSimpleServer>(
        quic::test::crypto_test_utils::ProofSourceForTesting(),
        quic::QuicConfig(), quic::QuicCryptoServerConfig::ConfigOptions(),
        AllSupportedQuicVersions(), &capturing_backend_);
    ASSERT_TRUE(server_->CreateUDPSocketAndListen(
        quic::QuicSocketAddress(quiche::QuicheIpAddress::Any6(), /*port=*/0)));
    port_ = server_->server_address().port();
  }

 protected:
  HeaderCapturingBackend capturing_backend_;
};

// Verify that additional_headers with mixed casing are lowercased and that
// duplicate names (differing only in case) have their values combined.
TEST_F(DedicatedWebTransportHttp3HeadersTest,
       AdditionalHeadersCasingAndDuplicates) {
  StartServerWithCapture();
  WebTransportParameters parameters;
  parameters.additional_headers = {
      {"X-Custom", "first"},
      {"x-custom", "second"},
  };
  client_ = std::make_unique<DedicatedWebTransportHttp3Client>(
      GetURL("/echo"), origin_, &visitor_, anonymization_key_,
      handles::kInvalidNetworkHandle, context_.get(), parameters);

  EXPECT_CALL(visitor_, OnBeforeConnect);
  EXPECT_CALL(visitor_, OnConnected).WillOnce(StopRunning());
  client_->Connect();
  Run();
  ASSERT_TRUE(client_->session() != nullptr);

  // Find x-custom in the captured headers. With AppendValueOrAddHeader, both
  // values are combined with a null separator (quiche's internal format).
  // With operator[], only "second" would be present.
  std::string custom_value;
  for (const auto& [name, value] : capturing_backend_.captured_headers()) {
    if (name == "x-custom") {
      custom_value = value;
      break;
    }
  }
  EXPECT_THAT(custom_value, ::testing::HasSubstr("first"));
  EXPECT_THAT(custom_value, ::testing::HasSubstr("second"));
}

}  // namespace
}  // namespace net::test