test_doh_server.h 3.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102
  1. // Copyright 2021 The Chromium Authors. All rights reserved.
  2. // Use of this source code is governed by a BSD-style license that can be
  3. // found in the LICENSE file.
  4. #ifndef NET_TEST_TEST_DOH_SERVER_H_
  5. #define NET_TEST_TEST_DOH_SERVER_H_
  6. #include <cstdint>
  7. #include <map>
  8. #include <memory>
  9. #include <string>
  10. #include <utility>
  11. #include "base/containers/span.h"
  12. #include "base/strings/string_piece.h"
  13. #include "base/synchronization/lock.h"
  14. #include "base/thread_annotations.h"
  15. #include "base/time/time.h"
  16. #include "net/base/ip_address.h"
  17. #include "net/dns/dns_response.h"
  18. #include "net/test/embedded_test_server/embedded_test_server.h"
  19. #include "net/test/embedded_test_server/http_response.h"
  20. namespace net {
  21. // TestDohServer is a test DoH server. It allows tests to specify DNS behavior
  22. // at the level of individual DNS records.
  23. class TestDohServer {
  24. public:
  25. TestDohServer();
  26. ~TestDohServer();
  27. // Configures the hostname the DoH server serves from. If not specified, the
  28. // server is accessed over 127.0.0.1. This determines the TLS certificate
  29. // used, and the hostname in `GetTemplate`.
  30. void SetHostname(base::StringPiece name);
  31. // Configures whether the server should fail all requests with an HTTP error.
  32. void SetFailRequests(bool fail_requests);
  33. // Adds `address` to the set of A (or AAAA, if IPv6) responses when querying
  34. // `name`. This is a convenience wrapper over `AddRecord`.
  35. void AddAddressRecord(base::StringPiece name,
  36. const IPAddress& address,
  37. base::TimeDelta ttl = base::Days(1));
  38. // Adds `record` to the set of records served by this server.
  39. void AddRecord(const DnsResourceRecord& record);
  40. // Starts the test server and returns true on success or false on failure.
  41. //
  42. // Note this method starts a background thread. In some tests, such as
  43. // browser_tests, the process is required to be single-threaded in the early
  44. // stages of test setup. Tests that call `GetTemplate` at that point should
  45. // call `InitializeAndListen` before `GetTemplate`, followed by
  46. // `StartAcceptingConnections` when threads are allowed. See
  47. // `EmbeddedTestServer` for an example.
  48. [[nodiscard]] bool Start();
  49. // Initializes the listening socket for the test server, allocating a
  50. // listening port, and returns true on success or false on failure. Call
  51. // `StartAcceptingConnections` to finish initialization.
  52. [[nodiscard]] bool InitializeAndListen();
  53. // Spawns a background thread and begins accepting connections. This method
  54. // must be called after `InitializeAndListen`.
  55. void StartAcceptingConnections();
  56. // Shuts down the server and waits until the shutdown is complete.
  57. [[nodiscard]] bool ShutdownAndWaitUntilComplete();
  58. // Returns the number of queries served so far.
  59. int QueriesServed();
  60. // Returns the URI template to connect to this server. The server's listening
  61. // port must have been allocated with `Start` or `InitializeAndListen` before
  62. // calling this function.
  63. std::string GetTemplate();
  64. // Behaves like `GetTemplate`, but returns a template without the "dns" URL
  65. // and thus can only be used with POST.
  66. std::string GetPostOnlyTemplate();
  67. private:
  68. std::unique_ptr<test_server::HttpResponse> HandleRequest(
  69. const test_server::HttpRequest& request);
  70. absl::optional<std::string> hostname_;
  71. base::Lock lock_;
  72. // The following fields are accessed from a background thread and protected by
  73. // `lock_`.
  74. bool fail_requests_ GUARDED_BY(lock_) = false;
  75. // Maps from query name and query type to a record set.
  76. std::multimap<std::pair<std::string, uint16_t>, DnsResourceRecord> records_
  77. GUARDED_BY(lock_);
  78. int queries_served_ GUARDED_BY(lock_) = 0;
  79. EmbeddedTestServer server_{EmbeddedTestServer::TYPE_HTTPS};
  80. };
  81. } // namespace net
  82. #endif // NET_TEST_TEST_DOH_SERVER_H_