network_diagnostics: Add the UdpProber class

Network diagnostic routines will use the single-use UdpProber to send
data to a particular destination. The primary use case is to send a STUN
packet (as defined in RFC 5389) to a STUN server owned by Google and
validate that a response is received.

BUG=b/172994051
TEST=unit_tests --gtest_filter=UdpProberWithFakeNetworkContextTest.*

Change-Id: I0abfb70c04a93bef60d690a64f9d99189267b3e5
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2516257
Commit-Queue: Kartik Hegde <khegde@chromium.org>
Reviewed-by: Steven Bennetts <stevenjb@chromium.org>
Reviewed-by: Matt Menke <mmenke@chromium.org>
Cr-Commit-Position: refs/heads/master@{#827610}

chrome/browser/chromeos/BUILD.gn

@@ -1978,6 +1978,8 @@ source_set("chromeos") {
     "net/network_diagnostics/gateway_can_be_pinged_routine.h",
     "net/network_diagnostics/has_secure_wifi_connection_routine.cc",
     "net/network_diagnostics/has_secure_wifi_connection_routine.h",
+    "net/network_diagnostics/host_resolver.cc",
+    "net/network_diagnostics/host_resolver.h",
     "net/network_diagnostics/http_firewall_routine.cc",
     "net/network_diagnostics/http_firewall_routine.h",
     "net/network_diagnostics/http_request_manager.cc",
@@ -1998,6 +2000,8 @@ source_set("chromeos") {
     "net/network_diagnostics/signal_strength_routine.h",
     "net/network_diagnostics/tls_prober.cc",
     "net/network_diagnostics/tls_prober.h",
+    "net/network_diagnostics/udp_prober.cc",
+    "net/network_diagnostics/udp_prober.h",
     "net/network_health/network_health.cc",
     "net/network_health/network_health.h",
     "net/network_health/network_health_localized_strings.cc",
@@ -3593,8 +3597,11 @@ source_set("unit_tests") {
     "net/network_diagnostics/fake_network_context.h",
     "net/network_diagnostics/fake_tcp_connected_socket.cc",
     "net/network_diagnostics/fake_tcp_connected_socket.h",
+    "net/network_diagnostics/fake_udp_socket.cc",
+    "net/network_diagnostics/fake_udp_socket.h",
     "net/network_diagnostics/gateway_can_be_pinged_routine_unittest.cc",
     "net/network_diagnostics/has_secure_wifi_connection_routine_unittest.cc",
+    "net/network_diagnostics/host_resolver_unittest.cc",
     "net/network_diagnostics/http_firewall_routine_unittest.cc",
     "net/network_diagnostics/http_request_manager_unittest.cc",
     "net/network_diagnostics/https_firewall_routine_unittest.cc",
@@ -3605,6 +3612,7 @@ source_set("unit_tests") {
     "net/network_diagnostics/network_diagnostics_util_unittest.cc",
     "net/network_diagnostics/signal_strength_routine_unittest.cc",
     "net/network_diagnostics/tls_prober_unittest.cc",
+    "net/network_diagnostics/udp_prober_unittest.cc",
     "net/network_health/network_health_unittest.cc",
     "net/network_portal_detector_impl_unittest.cc",
     "net/network_pref_state_observer_unittest.cc",

chrome/browser/chromeos/net/network_diagnostics/fake_host_resolver.cc

@@ -43,13 +43,13 @@ void FakeHostResolver::ResolveHost(
     receiver_.reset();
     return;
   }
-  mojo::Remote<network::mojom::ResolveHostClient> response_client(
-      std::move(pending_response_client));
-  DnsResult* result = fake_dns_results_.front();
-  DCHECK(result);
-  fake_dns_results_.pop_front();
-  response_client->OnComplete(result->result_, result->resolve_error_info_,
-                              result->resolved_addresses_);
+  response_client_.Bind(std::move(pending_response_client));
+
+  DCHECK(fake_dns_result_);
+  response_client_->OnComplete(fake_dns_result_->result_,
+                               fake_dns_result_->resolve_error_info_,
+                               fake_dns_result_->resolved_addresses_);
+  fake_dns_result_.release();
 }
 
 void FakeHostResolver::MdnsListen(
@@ -60,5 +60,12 @@ void FakeHostResolver::MdnsListen(
   NOTIMPLEMENTED();
 }
 
+void FakeHostResolver::SetFakeDnsResult(
+    std::unique_ptr<DnsResult> fake_dns_result) {
+  DCHECK(!fake_dns_result_);
+
+  fake_dns_result_ = std::move(fake_dns_result);
+}
+
 }  // namespace network_diagnostics
 }  // namespace chromeos

chrome/browser/chromeos/net/network_diagnostics/fake_host_resolver.h

@@ -5,7 +5,8 @@
 #ifndef CHROME_BROWSER_CHROMEOS_NET_NETWORK_DIAGNOSTICS_FAKE_HOST_RESOLVER_H_
 #define CHROME_BROWSER_CHROMEOS_NET_NETWORK_DIAGNOSTICS_FAKE_HOST_RESOLVER_H_
 
-#include <deque>
+#include <memory>
+#include <utility>
 
 #include "base/optional.h"
 #include "mojo/public/cpp/bindings/pending_receiver.h"
@@ -49,10 +50,8 @@ class FakeHostResolver : public network::mojom::HostResolver {
       mojo::PendingRemote<network::mojom::MdnsListenClient> response_client,
       MdnsListenCallback callback) override;
 
-  // Sets the fake dns results.
-  void set_fake_dns_results(std::deque<DnsResult*> fake_dns_results) {
-    fake_dns_results_ = std::move(fake_dns_results);
-  }
+  // Sets the fake DNS result for single host resolutions.
+  void SetFakeDnsResult(std::unique_ptr<DnsResult> fake_dns_result);
 
   // If set to true, the binding pipe will be disconnected when attempting to
   // connect.
@@ -61,10 +60,12 @@ class FakeHostResolver : public network::mojom::HostResolver {
   }
 
  private:
+  // Handles calls to the HostResolver.
   mojo::Receiver<network::mojom::HostResolver> receiver_;
-  // Use the list of fake dns results to fake different responses for multiple
-  // calls to the host_resolver's ResolveHost().
-  std::deque<DnsResult*> fake_dns_results_;
+  // Responds to calls made to |this|.
+  mojo::Remote<network::mojom::ResolveHostClient> response_client_;
+  // Use the |fake_dns_result| to fake a single host resolution.
+  std::unique_ptr<DnsResult> fake_dns_result_ = nullptr;
   // Used to mimic the scenario where network::mojom::HostResolver receiver
   // is disconnected.
   bool disconnect_ = false;

chrome/browser/chromeos/net/network_diagnostics/fake_network_context.cc

@@ -6,6 +6,8 @@
 
 #include <utility>
 
+#include "base/logging.h"
+
 namespace chromeos {
 namespace network_diagnostics {
 
@@ -30,8 +32,8 @@ void FakeNetworkContext::CreateHostResolver(
     mojo::PendingReceiver<network::mojom::HostResolver> receiver) {
   DCHECK(!resolver_);
   resolver_ = std::make_unique<FakeHostResolver>(std::move(receiver));
-  resolver_->set_fake_dns_results(std::move(fake_dns_results_));
   resolver_->set_disconnect_during_host_resolution(host_resolution_disconnect_);
+  resolver_->SetFakeDnsResult(std::move(fake_dns_result_));
 }
 
 void FakeNetworkContext::CreateTCPConnectedSocket(
@@ -59,6 +61,24 @@ void FakeNetworkContext::CreateTCPConnectedSocket(
                           mojo::ScopedDataPipeProducerHandle());
 }
 
+void FakeNetworkContext::CreateUDPSocket(
+    mojo::PendingReceiver<network::mojom::UDPSocket> receiver,
+    mojo::PendingRemote<network::mojom::UDPSocketListener> listener) {
+  if (udp_connection_attempt_disconnect_) {
+    receiver.reset();
+    listener.reset();
+    return;
+  }
+
+  // Bind the receiver if UDP connection is successful.
+  if (udp_connect_code_ == net::OK) {
+    DCHECK(fake_udp_socket_);
+
+    fake_udp_socket_->BindReceiver(std::move(receiver));
+    fake_udp_socket_->BindRemote(std::move(listener));
+  }
+}
+
 void FakeNetworkContext::SetTCPConnectCode(
     base::Optional<net::Error>& tcp_connect_code) {
   if (tcp_connect_code.has_value()) {
@@ -74,5 +94,41 @@ void FakeNetworkContext::SetTLSUpgradeCode(
   }
 }
 
+void FakeNetworkContext::SetUdpConnectCode(net::Error udp_connect_code) {
+  fake_udp_socket_ = std::make_unique<FakeUdpSocket>();
+  fake_udp_socket_->set_udp_connect_code(udp_connect_code);
+}
+
+void FakeNetworkContext::SetUdpSendCode(net::Error udp_send_code) {
+  DCHECK(fake_udp_socket_);
+
+  fake_udp_socket_->set_udp_send_code(udp_send_code);
+}
+
+void FakeNetworkContext::SetDisconnectDuringUdpSendAttempt(bool disconnect) {
+  DCHECK(fake_udp_socket_);
+
+  fake_udp_socket_->set_disconnect_during_udp_send_attempt(disconnect);
+}
+
+void FakeNetworkContext::SetUdpOnReceivedCode(net::Error udp_on_received_code) {
+  DCHECK(fake_udp_socket_);
+
+  fake_udp_socket_->set_udp_on_received_code(udp_on_received_code);
+}
+
+void FakeNetworkContext::SetUdpOnReceivedData(
+    base::span<const uint8_t> udp_on_received_data) {
+  DCHECK(fake_udp_socket_);
+
+  fake_udp_socket_->set_udp_on_received_data(std::move(udp_on_received_data));
+}
+
+void FakeNetworkContext::SetDisconnectDuringUdpReceiveAttempt(bool disconnect) {
+  DCHECK(fake_udp_socket_);
+
+  fake_udp_socket_->set_disconnect_during_udp_receive_attempt(disconnect);
+}
+
 }  // namespace network_diagnostics
 }  // namespace chromeos

chrome/browser/chromeos/net/network_diagnostics/fake_network_context.h

@@ -7,9 +7,13 @@
 
 #include <deque>
 #include <memory>
+#include <utility>
 
+#include "base/containers/span.h"
+#include "base/optional.h"
 #include "chrome/browser/chromeos/net/network_diagnostics/fake_host_resolver.h"
 #include "chrome/browser/chromeos/net/network_diagnostics/fake_tcp_connected_socket.h"
+#include "chrome/browser/chromeos/net/network_diagnostics/fake_udp_socket.h"
 #include "services/network/test/test_network_context.h"
 
 namespace chromeos {
@@ -20,6 +24,8 @@ namespace network_diagnostics {
 class FakeNetworkContext : public network::TestNetworkContext {
  public:
   FakeNetworkContext();
+  FakeNetworkContext(const FakeNetworkContext&) = delete;
+  FakeNetworkContext& operator=(const FakeNetworkContext&) = delete;
   ~FakeNetworkContext() override;
 
   // network::TestNetworkContext:
@@ -36,16 +42,40 @@ class FakeNetworkContext : public network::TestNetworkContext {
       mojo::PendingRemote<network::mojom::SocketObserver> observer,
       CreateTCPConnectedSocketCallback callback) override;
 
-  // Sets the fake TCP connect code.
+  void CreateUDPSocket(
+      mojo::PendingReceiver<network::mojom::UDPSocket> receiver,
+      mojo::PendingRemote<network::mojom::UDPSocketListener> listener) override;
+
+  // Sets the fake TCP connect code. TODO(khegde): Change this to
+  // SetTCPConnectCompleteCode.
   void SetTCPConnectCode(base::Optional<net::Error>& tcp_connect_code);
 
   // Sets the fake TLS upgrade code.
   void SetTLSUpgradeCode(base::Optional<net::Error>& tls_upgrade_code);
 
-  // Sets the fake dns results.
-  void set_fake_dns_results(
-      std::deque<FakeHostResolver::DnsResult*> fake_dns_results) {
-    fake_dns_results_ = std::move(fake_dns_results);
+  // Sets the fake UDP connect code.
+  void SetUdpConnectCode(net::Error udp_connect_code);
+
+  // Sets the fake UDP send code.
+  void SetUdpSendCode(net::Error udp_send_code);
+
+  // Sets the state to mimic a fake disconnect during a UDP send attempt.
+  void SetDisconnectDuringUdpSendAttempt(bool disconnect);
+
+  // Sets the fake UDP on received code.
+  void SetUdpOnReceivedCode(net::Error udp_on_received_code);
+
+  // Sets the fake UDP on received data.
+  void SetUdpOnReceivedData(base::span<const uint8_t> udp_on_received_data);
+
+  // Sets the state to mimic a fake disconnect after receiving successful send
+  // confirmation, but before receiving any data.
+  void SetDisconnectDuringUdpReceiveAttempt(bool disconnect);
+
+  // Sets the fake DNS result. Used to test a single host resolution.
+  void set_fake_dns_result(
+      std::unique_ptr<FakeHostResolver::DnsResult> fake_dns_result) {
+    fake_dns_result_ = std::move(fake_dns_result);
   }
 
   // If set to true, the binding pipe will be disconnected when attempting to
@@ -66,15 +96,28 @@ class FakeNetworkContext : public network::TestNetworkContext {
     tls_upgrade_attempt_disconnect_ = disconnect;
   }
 
+  // If set to true, the binding pipe will be disconnected when attempting to
+  // connect.
+  void set_disconnect_during_udp_connection_attempt(bool disconnect) {
+    udp_connection_attempt_disconnect_ = disconnect;
+  }
+
  private:
   // Fake host resolver.
   std::unique_ptr<FakeHostResolver> resolver_;
   // Fake DNS lookup results.
   std::deque<FakeHostResolver::DnsResult*> fake_dns_results_;
+  // Fake DNS lookup result.
+  std::unique_ptr<FakeHostResolver::DnsResult> fake_dns_result_;
   // Provides the TCP socket functionality for tests.
   std::unique_ptr<FakeTCPConnectedSocket> fake_tcp_connected_socket_;
+  // Provides the UDP socket functionality for tests.
+  std::unique_ptr<FakeUdpSocket> fake_udp_socket_;
   // TCP connect code corresponding to the connection attempt.
   net::Error tcp_connect_code_ = net::OK;
+  // UDP connect code corresponding to the connection attempt.
+  net::Error udp_connect_code_ = net::OK;
+
   // Used to mimic the scenario where network::mojom::HostResolver receiver
   // is disconnected.
   bool host_resolution_disconnect_ = false;
@@ -84,6 +127,9 @@ class FakeNetworkContext : public network::TestNetworkContext {
   // Used to mimic the scenario where network::mojom::TLSClientSocket receiver
   // is disconnected.
   bool tls_upgrade_attempt_disconnect_ = false;
+  // Used to mimic the scenario where network::mojom::UDPSocket receiver is
+  // disconnected while connecting.
+  bool udp_connection_attempt_disconnect_ = false;
 };
 
 }  // namespace network_diagnostics

chrome/browser/chromeos/net/network_diagnostics/fake_udp_socket.cc

+// Copyright 2020 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 "chrome/browser/chromeos/net/network_diagnostics/fake_udp_socket.h"
+
+#include <utility>
+
+#include "base/optional.h"
+#include "net/base/ip_endpoint.h"
+
+namespace chromeos {
+namespace network_diagnostics {
+
+namespace {}  // namespace
+
+FakeUdpSocket::FakeUdpSocket() = default;
+
+FakeUdpSocket::~FakeUdpSocket() = default;
+
+void FakeUdpSocket::Connect(const net::IPEndPoint& remote_addr,
+                            network::mojom::UDPSocketOptionsPtr options,
+                            ConnectCallback callback) {
+  if (mojo_disconnect_on_connect_) {
+    receiver_.reset();
+    return;
+  }
+  std::move(callback).Run(udp_connect_code_, net::IPEndPoint());
+}
+
+void FakeUdpSocket::Bind(const net::IPEndPoint& local_addr,
+                         network::mojom::UDPSocketOptionsPtr options,
+                         BindCallback callback) {
+  NOTREACHED();
+}
+
+void FakeUdpSocket::SetBroadcast(bool broadcast,
+                                 SetBroadcastCallback callback) {
+  NOTREACHED();
+}
+
+void FakeUdpSocket::SetSendBufferSize(int32_t send_buffer_size,
+                                      SetSendBufferSizeCallback callback) {
+  NOTREACHED();
+}
+
+void FakeUdpSocket::SetReceiveBufferSize(int32_t receive_buffer_size,
+                                         SetSendBufferSizeCallback callback) {
+  NOTREACHED();
+}
+
+void FakeUdpSocket::JoinGroup(const net::IPAddress& group_address,
+                              JoinGroupCallback callback) {
+  NOTREACHED();
+}
+
+void FakeUdpSocket::LeaveGroup(const net::IPAddress& group_address,
+                               LeaveGroupCallback callback) {
+  NOTREACHED();
+}
+
+void FakeUdpSocket::ReceiveMore(uint32_t num_additional_datagrams) {
+  DCHECK(remote_.is_bound());
+
+  if (mojo_disconnect_on_receive_) {
+    remote_.reset();
+    return;
+  }
+
+  remote_->OnReceived(udp_on_received_code_, net::IPEndPoint(),
+                      std::move(udp_on_received_data_));
+}
+
+void FakeUdpSocket::ReceiveMoreWithBufferSize(uint32_t num_additional_datagrams,
+                                              uint32_t buffer_size) {
+  NOTREACHED();
+}
+
+void FakeUdpSocket::SendTo(
+    const net::IPEndPoint& dest_addr,
+    base::span<const uint8_t> data,
+    const net::MutableNetworkTrafficAnnotationTag& traffic_annotation,
+    SendToCallback callback) {
+  NOTREACHED();
+}
+
+void FakeUdpSocket::Send(
+    base::span<const uint8_t> data,
+    const net::MutableNetworkTrafficAnnotationTag& traffic_annotation,
+    SendCallback callback) {
+  if (mojo_disconnect_on_send_) {
+    receiver_.reset();
+    return;
+  }
+  std::move(callback).Run(udp_send_code_);
+}
+
+void FakeUdpSocket::Close() {
+  NOTREACHED();
+}
+
+void FakeUdpSocket::BindReceiver(
+    mojo::PendingReceiver<network::mojom::UDPSocket> socket) {
+  DCHECK(!receiver_.is_bound());
+
+  receiver_.Bind(std::move(socket));
+}
+
+void FakeUdpSocket::BindRemote(
+    mojo::PendingRemote<network::mojom::UDPSocketListener> socket_listener) {
+  DCHECK(!remote_.is_bound());
+
+  remote_.Bind(std::move(socket_listener));
+}
+
+}  // namespace network_diagnostics
+}  // namespace chromeos

chrome/browser/chromeos/net/network_diagnostics/fake_udp_socket.h

+// Copyright 2020 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 CHROME_BROWSER_CHROMEOS_NET_NETWORK_DIAGNOSTICS_FAKE_UDP_SOCKET_H_
+#define CHROME_BROWSER_CHROMEOS_NET_NETWORK_DIAGNOSTICS_FAKE_UDP_SOCKET_H_
+
+#include "mojo/public/cpp/bindings/pending_receiver.h"
+#include "mojo/public/cpp/bindings/pending_remote.h"
+#include "mojo/public/cpp/bindings/receiver.h"
+#include "mojo/public/cpp/bindings/remote.h"
+#include "net/base/net_errors.h"
+#include "services/network/public/mojom/udp_socket.mojom.h"
+
+namespace chromeos {
+namespace network_diagnostics {
+
+// Provides some UDP socket functionality in tests. Most methods, unless
+// otherwise noted, are not expected to be called in tests or in a non Chrome OS
+// environment.
+class FakeUdpSocket : public network::mojom::UDPSocket {
+ public:
+  FakeUdpSocket();
+  FakeUdpSocket(const FakeUdpSocket&) = delete;
+  FakeUdpSocket& operator=(const FakeUdpSocket&) = delete;
+  ~FakeUdpSocket() override;
+
+  // network::mojom::UDPSocket:
+  // Used in the fake.
+  void Connect(const net::IPEndPoint& remote_addr,
+               network::mojom::UDPSocketOptionsPtr options,
+               ConnectCallback callback) override;
+  void Bind(const net::IPEndPoint& local_addr,
+            network::mojom::UDPSocketOptionsPtr options,
+            BindCallback callback) override;
+  void SetBroadcast(bool broadcast, SetBroadcastCallback callback) override;
+  void SetSendBufferSize(int32_t send_buffer_size,
+                         SetSendBufferSizeCallback callback) override;
+  void SetReceiveBufferSize(int32_t receive_buffer_size,
+                            SetSendBufferSizeCallback callback) override;
+  void JoinGroup(const net::IPAddress& group_address,
+                 JoinGroupCallback callback) override;
+  void LeaveGroup(const net::IPAddress& group_address,
+                  LeaveGroupCallback callback) override;
+  // Used in the fake.
+  void ReceiveMore(uint32_t num_additional_datagrams) override;
+  void ReceiveMoreWithBufferSize(uint32_t num_additional_datagrams,
+                                 uint32_t buffer_size) override;
+  void SendTo(const net::IPEndPoint& dest_addr,
+              base::span<const uint8_t> data,
+              const net::MutableNetworkTrafficAnnotationTag& traffic_annotation,
+              SendToCallback callback) override;
+  // Used in the fake.
+  void Send(base::span<const uint8_t> data,
+            const net::MutableNetworkTrafficAnnotationTag& traffic_annotation,
+            SendCallback callback) override;
+  void Close() override;
+
+  // Binds the pending receiver to |this|.
+  void BindReceiver(mojo::PendingReceiver<network::mojom::UDPSocket> socket);
+
+  // Binds the pending remote to a remote held by |this|.
+  void BindRemote(
+      mojo::PendingRemote<network::mojom::UDPSocketListener> socket_listener);
+
+  void set_udp_connect_code(net::Error udp_connect_code) {
+    udp_connect_code_ = udp_connect_code;
+  }
+
+  void set_udp_send_code(net::Error udp_send_code) {
+    udp_send_code_ = udp_send_code;
+  }
+
+  void set_udp_on_received_code(net::Error udp_on_received_code) {
+    udp_on_received_code_ = udp_on_received_code;
+  }
+
+  void set_udp_on_received_data(
+      base::span<const uint8_t> udp_on_received_data) {
+    udp_on_received_data_ = std::move(udp_on_received_data);
+  }
+
+  void set_disconnect_during_udp_connection_attempt(bool disconnect) {
+    mojo_disconnect_on_connect_ = disconnect;
+  }
+
+  void set_disconnect_during_udp_send_attempt(bool disconnect) {
+    mojo_disconnect_on_send_ = disconnect;
+  }
+
+  void set_disconnect_during_udp_receive_attempt(bool disconnect) {
+    mojo_disconnect_on_receive_ = disconnect;
+  }
+
+ private:
+  mojo::Receiver<network::mojom::UDPSocket> receiver_{this};
+  mojo::Remote<network::mojom::UDPSocketListener> remote_;
+  net::Error udp_connect_code_ = net::ERR_FAILED;
+  net::Error udp_send_code_ = net::ERR_FAILED;
+  net::Error udp_on_received_code_ = net::ERR_FAILED;
+  base::span<const uint8_t> udp_on_received_data_ = {};
+  bool mojo_disconnect_on_connect_ = false;
+  bool mojo_disconnect_on_send_ = false;
+  bool mojo_disconnect_on_receive_ = false;
+};
+
+}  // namespace network_diagnostics
+}  // namespace chromeos
+
+#endif  // CHROME_BROWSER_CHROMEOS_NET_NETWORK_DIAGNOSTICS_FAKE_UDP_SOCKET_H_

chrome/browser/chromeos/net/network_diagnostics/host_resolver.cc

+// Copyright 2020 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 "chrome/browser/chromeos/net/network_diagnostics/host_resolver.h"
+
+#include <utility>
+
+#include "base/bind.h"
+#include "base/optional.h"
+#include "net/base/net_errors.h"
+#include "net/base/network_isolation_key.h"
+#include "net/dns/public/dns_config_overrides.h"
+
+namespace chromeos {
+namespace network_diagnostics {
+
+HostResolver::ResolutionResult::ResolutionResult(
+    int result,
+    const net::ResolveErrorInfo& resolve_error_info,
+    const base::Optional<net::AddressList>& resolved_addresses)
+    : result(result),
+      resolve_error_info(resolve_error_info),
+      resolved_addresses(resolved_addresses) {}
+
+HostResolver::ResolutionResult::~ResolutionResult() = default;
+
+HostResolver::HostResolver(const GURL& url,
+                           network::mojom::NetworkContext* network_context,
+                           OnResolutionComplete callback)
+    : HostResolver(net::HostPortPair::FromURL(url),
+                   network_context,
+                   std::move(callback)) {}
+
+HostResolver::HostResolver(const net::HostPortPair& host_port_pair,
+                           network::mojom::NetworkContext* network_context,
+                           OnResolutionComplete callback)
+    : callback_(std::move(callback)) {
+  DCHECK(network_context);
+  DCHECK(callback_);
+
+  network_context->CreateHostResolver(
+      net::DnsConfigOverrides(), host_resolver_.BindNewPipeAndPassReceiver());
+  // Disconnect handler will be invoked if the network service crashes.
+  host_resolver_.set_disconnect_handler(base::BindOnce(
+      &HostResolver::OnMojoConnectionError, base::Unretained(this)));
+
+  network::mojom::ResolveHostParametersPtr parameters =
+      network::mojom::ResolveHostParameters::New();
+  parameters->dns_query_type = net::DnsQueryType::A;
+  parameters->source = net::HostResolverSource::DNS;
+  parameters->cache_usage =
+      network::mojom::ResolveHostParameters::CacheUsage::DISALLOWED;
+
+  host_resolver_->ResolveHost(
+      host_port_pair, net::NetworkIsolationKey::CreateTransient(),
+      std::move(parameters), receiver_.BindNewPipeAndPassRemote());
+}
+
+HostResolver::~HostResolver() = default;
+
+void HostResolver::OnComplete(
+    int result,
+    const net::ResolveErrorInfo& resolve_error_info,
+    const base::Optional<net::AddressList>& resolved_addresses) {
+  receiver_.reset();
+  host_resolver_.reset();
+
+  ResolutionResult resolution_result{result, resolve_error_info,
+                                     resolved_addresses};
+  std::move(callback_).Run(resolution_result);
+}
+
+void HostResolver::OnMojoConnectionError() {
+  OnComplete(net::ERR_NAME_NOT_RESOLVED, net::ResolveErrorInfo(net::ERR_FAILED),
+             base::nullopt);
+}
+
+}  // namespace network_diagnostics
+}  // namespace chromeos

chrome/browser/chromeos/net/network_diagnostics/host_resolver.h

+// Copyright 2020 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 CHROME_BROWSER_CHROMEOS_NET_NETWORK_DIAGNOSTICS_HOST_RESOLVER_H_
+#define CHROME_BROWSER_CHROMEOS_NET_NETWORK_DIAGNOSTICS_HOST_RESOLVER_H_
+
+#include "base/callback.h"
+#include "mojo/public/cpp/bindings/receiver.h"
+#include "mojo/public/cpp/bindings/remote.h"
+#include "net/base/host_port_pair.h"
+#include "net/base/net_errors.h"
+#include "services/network/public/cpp/resolve_host_client_base.h"
+#include "services/network/public/mojom/network_context.mojom.h"
+#include "url/gurl.h"
+
+namespace chromeos {
+namespace network_diagnostics {
+
+// Performs a DNS host resolution. This is a single-use class.
+class HostResolver : public network::ResolveHostClientBase {
+ public:
+  struct ResolutionResult {
+    ResolutionResult(
+        int result,
+        const net::ResolveErrorInfo& resolve_error_info,
+        const base::Optional<net::AddressList>& resolved_addresses);
+    ~ResolutionResult();
+
+    int result;
+    net::ResolveErrorInfo resolve_error_info;
+    base::Optional<net::AddressList> resolved_addresses;
+  };
+  using OnResolutionComplete = base::OnceCallback<void(ResolutionResult&)>;
+
+  // Performs the DNS resolution of a specified |url|. Note that |callback|
+  // will not be called until construction is complete.
+  HostResolver(const GURL& url,
+               network::mojom::NetworkContext* network_context,
+               OnResolutionComplete callback);
+
+  // Performs the DNS resolution of a specified |host_port_pair|. Note that
+  // |callback| will not be called until construction is complete.
+  HostResolver(const net::HostPortPair& host_port_pair,
+               network::mojom::NetworkContext* network_context,
+               OnResolutionComplete callback);
+
+  HostResolver(const HostResolver&) = delete;
+  HostResolver& operator=(const HostResolver&) = delete;
+  ~HostResolver() override;
+
+  // network::mojom::ResolveHostClient:
+  void OnComplete(
+      int result,
+      const net::ResolveErrorInfo& resolve_error_info,
+      const base::Optional<net::AddressList>& resolved_addresses) override;
+
+ private:
+  // Handles Mojo connection errors during host resolution.
+  void OnMojoConnectionError();
+
+  // Callback invoked once resolution is complete.
+  OnResolutionComplete callback_;
+  // Receiver endpoint to handle host resolution results.
+  mojo::Receiver<network::mojom::ResolveHostClient> receiver_{this};
+  // Remote endpoint that calls the network service.
+  mojo::Remote<network::mojom::HostResolver> host_resolver_;
+};
+
+}  // namespace network_diagnostics
+}  // namespace chromeos
+
+#endif  // CHROME_BROWSER_CHROMEOS_NET_NETWORK_DIAGNOSTICS_HOST_RESOLVER_H_

chrome/browser/chromeos/net/network_diagnostics/host_resolver_unittest.cc

+// Copyright 2020 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 "chrome/browser/chromeos/net/network_diagnostics/host_resolver.h"
+
+#include <memory>
+
+#include "base/bind.h"
+#include "base/run_loop.h"
+#include "chrome/browser/chromeos/net/network_diagnostics/fake_host_resolver.h"
+#include "chrome/browser/chromeos/net/network_diagnostics/fake_network_context.h"
+#include "content/public/test/browser_task_environment.h"
+#include "net/base/address_list.h"
+#include "net/base/host_port_pair.h"
+#include "net/base/ip_address.h"
+#include "net/base/ip_endpoint.h"
+#include "net/dns/public/resolve_error_info.h"
+#include "testing/gtest/include/gtest/gtest.h"
+#include "url/gurl.h"
+
+namespace chromeos {
+namespace network_diagnostics {
+
+class HostResolverTest : public ::testing::Test {
+ public:
+  HostResolverTest() = default;
+  HostResolverTest(const HostResolverTest&) = delete;
+  HostResolverTest& operator=(const HostResolverTest&) = delete;
+
+  void InitializeNetworkContext(
+      std::unique_ptr<FakeHostResolver::DnsResult> fake_dns_result) {
+    fake_network_context_.set_fake_dns_result(std::move(fake_dns_result));
+  }
+
+  FakeNetworkContext* fake_network_context() { return &fake_network_context_; }
+
+ protected:
+  const net::HostPortPair kFakeHostPortPair =
+      net::HostPortPair::FromString("fake_stun_server.com:80");
+  const GURL kFakeUrl{"https://www.FAKE_HOST_NAME.com:1234/"};
+  const net::IPEndPoint kFakeIPAddress{
+      net::IPEndPoint(net::IPAddress::IPv4Localhost(), /*port=*/1234)};
+  std::unique_ptr<HostResolver> host_resolver_;
+
+ private:
+  content::BrowserTaskEnvironment task_environment_;
+  FakeNetworkContext fake_network_context_;
+};
+
+TEST_F(HostResolverTest, TestSuccessfulResolutionWithUrl) {
+  auto address_list = net::AddressList(kFakeIPAddress);
+  auto fake_dns_result = std::make_unique<FakeHostResolver::DnsResult>(
+      net::OK, net::ResolveErrorInfo(net::OK), address_list);
+  InitializeNetworkContext(std::move(fake_dns_result));
+  HostResolver::ResolutionResult resolution_result{
+      net::ERR_FAILED, net::ResolveErrorInfo(net::OK), base::nullopt};
+  base::RunLoop run_loop;
+  host_resolver_ = std::make_unique<HostResolver>(
+      kFakeUrl, fake_network_context(),
+      base::BindOnce(
+          [](HostResolver::ResolutionResult* resolution_result,
+             base::OnceClosure quit_closure,
+             HostResolver::ResolutionResult& res_result) {
+            resolution_result->result = res_result.result;
+            resolution_result->resolve_error_info =
+                res_result.resolve_error_info;
+            resolution_result->resolved_addresses =
+                res_result.resolved_addresses;
+            std::move(quit_closure).Run();
+          },
+          &resolution_result, run_loop.QuitClosure()));
+  run_loop.Run();
+
+  EXPECT_EQ(resolution_result.result, net::OK);
+  EXPECT_EQ(resolution_result.resolve_error_info,
+            net::ResolveErrorInfo(net::OK));
+  EXPECT_EQ(resolution_result.resolved_addresses.value().size(), 1);
+  EXPECT_EQ(resolution_result.resolved_addresses.value().front(),
+            address_list.front());
+}
+
+TEST_F(HostResolverTest, TestSuccessfulResolutionWithHostPortPair) {
+  auto address_list = net::AddressList(kFakeIPAddress);
+  auto fake_dns_result = std::make_unique<FakeHostResolver::DnsResult>(
+      net::OK, net::ResolveErrorInfo(net::OK), address_list);
+  InitializeNetworkContext(std::move(fake_dns_result));
+  HostResolver::ResolutionResult resolution_result{
+      net::ERR_FAILED, net::ResolveErrorInfo(net::OK), base::nullopt};
+  base::RunLoop run_loop;
+  host_resolver_ = std::make_unique<HostResolver>(
+      kFakeHostPortPair, fake_network_context(),
+      base::BindOnce(
+          [](HostResolver::ResolutionResult* resolution_result,
+             base::OnceClosure quit_closure,
+             HostResolver::ResolutionResult& res_result) {
+            resolution_result->result = res_result.result;
+            resolution_result->resolve_error_info =
+                res_result.resolve_error_info;
+            resolution_result->resolved_addresses =
+                res_result.resolved_addresses;
+            std::move(quit_closure).Run();
+          },
+          &resolution_result, run_loop.QuitClosure()));
+  run_loop.Run();
+
+  EXPECT_EQ(resolution_result.result, net::OK);
+  EXPECT_EQ(resolution_result.resolve_error_info,
+            net::ResolveErrorInfo(net::OK));
+  EXPECT_EQ(resolution_result.resolved_addresses.value().size(), 1);
+  EXPECT_EQ(resolution_result.resolved_addresses.value().front(),
+            address_list.front());
+}
+
+TEST_F(HostResolverTest, TestFailedHostResolution) {
+  auto fake_dns_result = std::make_unique<FakeHostResolver::DnsResult>(
+      net::ERR_NAME_NOT_RESOLVED,
+      net::ResolveErrorInfo(net::ERR_NAME_NOT_RESOLVED), base::nullopt);
+  InitializeNetworkContext(std::move(fake_dns_result));
+  HostResolver::ResolutionResult resolution_result{
+      net::ERR_FAILED, net::ResolveErrorInfo(net::OK), base::nullopt};
+  base::RunLoop run_loop;
+  host_resolver_ = std::make_unique<HostResolver>(
+      kFakeHostPortPair, fake_network_context(),
+      base::BindOnce(
+          [](HostResolver::ResolutionResult* resolution_result,
+             base::OnceClosure quit_closure,
+             HostResolver::ResolutionResult& res_result) {
+            resolution_result->result = res_result.result;
+            resolution_result->resolve_error_info =
+                res_result.resolve_error_info;
+            resolution_result->resolved_addresses =
+                res_result.resolved_addresses;
+            std::move(quit_closure).Run();
+          },
+          &resolution_result, run_loop.QuitClosure()));
+  run_loop.Run();
+
+  EXPECT_EQ(resolution_result.result, net::ERR_NAME_NOT_RESOLVED);
+  EXPECT_EQ(resolution_result.resolve_error_info,
+            net::ResolveErrorInfo(net::ERR_NAME_NOT_RESOLVED));
+  ASSERT_FALSE(resolution_result.resolved_addresses.has_value());
+}
+
+TEST_F(HostResolverTest, TestMojoDisconnectDuringHostResolution) {
+  InitializeNetworkContext(/*fake_dns_result=*/{});
+  fake_network_context()->set_disconnect_during_host_resolution(true);
+  HostResolver::ResolutionResult resolution_result{
+      net::ERR_FAILED, net::ResolveErrorInfo(net::OK), base::nullopt};
+  base::RunLoop run_loop;
+  host_resolver_ = std::make_unique<HostResolver>(
+      kFakeHostPortPair, fake_network_context(),
+      base::BindOnce(
+          [](HostResolver::ResolutionResult* resolution_result,
+             base::OnceClosure quit_closure,
+             HostResolver::ResolutionResult& res_result) {
+            resolution_result->result = res_result.result;
+            resolution_result->resolve_error_info =
+                res_result.resolve_error_info;
+            resolution_result->resolved_addresses =
+                res_result.resolved_addresses;
+            std::move(quit_closure).Run();
+          },
+          &resolution_result, run_loop.QuitClosure()));
+  run_loop.Run();
+
+  EXPECT_EQ(resolution_result.result, net::ERR_NAME_NOT_RESOLVED);
+  EXPECT_EQ(resolution_result.resolve_error_info,
+            net::ResolveErrorInfo(net::ERR_FAILED));
+  ASSERT_FALSE(resolution_result.resolved_addresses.has_value());
+}
+
+}  // namespace network_diagnostics
+}  // namespace chromeos

chrome/browser/chromeos/net/network_diagnostics/network_diagnostics_util.cc

@@ -149,6 +149,39 @@ Profile* GetUserProfile() {
   return profile;
 }
 
+const std::array<uint8_t, kStunHeaderSize>& GetStunHeader() {
+  static std::array<uint8_t, kStunHeaderSize> stun_header = {
+      0x00, 0x01, 0x00, 0x00, 0x21, 0x12, 0xa4, 0x42, 0x79, 0x64,
+      0x66, 0x36, 0x66, 0x53, 0x42, 0x73, 0x76, 0x77, 0x76, 0x75};
+
+  return stun_header;
+}
+
+net::NetworkTrafficAnnotationTag GetStunNetworkAnnotationTag() {
+  return net::DefineNetworkTrafficAnnotation("network_diagnostics_routines",
+                                             R"(
+      semantics {
+        sender: "NetworkDiagnosticsRoutines"
+        description:
+            "Routines send network traffic to hosts in order to "
+            "validate the internet connection on a device."
+        trigger:
+            "A routine attempts a socket connection or makes an http/s "
+            "request."
+        data:
+          "For UDP connections, data is sent along with the origin "
+          "(scheme-host-port). The primary purpose of the UDP prober is to "
+          "send a STUN packet header to a STUN server. For TCP connections, "
+          "only the origin is sent. No user identifier is sent along with the "
+          "data."
+        destination: GOOGLE_OWNED_SERVICE
+      }
+      policy {
+        cookies_allowed: NO
+      }
+  )");
+}
+
 }  // namespace util
 
 }  // namespace network_diagnostics

chrome/browser/chromeos/net/network_diagnostics/network_diagnostics_util.h

@@ -5,9 +5,12 @@
 #ifndef CHROME_BROWSER_CHROMEOS_NET_NETWORK_DIAGNOSTICS_NETWORK_DIAGNOSTICS_UTIL_H_
 #define CHROME_BROWSER_CHROMEOS_NET_NETWORK_DIAGNOSTICS_NETWORK_DIAGNOSTICS_UTIL_H_
 
+#include <array>
+#include <cstdint>
 #include <string>
 #include <vector>
 
+#include "net/traffic_annotation/network_traffic_annotation.h"
 #include "url/gurl.h"
 
 class Profile;
@@ -20,6 +23,9 @@ namespace util {
 // Generate 204 path.
 extern const char kGenerate204Path[];
 
+// STUN packet header size.
+constexpr int kStunHeaderSize = 20;
+
 // Returns the Gstatic host suffix. Network diagnostic routines attach a random
 // prefix to |kGstaticHostSuffix| to get a complete hostname.
 const char* GetGstaticHostSuffix();
@@ -78,6 +84,12 @@ std::vector<GURL> GetRandomHostsWithSchemeAndPortAndGenerate204Path(
 // Returns the profile associated with this account.
 Profile* GetUserProfile();
 
+// Returns a STUN packet with a header defined in RFC 5389.
+const std::array<uint8_t, kStunHeaderSize>& GetStunHeader();
+
+// Returns the traffic annotation tag for STUN traffic.
+net::NetworkTrafficAnnotationTag GetStunNetworkAnnotationTag();
+
 }  // namespace util
 
 }  // namespace network_diagnostics

chrome/browser/chromeos/net/network_diagnostics/tls_prober.cc

@@ -54,89 +54,6 @@ net::NetworkTrafficAnnotationTag GetTrafficAnnotationTag() {
 
 }  // namespace
 
-class TlsProber::HostResolver : public network::ResolveHostClientBase {
- public:
-  HostResolver(network::mojom::NetworkContext* network_context,
-               TlsProber* tls_prober);
-  HostResolver(const HostResolver&) = delete;
-  HostResolver& operator=(const HostResolver&) = delete;
-  ~HostResolver() override;
-
-  // network::mojom::ResolveHostClient:
-  void OnComplete(
-      int result,
-      const net::ResolveErrorInfo& resolve_error_info,
-      const base::Optional<net::AddressList>& resolved_addresses) override;
-
-  // Performs the DNS resolution.
-  void Run(const GURL& url);
-
-  network::mojom::NetworkContext* network_context() const {
-    return network_context_;
-  }
-
- private:
-  void CreateHostResolver();
-  void OnMojoConnectionError();
-
-  network::mojom::NetworkContext* network_context_ = nullptr;  // Unowned
-  TlsProber* tls_prober_;                                      // Unowned
-  mojo::Receiver<network::mojom::ResolveHostClient> receiver_{this};
-  mojo::Remote<network::mojom::HostResolver> host_resolver_;
-};
-
-TlsProber::HostResolver::HostResolver(
-    network::mojom::NetworkContext* network_context,
-    TlsProber* tls_prober)
-    : network_context_(network_context), tls_prober_(tls_prober) {
-  DCHECK(network_context_);
-  DCHECK(tls_prober_);
-}
-
-TlsProber::HostResolver::~HostResolver() = default;
-
-void TlsProber::HostResolver::OnComplete(
-    int result,
-    const net::ResolveErrorInfo& resolve_error_info,
-    const base::Optional<net::AddressList>& resolved_addresses) {
-  receiver_.reset();
-  host_resolver_.reset();
-
-  tls_prober_->OnHostResolutionComplete(result, resolve_error_info,
-                                        resolved_addresses);
-}
-
-void TlsProber::HostResolver::Run(const GURL& url) {
-  CreateHostResolver();
-  DCHECK(host_resolver_);
-  DCHECK(!receiver_.is_bound());
-
-  network::mojom::ResolveHostParametersPtr parameters =
-      network::mojom::ResolveHostParameters::New();
-  parameters->dns_query_type = net::DnsQueryType::A;
-  parameters->source = net::HostResolverSource::DNS;
-  parameters->cache_usage =
-      network::mojom::ResolveHostParameters::CacheUsage::DISALLOWED;
-
-  host_resolver_->ResolveHost(net::HostPortPair::FromURL(url),
-                              net::NetworkIsolationKey::CreateTransient(),
-                              std::move(parameters),
-                              receiver_.BindNewPipeAndPassRemote());
-}
-
-void TlsProber::HostResolver::CreateHostResolver() {
-  network_context()->CreateHostResolver(
-      net::DnsConfigOverrides(), host_resolver_.BindNewPipeAndPassReceiver());
-  // Disconnect handler will be invoked if the network service crashes.
-  host_resolver_.set_disconnect_handler(base::BindOnce(
-      &HostResolver::OnMojoConnectionError, base::Unretained(this)));
-}
-
-void TlsProber::HostResolver::OnMojoConnectionError() {
-  OnComplete(net::ERR_NAME_NOT_RESOLVED, net::ResolveErrorInfo(net::ERR_FAILED),
-             base::nullopt);
-}
-
 TlsProber::TlsProber(NetworkContextGetter network_context_getter,
                      const GURL& url,
                      TlsProbeCompleteCallback callback)
@@ -151,9 +68,10 @@ TlsProber::TlsProber(NetworkContextGetter network_context_getter,
       network_context_getter_.Run();
   DCHECK(network_context);
 
-  host_resolver_ = std::make_unique<HostResolver>(network_context, this);
-  DCHECK(host_resolver_);
-  host_resolver_->Run(url);
+  host_resolver_ = std::make_unique<HostResolver>(
+      net::HostPortPair::FromURL(url_), network_context,
+      base::BindOnce(&TlsProber::OnHostResolutionComplete,
+                     weak_factory_.GetWeakPtr()));
 }
 
 TlsProber::TlsProber() = default;
@@ -161,15 +79,15 @@ TlsProber::TlsProber() = default;
 TlsProber::~TlsProber() = default;
 
 void TlsProber::OnHostResolutionComplete(
-    int result,
-    const net::ResolveErrorInfo& resolve_error_info,
-    const base::Optional<net::AddressList>& resolved_addresses) {
+    HostResolver::ResolutionResult& resolution_result) {
   DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
 
-  bool success = result == net::OK && !resolved_addresses->empty() &&
-                 resolved_addresses.has_value();
+  host_resolver_.reset();
+  bool success = resolution_result.result == net::OK &&
+                 !resolution_result.resolved_addresses->empty() &&
+                 resolution_result.resolved_addresses.has_value();
   if (!success) {
-    OnDone(result, ProbeExitEnum::kDnsFailure);
+    OnDone(resolution_result.result, ProbeExitEnum::kDnsFailure);
     return;
   }
 
@@ -187,7 +105,8 @@ void TlsProber::OnHostResolutionComplete(
   DCHECK(network_context);
 
   network_context->CreateTCPConnectedSocket(
-      /*local_addr=*/base::nullopt, resolved_addresses.value(),
+      /*local_addr=*/base::nullopt,
+      resolution_result.resolved_addresses.value(),
       /*options=*/nullptr,
       net::MutableNetworkTrafficAnnotationTag(GetTrafficAnnotationTag()),
       std::move(pending_receiver), /*observer=*/mojo::NullRemote(),

chrome/browser/chromeos/net/network_diagnostics/tls_prober.h

@@ -11,6 +11,7 @@
 #include "base/memory/scoped_refptr.h"
 #include "base/memory/weak_ptr.h"
 #include "base/sequenced_task_runner.h"
+#include "chrome/browser/chromeos/net/network_diagnostics/host_resolver.h"
 #include "mojo/public/cpp/bindings/remote.h"
 #include "net/base/completion_once_callback.h"
 #include "services/network/public/mojom/network_context.mojom.h"
@@ -18,10 +19,6 @@
 #include "services/network/public/mojom/tls_socket.mojom.h"
 #include "url/gurl.h"
 
-namespace net {
-class AddressList;
-}
-
 namespace chromeos {
 namespace network_diagnostics {
 
@@ -39,7 +36,6 @@ class TlsProber {
     kMojoDisconnectFailure,
     kSuccess,
   };
-  class HostResolver;
   using NetworkContextGetter =
       base::RepeatingCallback<network::mojom::NetworkContext*()>;
   using OnConnectCompleteOnUIThreadCallback = base::OnceCallback<
@@ -62,9 +58,7 @@ class TlsProber {
 
   // Processes the results of the DNS resolution done by |host_resolver_|.
   void OnHostResolutionComplete(
-      int result,
-      const net::ResolveErrorInfo& resolve_error_info,
-      const base::Optional<net::AddressList>& resolved_addresses);
+      HostResolver::ResolutionResult& resolution_result);
 
  protected:
   // Test-only constructor.

chrome/browser/chromeos/net/network_diagnostics/tls_prober_unittest.cc

@@ -4,7 +4,6 @@
 
 #include "chrome/browser/chromeos/net/network_diagnostics/tls_prober.h"
 
-#include <deque>
 #include <memory>
 #include <utility>
 #include <vector>
@@ -48,11 +47,11 @@ class TlsProberWithFakeNetworkContextTest : public ::testing::Test {
       const TlsProberWithFakeNetworkContextTest&) = delete;
 
   void InitializeProberNetworkContext(
-      std::deque<FakeHostResolver::DnsResult*> fake_dns_result,
+      std::unique_ptr<FakeHostResolver::DnsResult> fake_dns_result,
       base::Optional<net::Error> tcp_connect_code,
       base::Optional<net::Error> tls_upgrade_code) {
     fake_network_context_ = std::make_unique<FakeNetworkContext>();
-    fake_network_context_->set_fake_dns_results(std::move(fake_dns_result));
+    fake_network_context_->set_fake_dns_result(std::move(fake_dns_result));
     fake_network_context_->SetTCPConnectCode(tcp_connect_code);
     fake_network_context_->SetTLSUpgradeCode(tls_upgrade_code);
   }
@@ -65,8 +64,7 @@ class TlsProberWithFakeNetworkContextTest : public ::testing::Test {
             [](network::mojom::NetworkContext* network_context) {
               return network_context;
             },
-            static_cast<network::mojom::NetworkContext*>(
-                fake_network_context_.get())),
+            fake_network_context_.get()),
         url, std::move(callback));
   }
 
@@ -87,10 +85,9 @@ class TlsProberWithFakeNetworkContextTest : public ::testing::Test {
 
 TEST_F(TlsProberWithFakeNetworkContextTest,
        SocketConnectedAndUpgradedSuccessfully) {
-  auto dns_result = std::make_unique<FakeHostResolver::DnsResult>(
+  auto fake_dns_result = std::make_unique<FakeHostResolver::DnsResult>(
       net::OK, net::ResolveErrorInfo(net::OK),
       net::AddressList(kFakeIPAddress));
-  std::deque<FakeHostResolver::DnsResult*> fake_dns_result = {dns_result.get()};
   net::Error tcp_connect_code = net::OK;
   net::Error tls_upgrade_code = net::OK;
   InitializeProberNetworkContext(std::move(fake_dns_result), tcp_connect_code,
@@ -114,10 +111,9 @@ TEST_F(TlsProberWithFakeNetworkContextTest,
 }
 
 TEST_F(TlsProberWithFakeNetworkContextTest, FailedDnsLookup) {
-  auto dns_result = std::make_unique<FakeHostResolver::DnsResult>(
+  auto fake_dns_result = std::make_unique<FakeHostResolver::DnsResult>(
       net::ERR_NAME_NOT_RESOLVED,
       net::ResolveErrorInfo(net::ERR_NAME_NOT_RESOLVED), net::AddressList());
-  std::deque<FakeHostResolver::DnsResult*> fake_dns_result = {dns_result.get()};
   // Neither TCP connect nor TLS upgrade should not be called in this scenario.
   InitializeProberNetworkContext(std::move(fake_dns_result),
                                  /*tcp_connect_code=*/base::nullopt,
@@ -141,10 +137,9 @@ TEST_F(TlsProberWithFakeNetworkContextTest, FailedDnsLookup) {
 }
 
 TEST_F(TlsProberWithFakeNetworkContextTest, MojoDisconnectDuringDnsLookup) {
-  // Host resolution will not be successful due to Mojo disconnect.
-  std::deque<FakeHostResolver::DnsResult*> fake_dns_result = {};
-  // Neither TCP connect nor TLS upgrade should not be called in this scenario.
-  InitializeProberNetworkContext(std::move(fake_dns_result),
+  // Host resolution will not be successful due to Mojo disconnect. Neither TCP
+  // connect nor TLS upgrade should not be called in this scenario.
+  InitializeProberNetworkContext(/*fake_dns_result=*/{},
                                  /*tcp_connect_code=*/base::nullopt,
                                  /*tls_upgrade_code=*/base::nullopt);
   fake_network_context()->set_disconnect_during_host_resolution(true);
@@ -167,10 +162,9 @@ TEST_F(TlsProberWithFakeNetworkContextTest, MojoDisconnectDuringDnsLookup) {
 }
 
 TEST_F(TlsProberWithFakeNetworkContextTest, FailedTcpConnection) {
-  auto dns_result = std::make_unique<FakeHostResolver::DnsResult>(
+  auto fake_dns_result = std::make_unique<FakeHostResolver::DnsResult>(
       net::OK, net::ResolveErrorInfo(net::OK),
       net::AddressList(kFakeIPAddress));
-  std::deque<FakeHostResolver::DnsResult*> fake_dns_result = {dns_result.get()};
   net::Error tcp_connect_code = net::ERR_CONNECTION_FAILED;
   // TLS upgrade should not be called in this scenario.
   InitializeProberNetworkContext(std::move(fake_dns_result), tcp_connect_code,
@@ -194,10 +188,9 @@ TEST_F(TlsProberWithFakeNetworkContextTest, FailedTcpConnection) {
 }
 
 TEST_F(TlsProberWithFakeNetworkContextTest, FailedTlsUpgrade) {
-  auto dns_result = std::make_unique<FakeHostResolver::DnsResult>(
+  auto fake_dns_result = std::make_unique<FakeHostResolver::DnsResult>(
       net::OK, net::ResolveErrorInfo(net::OK),
       net::AddressList(kFakeIPAddress));
-  std::deque<FakeHostResolver::DnsResult*> fake_dns_result = {dns_result.get()};
   net::Error tcp_connect_code = net::OK;
   net::Error tls_upgrade_code = net::ERR_SSL_PROTOCOL_ERROR;
   InitializeProberNetworkContext(std::move(fake_dns_result), tcp_connect_code,
@@ -222,10 +215,9 @@ TEST_F(TlsProberWithFakeNetworkContextTest, FailedTlsUpgrade) {
 
 TEST_F(TlsProberWithFakeNetworkContextTest,
        MojoDisconnectedDuringTcpConnectionAttempt) {
-  auto dns_result = std::make_unique<FakeHostResolver::DnsResult>(
+  auto fake_dns_result = std::make_unique<FakeHostResolver::DnsResult>(
       net::OK, net::ResolveErrorInfo(net::OK),
       net::AddressList(kFakeIPAddress));
-  std::deque<FakeHostResolver::DnsResult*> fake_dns_result = {dns_result.get()};
   // Since the TCP connection is disconnected, no connection codes are needed.
   InitializeProberNetworkContext(std::move(fake_dns_result),
                                  /*tcp_connect_code=*/base::nullopt,
@@ -251,10 +243,9 @@ TEST_F(TlsProberWithFakeNetworkContextTest,
 
 TEST_F(TlsProberWithFakeNetworkContextTest,
        MojoDisconnectedDuringTlsUpgradeAttempt) {
-  auto dns_result = std::make_unique<FakeHostResolver::DnsResult>(
+  auto fake_dns_result = std::make_unique<FakeHostResolver::DnsResult>(
       net::OK, net::ResolveErrorInfo(net::OK),
       net::AddressList(kFakeIPAddress));
-  std::deque<FakeHostResolver::DnsResult*> fake_dns_result = {dns_result.get()};
   net::Error tcp_connect_code = net::OK;
   // TLS upgrade attempt will fail due to disconnection. |tls_upgrade_code|
   // is only populated to correctly initialize the FakeNetworkContext instance.

chrome/browser/chromeos/net/network_diagnostics/udp_prober.cc

+// Copyright 2020 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 "chrome/browser/chromeos/net/network_diagnostics/udp_prober.h"
+
+#include <utility>
+
+#include "base/logging.h"
+#include "chrome/browser/profiles/profile_manager.h"
+#include "content/public/browser/browser_task_traits.h"
+#include "content/public/browser/browser_thread.h"
+#include "mojo/public/cpp/bindings/pending_receiver.h"
+#include "mojo/public/cpp/bindings/pending_remote.h"
+#include "net/base/address_list.h"
+#include "net/base/host_port_pair.h"
+#include "net/base/net_errors.h"
+#include "net/dns/public/resolve_error_info.h"
+#include "net/traffic_annotation/network_traffic_annotation.h"
+
+namespace chromeos {
+namespace network_diagnostics {
+
+UdpProber::UdpProber(NetworkContextGetter network_context_getter,
+                     net::HostPortPair host_port_pair,
+                     base::span<const uint8_t> data,
+                     net::NetworkTrafficAnnotationTag tag,
+                     UdpProbeCompleteCallback callback)
+    : network_context_getter_(std::move(network_context_getter)),
+      host_port_pair_(host_port_pair),
+      data_(std::move(data)),
+      tag_(tag),
+      callback_(std::move(callback)) {
+  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
+  DCHECK(!data_.empty());
+  DCHECK(callback_);
+  DCHECK(!host_port_pair_.IsEmpty());
+
+  network::mojom::NetworkContext* network_context =
+      network_context_getter_.Run();
+  DCHECK(network_context);
+
+  host_resolver_ = std::make_unique<HostResolver>(
+      host_port_pair_, network_context,
+      base::BindOnce(&UdpProber::OnHostResolutionComplete,
+                     weak_factory_.GetWeakPtr()));
+}
+
+UdpProber::~UdpProber() = default;
+
+void UdpProber::OnHostResolutionComplete(
+    HostResolver::ResolutionResult& resolution_result) {
+  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
+
+  bool success = resolution_result.result == net::OK &&
+                 !resolution_result.resolved_addresses->empty() &&
+                 resolution_result.resolved_addresses.has_value();
+  if (!success) {
+    OnDone(resolution_result.result, ProbeExitEnum::kDnsFailure);
+    return;
+  }
+
+  network::mojom::NetworkContext* network_context =
+      network_context_getter_.Run();
+  DCHECK(network_context);
+
+  auto pending_receiver = udp_socket_remote_.BindNewPipeAndPassReceiver();
+  udp_socket_remote_.set_disconnect_handler(
+      base::BindOnce(&UdpProber::OnDisconnect, weak_factory_.GetWeakPtr()));
+
+  auto pending_remote =
+      udp_socket_listener_receiver_.BindNewPipeAndPassRemote();
+  udp_socket_listener_receiver_.set_disconnect_handler(
+      base::BindOnce(&UdpProber::OnDisconnect, weak_factory_.GetWeakPtr()));
+
+  network_context->CreateUDPSocket(std::move(pending_receiver),
+                                   std::move(pending_remote));
+  udp_socket_remote_->Connect(
+      resolution_result.resolved_addresses.value().front(), nullptr,
+      base::BindOnce(&UdpProber::OnConnectComplete,
+                     weak_factory_.GetWeakPtr()));
+}
+
+void UdpProber::OnConnectComplete(
+    int result,
+    const base::Optional<net::IPEndPoint>& local_addr_out) {
+  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
+
+  if (result != net::OK) {
+    OnDone(result, ProbeExitEnum::kConnectFailure);
+    return;
+  }
+  udp_socket_remote_->Send(
+      std::move(data_), net::MutableNetworkTrafficAnnotationTag(tag_),
+      base::BindOnce(&UdpProber::OnSendComplete, weak_factory_.GetWeakPtr()));
+}
+
+void UdpProber::OnSendComplete(int result) {
+  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
+
+  if (result != net::OK) {
+    OnDone(result, ProbeExitEnum::kSendFailure);
+    return;
+  }
+  udp_socket_remote_->ReceiveMore(/*num_additional_datagrams=*/1);
+}
+
+void UdpProber::OnReceived(int32_t result,
+                           const base::Optional<net::IPEndPoint>& src_ip,
+                           base::Optional<base::span<const uint8_t>> data) {
+  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
+
+  if (result != net::OK) {
+    OnDone(result, ProbeExitEnum::kNetworkErrorOnReceiveFailure);
+    return;
+  }
+
+  // The UdpProber instance is only interested in validating whether
+  // data can be received from the destination host.
+  if (!data.has_value() || data.value().empty()) {
+    // Note that net::ERR_FAILED is reported even if |result| is net::OK
+    // when no data is received.
+    OnDone(net::ERR_FAILED, ProbeExitEnum::kNoDataReceivedFailure);
+    return;
+  }
+  OnDone(net::OK, ProbeExitEnum::kSuccess);
+}
+
+void UdpProber::OnDone(int result, ProbeExitEnum probe_exit_enum) {
+  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
+
+  // Invalidate pending callbacks.
+  weak_factory_.InvalidateWeakPtrs();
+  // Destroy the socket connection.
+  udp_socket_listener_receiver_.reset();
+  udp_socket_remote_.reset();
+  // Reset the host resolver.
+  host_resolver_.reset();
+
+  std::move(callback_).Run(result, probe_exit_enum);
+}
+
+void UdpProber::OnDisconnect() {
+  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
+
+  OnDone(net::ERR_FAILED, ProbeExitEnum::kMojoDisconnectFailure);
+}
+
+}  // namespace network_diagnostics
+}  // namespace chromeos

chrome/browser/chromeos/net/network_diagnostics/udp_prober.h

+// Copyright 2020 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 CHROME_BROWSER_CHROMEOS_NET_NETWORK_DIAGNOSTICS_UDP_PROBER_H_
+#define CHROME_BROWSER_CHROMEOS_NET_NETWORK_DIAGNOSTICS_UDP_PROBER_H_
+
+#include <cstdint>
+
+#include "base/callback.h"
+#include "base/containers/span.h"
+#include "base/memory/weak_ptr.h"
+#include "base/optional.h"
+#include "chrome/browser/chromeos/net/network_diagnostics/host_resolver.h"
+#include "mojo/public/cpp/bindings/receiver.h"
+#include "mojo/public/cpp/bindings/remote.h"
+#include "net/base/address_list.h"
+#include "net/base/ip_endpoint.h"
+#include "net/dns/public/resolve_error_info.h"
+#include "net/traffic_annotation/network_traffic_annotation.h"
+#include "services/network/public/mojom/network_context.mojom.h"
+#include "services/network/public/mojom/udp_socket.mojom.h"
+#include "url/gurl.h"
+
+namespace chromeos {
+namespace network_diagnostics {
+
+// Uses a UDP socket to send data to a remote destination. After sending data,
+// the prober listens for received data. It confirms that data was received but
+// does not validate the content, hence no data is parsed. Used by network
+// diagnostic routines.
+class UdpProber : public network::mojom::UDPSocketListener {
+ public:
+  // Lists the ways a prober may end. The callback passed into the prober's
+  // constructor is invoked while exiting.
+  enum ProbeExitEnum {
+    kDnsFailure,
+    kConnectFailure,
+    kSendFailure,
+    kNetworkErrorOnReceiveFailure,
+    kMojoDisconnectFailure,
+    kNoDataReceivedFailure,
+    kSuccess,
+  };
+
+  using NetworkContextGetter =
+      base::RepeatingCallback<network::mojom::NetworkContext*()>;
+  using ConnectCallback = base::OnceCallback<
+      void(int result, const base::Optional<net::IPEndPoint>& local_addr_out)>;
+  using SendCallback = base::OnceCallback<void(int result)>;
+  using UdpProbeCompleteCallback =
+      base::OnceCallback<void(int result, ProbeExitEnum probe_exit_enum)>;
+
+  // Establishes a UDP connection and sends |data| to |host_port_pair|. The
+  // traffic sent by the prober is described by |tag|. Note that the constructor
+  // will not invoke |callback|, which is passed into UdpProber during
+  // construction. This ensures the UdpProber instance is constructed before
+  // |callback| is invoked. The UdpProber must be created on the UI thread and
+  // will invoke |callback| on the UI thread. |network_context_getter| will be
+  // invoked on the UI thread.
+  UdpProber(NetworkContextGetter network_context_getter,
+            net::HostPortPair host_port_pair,
+            base::span<const uint8_t> data,
+            net::NetworkTrafficAnnotationTag tag,
+            UdpProbeCompleteCallback callback);
+  UdpProber(const UdpProber&) = delete;
+  UdpProber& operator=(const UdpProber&) = delete;
+  ~UdpProber() override;
+
+ private:
+  // Processes the results of the DNS resolution done by |host_resolver_|.
+  void OnHostResolutionComplete(
+      HostResolver::ResolutionResult& resolution_result);
+
+  // On success, the UDP socket is connected to the destination and is ready to
+  // send data. On failure, the UdpProber exits with a failure.
+  void OnConnectComplete(int result,
+                         const base::Optional<net::IPEndPoint>& local_addr_out);
+
+  // On success, the UDP socket is ready to receive data. So long as the
+  // received data is not empty, it is considered valid. The content itself is
+  // not verified.
+  void OnSendComplete(int result);
+
+  // network::mojom::UDPSocketListener:
+  void OnReceived(int32_t result,
+                  const base::Optional<net::IPEndPoint>& src_ip,
+                  base::Optional<base::span<const uint8_t>> data) override;
+
+  // Signals the end of the probe. Manages the clean up and returns a response
+  // to the caller.
+  void OnDone(int result, ProbeExitEnum probe_exit_enum);
+
+  // Handles disconnects on the UDPSocket remote and UDPSocketListener receiver.
+  void OnDisconnect();
+
+  // Gets the active profile-specific network context.
+  NetworkContextGetter network_context_getter_;
+  // Contains the hostname and port.
+  net::HostPortPair host_port_pair_;
+  // Data to be sent to the destination.
+  base::span<const uint8_t> data_;
+  // Network annotation tag describing the socket traffic.
+  net::NetworkTrafficAnnotationTag tag_;
+  // Host resolver used for DNS lookup.
+  std::unique_ptr<HostResolver> host_resolver_;
+  // Stores the callback invoked once probe is complete or interrupted.
+  UdpProbeCompleteCallback callback_;
+  // Holds the UDPSocket remote.
+  mojo::Remote<network::mojom::UDPSocket> udp_socket_remote_;
+  // Listens to the response from hostname specified by |url_|.
+  mojo::Receiver<network::mojom::UDPSocketListener>
+      udp_socket_listener_receiver_{this};
+
+  // Must be the last member so that any callbacks taking a weak pointer to this
+  // instance are invalidated first during object destruction.
+  base::WeakPtrFactory<UdpProber> weak_factory_{this};
+};
+
+}  // namespace network_diagnostics
+}  // namespace chromeos
+
+#endif  // CHROME_BROWSER_CHROMEOS_NET_NETWORK_DIAGNOSTICS_UDP_PROBER_H_