net/tools/quic/test_tools/quic_test_client.h

// 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.

#ifndef NET_TOOLS_QUIC_TEST_TOOLS_QUIC_TEST_CLIENT_H_
#define NET_TOOLS_QUIC_TEST_TOOLS_QUIC_TEST_CLIENT_H_

#include <stddef.h>
#include <stdint.h>

#include <cstdint>
#include <memory>
#include <string>

#include "base/macros.h"
#include "net/base/ip_address.h"
#include "net/base/ip_endpoint.h"
#include "net/quic/proto/cached_network_parameters.pb.h"
#include "net/quic/quic_framer.h"
#include "net/quic/quic_packet_creator.h"
#include "net/quic/quic_protocol.h"
#include "net/tools/balsa/balsa_frame.h"
#include "net/tools/epoll_server/epoll_server.h"
#include "net/tools/quic/quic_client.h"
#include "net/tools/quic/test_tools/simple_client.h"

using base::StringPiece;

namespace net {

class ProofVerifier;

class QuicPacketWriterWrapper;

namespace test {

class HTTPMessage;
class MockableQuicClient;

// A quic client which allows mocking out writes.
class MockableQuicClient : public QuicClient {
 public:
  MockableQuicClient(IPEndPoint server_address,
                     const QuicServerId& server_id,
                     const QuicVersionVector& supported_versions,
                     EpollServer* epoll_server);

  MockableQuicClient(IPEndPoint server_address,
                     const QuicServerId& server_id,
                     const QuicConfig& config,
                     const QuicVersionVector& supported_versions,
                     EpollServer* epoll_server);

  ~MockableQuicClient() override;
  QuicPacketWriter* CreateQuicPacketWriter() override;
  QuicConnectionId GenerateNewConnectionId() override;
  void UseWriter(QuicPacketWriterWrapper* writer);
  void UseConnectionId(QuicConnectionId connection_id);
  void SendCachedNetworkParamaters(
      const CachedNetworkParameters& cached_network_params) {
    cached_network_paramaters_ = cached_network_params;
  }

 private:
  QuicConnectionId override_connection_id_;  // ConnectionId to use, if nonzero
  QuicPacketWriterWrapper* test_writer_;
  CachedNetworkParameters cached_network_paramaters_;

  DISALLOW_COPY_AND_ASSIGN(MockableQuicClient);
};

// A toy QUIC client used for testing, mostly following the SimpleClient APIs.
class QuicTestClient : public test::SimpleClient,
                       public QuicSpdyStream::Visitor,
                       public QuicClientPushPromiseIndex::Delegate {
 public:
  QuicTestClient(IPEndPoint server_address,
                 const std::string& server_hostname,
                 const QuicVersionVector& supported_versions);
  QuicTestClient(IPEndPoint server_address,
                 const std::string& server_hostname,
                 const QuicConfig& config,
                 const QuicVersionVector& supported_versions);

  ~QuicTestClient() override;

  // Sets the |user_agent_id| of the |client_|.
  void SetUserAgentID(const std::string& user_agent_id);

  // Wraps data in a quic packet and sends it.
  ssize_t SendData(const std::string& data, bool last_data);
  // As above, but |delegate| will be notified when |data| is ACKed.
  ssize_t SendData(const std::string& data,
                   bool last_data,
                   QuicAckListenerInterface* delegate);

  // From SimpleClient
  // Clears any outstanding state and sends a simple GET of 'uri' to the
  // server.  Returns 0 if the request failed and no bytes were written.
  ssize_t SendRequest(const std::string& uri) override;
  // Sends requests for all the urls and waits for the responses.  To process
  // the individual responses as they are returned, the caller should use the
  // set the response_listener on the client().
  void SendRequestsAndWaitForResponses(
      const std::vector<std::string>& url_list);
  ssize_t SendMessage(const HTTPMessage& message) override;
  std::string SendCustomSynchronousRequest(const HTTPMessage& message) override;
  std::string SendSynchronousRequest(const std::string& uri) override;
  void Connect() override;
  void ResetConnection() override;
  void Disconnect() override;
  IPEndPoint local_address() const override;
  void ClearPerRequestState() override;
  void WaitForResponseForMs(int timeout_ms) override;
  void WaitForInitialResponseForMs(int timeout_ms) override;
  ssize_t Send(const void* buffer, size_t size) override;
  bool response_complete() const override;
  bool response_headers_complete() const override;
  const BalsaHeaders* response_headers() const override;
  int64_t response_size() const override;
  int response_header_size() const override;
  int64_t response_body_size() const override;
  size_t bytes_read() const override;
  size_t bytes_written() const override;
  bool buffer_body() const override;
  void set_buffer_body(bool buffer_body) override;
  bool ServerInLameDuckMode() const override;
  const std::string& response_body() override;
  bool connected() const override;
  // These functions are all unimplemented functions from SimpleClient, and log
  // DFATAL if called by users of SimpleClient.
  ssize_t SendAndWaitForResponse(const void* buffer, size_t size) override;
  void Bind(IPEndPoint* local_address) override;
  void MigrateSocket(const IPAddress& new_host) override;
  std::string SerializeMessage(const HTTPMessage& message) override;
  IPAddress bind_to_address() const override;
  void set_bind_to_address(const IPAddress& address) override;
  const IPEndPoint& address() const override;
  size_t requests_sent() const override;

  // Returns the response trailers as received by the |stream_|.
  const SpdyHeaderBlock& response_trailers() const;

  // From QuicSpdyStream::Visitor
  void OnClose(QuicSpdyStream* stream) override;

  // From QuicClientPushPromiseIndex::Delegate
  bool CheckVary(const SpdyHeaderBlock& client_request,
                 const SpdyHeaderBlock& promise_request,
                 const SpdyHeaderBlock& promise_response) override;
  void OnRendezvousResult(QuicSpdyStream*) override;

  // Configures client_ to take ownership of and use the writer.
  // Must be called before initial connect.
  void UseWriter(QuicPacketWriterWrapper* writer);
  // If the given ConnectionId is nonzero, configures client_ to use a specific
  // ConnectionId instead of a random one.
  void UseConnectionId(QuicConnectionId connection_id);

  // Returns nullptr if the maximum number of streams have already been created.
  QuicSpdyClientStream* GetOrCreateStream();

  // Calls GetOrCreateStream(), sends the request on the stream, and
  // stores the request in case it needs to be resent.  If |headers| is
  // null, only the body will be sent on the stream.
  ssize_t GetOrCreateStreamAndSendRequest(const BalsaHeaders* headers,
                                          StringPiece body,
                                          bool fin,
                                          QuicAckListenerInterface* delegate);

  QuicRstStreamErrorCode stream_error() { return stream_error_; }
  QuicErrorCode connection_error();

  MockableQuicClient* client();

  // cert_common_name returns the common name value of the server's certificate,
  // or the empty string if no certificate was presented.
  const std::string& cert_common_name() const;

  // cert_sct returns the signed timestamp of the server's certificate,
  // or the empty string if no signed timestamp was presented.
  const std::string& cert_sct() const;

  // Get the server config map.
  QuicTagValueMap GetServerConfig() const;

  void set_auto_reconnect(bool reconnect) { auto_reconnect_ = reconnect; }

  void set_priority(SpdyPriority priority) { priority_ = priority; }

  void WaitForWriteToFlush();

  EpollServer* epoll_server() { return &epoll_server_; }

  void set_allow_bidirectional_data(bool value) {
    allow_bidirectional_data_ = value;
  }

  bool allow_bidirectional_data() const { return allow_bidirectional_data_; }

  size_t num_requests() const { return num_requests_; }

  size_t num_responses() const { return num_responses_; }

  // Explicitly set the SNI value for this client, overriding the default
  // behavior which extracts the SNI value from the request URL.
  void OverrideSni(const std::string& sni) {
    override_sni_set_ = true;
    override_sni_ = sni;
  }

 protected:
  QuicTestClient();

  void Initialize();

  void set_client(MockableQuicClient* client) { client_.reset(client); }

 private:
  class TestClientDataToResend : public QuicClient::QuicDataToResend {
   public:
    TestClientDataToResend(BalsaHeaders* headers,
                           StringPiece body,
                           bool fin,
                           QuicTestClient* test_client,
                           QuicAckListenerInterface* delegate)
        : QuicClient::QuicDataToResend(headers, body, fin),
          test_client_(test_client),
          delegate_(delegate) {}

    ~TestClientDataToResend() override {}

    void Resend() override;

   protected:
    QuicTestClient* test_client_;
    QuicAckListenerInterface* delegate_;
  };

  // Given a uri, creates a simple HTTPMessage request message for testing.
  static void FillInRequest(const std::string& uri, HTTPMessage* message);

  bool HaveActiveStream();

  EpollServer epoll_server_;
  std::unique_ptr<MockableQuicClient> client_;  // The actual client
  QuicSpdyClientStream* stream_;

  QuicRstStreamErrorCode stream_error_;

  bool response_complete_;
  bool response_headers_complete_;
  mutable BalsaHeaders response_headers_;

  // Parsed response trailers (if present), copied from the stream in OnClose.
  SpdyHeaderBlock response_trailers_;

  SpdyPriority priority_;
  std::string response_;
  uint64_t bytes_read_;
  uint64_t bytes_written_;
  // The number of uncompressed HTTP header bytes received.
  int response_header_size_;
  // The number of HTTP body bytes received.
  int64_t response_body_size_;
  // True if we tried to connect already since the last call to Disconnect().
  bool connect_attempted_;
  // The client will auto-connect exactly once before sending data.  If
  // something causes a connection reset, it will not automatically reconnect
  // unless auto_reconnect_ is true.
  bool auto_reconnect_;
  // Should we buffer the response body? Defaults to true.
  bool buffer_body_;
  // When true allows the sending of a request to continue while the response is
  // arriving.
  bool allow_bidirectional_data_;
  // For async push promise rendezvous, validation may fail in which
  // case the request should be retried.
  std::unique_ptr<TestClientDataToResend> push_promise_data_to_resend_;
  // Number of requests/responses this client has sent/received.
  size_t num_requests_;
  size_t num_responses_;

  // If set, this value is used for the connection SNI, overriding the usual
  // logic which extracts the SNI from the request URL.
  bool override_sni_set_ = false;
  std::string override_sni_;

  DISALLOW_COPY_AND_ASSIGN(QuicTestClient);
};

}  // namespace test

}  // namespace net

#endif  // NET_TOOLS_QUIC_TEST_TOOLS_QUIC_TEST_CLIENT_H_