network_diagnostics: Add a timeout to UdpProber

Since UDP connections do not guarantee a response from the other side,
include a timeout to avoid infinitely waiting.

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

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

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

@@ -132,5 +132,23 @@ void FakeNetworkContext::SetDisconnectDuringUdpReceiveAttempt(bool disconnect) {
   fake_udp_socket_->set_disconnect_during_udp_receive_attempt(disconnect);
 }
 
+void FakeNetworkContext::SetTaskEnvironmentForTesting(
+    content::BrowserTaskEnvironment* task_environment) {
+  fake_udp_socket_->set_task_environment_for_testing(task_environment);
+}
+
+void FakeNetworkContext::SetUdpConnectionDelay(
+    base::TimeDelta connection_delay) {
+  fake_udp_socket_->set_udp_connection_delay(connection_delay);
+}
+
+void FakeNetworkContext::SetUdpSendDelay(base::TimeDelta send_delay) {
+  fake_udp_socket_->set_udp_send_delay(send_delay);
+}
+
+void FakeNetworkContext::SetUdpReceiveDelay(base::TimeDelta receive_delay) {
+  fake_udp_socket_->set_udp_receive_delay(receive_delay);
+}
+
 }  // namespace network_diagnostics
 }  // namespace chromeos

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

@@ -72,6 +72,19 @@ class FakeNetworkContext : public network::TestNetworkContext {
   // confirmation, but before receiving any data.
   void SetDisconnectDuringUdpReceiveAttempt(bool disconnect);
 
+  // Sets the task environment used in testing. Used to fast forward the clock.
+  void SetTaskEnvironmentForTesting(
+      content::BrowserTaskEnvironment* task_environment);
+
+  // Sets the UDP connection delay.
+  void SetUdpConnectionDelay(base::TimeDelta connection_delay);
+
+  // Sets the UDP send delay.
+  void SetUdpSendDelay(base::TimeDelta send_delay);
+
+  // Sets the UDP receive delay.
+  void SetUdpReceiveDelay(base::TimeDelta receive_delay);
+
   // 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) {

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

@@ -21,6 +21,9 @@ FakeUdpSocket::~FakeUdpSocket() = default;
 void FakeUdpSocket::Connect(const net::IPEndPoint& remote_addr,
                             network::mojom::UDPSocketOptionsPtr options,
                             ConnectCallback callback) {
+  DCHECK(task_environment_);
+
+  task_environment_->FastForwardBy(connection_delay_);
   if (mojo_disconnect_on_connect_) {
     receiver_.reset();
     return;
@@ -61,7 +64,9 @@ void FakeUdpSocket::LeaveGroup(const net::IPAddress& group_address,
 
 void FakeUdpSocket::ReceiveMore(uint32_t num_additional_datagrams) {
   DCHECK(remote_.is_bound());
+  DCHECK(task_environment_);
 
+  task_environment_->FastForwardBy(receive_delay_);
   if (mojo_disconnect_on_receive_) {
     remote_.reset();
     return;
@@ -88,6 +93,9 @@ void FakeUdpSocket::Send(
     base::span<const uint8_t> data,
     const net::MutableNetworkTrafficAnnotationTag& traffic_annotation,
     SendCallback callback) {
+  DCHECK(task_environment_);
+
+  task_environment_->FastForwardBy(send_delay_);
   if (mojo_disconnect_on_send_) {
     receiver_.reset();
     return;

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

@@ -5,6 +5,8 @@
 #ifndef CHROME_BROWSER_CHROMEOS_NET_NETWORK_DIAGNOSTICS_FAKE_UDP_SOCKET_H_
 #define CHROME_BROWSER_CHROMEOS_NET_NETWORK_DIAGNOSTICS_FAKE_UDP_SOCKET_H_
 
+#include "base/time/time.h"
+#include "content/public/test/browser_task_environment.h"
 #include "mojo/public/cpp/bindings/pending_receiver.h"
 #include "mojo/public/cpp/bindings/pending_remote.h"
 #include "mojo/public/cpp/bindings/receiver.h"
@@ -92,6 +94,23 @@ class FakeUdpSocket : public network::mojom::UDPSocket {
     mojo_disconnect_on_receive_ = disconnect;
   }
 
+  void set_task_environment_for_testing(
+      content::BrowserTaskEnvironment* task_environment) {
+    task_environment_ = task_environment;
+  }
+
+  void set_udp_connection_delay(base::TimeDelta connection_delay) {
+    connection_delay_ = connection_delay;
+  }
+
+  void set_udp_send_delay(base::TimeDelta send_delay) {
+    send_delay_ = send_delay;
+  }
+
+  void set_udp_receive_delay(base::TimeDelta receive_delay) {
+    receive_delay_ = receive_delay;
+  }
+
  private:
   mojo::Receiver<network::mojom::UDPSocket> receiver_{this};
   mojo::Remote<network::mojom::UDPSocketListener> remote_;
@@ -102,6 +121,10 @@ class FakeUdpSocket : public network::mojom::UDPSocket {
   bool mojo_disconnect_on_connect_ = false;
   bool mojo_disconnect_on_send_ = false;
   bool mojo_disconnect_on_receive_ = false;
+  content::BrowserTaskEnvironment* task_environment_;
+  base::TimeDelta connection_delay_;
+  base::TimeDelta send_delay_;
+  base::TimeDelta receive_delay_;
 };
 
 }  // namespace network_diagnostics

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

@@ -25,11 +25,13 @@ UdpProber::UdpProber(NetworkContextGetter network_context_getter,
                      net::HostPortPair host_port_pair,
                      base::span<const uint8_t> data,
                      net::NetworkTrafficAnnotationTag tag,
+                     base::TimeDelta timeout_after_host_resolution,
                      UdpProbeCompleteCallback callback)
     : network_context_getter_(std::move(network_context_getter)),
       host_port_pair_(host_port_pair),
       data_(std::move(data)),
       tag_(tag),
+      timeout_after_host_resolution_(timeout_after_host_resolution),
       callback_(std::move(callback)) {
   DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
   DCHECK(!data_.empty());
@@ -67,6 +69,7 @@ void UdpProber::OnHostResolutionComplete(
   auto pending_receiver = udp_socket_remote_.BindNewPipeAndPassReceiver();
   udp_socket_remote_.set_disconnect_handler(
       base::BindOnce(&UdpProber::OnDisconnect, weak_factory_.GetWeakPtr()));
+  DCHECK(udp_socket_remote_.is_bound());
 
   auto pending_remote =
       udp_socket_listener_receiver_.BindNewPipeAndPassRemote();
@@ -75,6 +78,9 @@ void UdpProber::OnHostResolutionComplete(
 
   network_context->CreateUDPSocket(std::move(pending_receiver),
                                    std::move(pending_remote));
+  timer_.Start(FROM_HERE, timeout_after_host_resolution_,
+               base::BindOnce(&UdpProber::OnDone, weak_factory_.GetWeakPtr(),
+                              net::ERR_TIMED_OUT, ProbeExitEnum::kTimeout));
   udp_socket_remote_->Connect(
       resolution_result.resolved_addresses.value().front(), nullptr,
       base::BindOnce(&UdpProber::OnConnectComplete,
@@ -85,7 +91,6 @@ 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;

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

@@ -11,6 +11,8 @@
 #include "base/containers/span.h"
 #include "base/memory/weak_ptr.h"
 #include "base/optional.h"
+#include "base/time/time.h"
+#include "base/timer/timer.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"
@@ -40,6 +42,7 @@ class UdpProber : public network::mojom::UDPSocketListener {
     kNetworkErrorOnReceiveFailure,
     kMojoDisconnectFailure,
     kNoDataReceivedFailure,
+    kTimeout,
     kSuccess,
   };
 
@@ -52,16 +55,22 @@ class UdpProber : public network::mojom::UDPSocketListener {
       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.
+  // traffic sent by the prober is described by |tag|. Since there is no
+  // guarantee the host specified by |host_port_pair| will respond to a UDP
+  // request, the prober will timeout with a failure after
+  // |timeout_after_host_resolution|. As such, the prober will run no longer
+  // than the time it takes to perform host resolution +
+  // |timeout_after_host_resolution|. 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,
+            base::TimeDelta timeout_after_host_resolution,
             UdpProbeCompleteCallback callback);
   UdpProber(const UdpProber&) = delete;
   UdpProber& operator=(const UdpProber&) = delete;
@@ -102,6 +111,10 @@ class UdpProber : public network::mojom::UDPSocketListener {
   base::span<const uint8_t> data_;
   // Network annotation tag describing the socket traffic.
   net::NetworkTrafficAnnotationTag tag_;
+  // Represents the time after host resolution.
+  base::TimeDelta timeout_after_host_resolution_;
+  // Times the prober.
+  base::OneShotTimer timer_;
   // Host resolver used for DNS lookup.
   std::unique_ptr<HostResolver> host_resolver_;
   // Stores the callback invoked once probe is complete or interrupted.