Sm refactor IsolatedPrerenderOriginProber to better support TLS

The origin prober needs to support setting up a TLS connection, for
which resolving DNS is the first step. This CL refactors the origin
prober to make that first step possible.

Bug: 1115731
Change-Id: I67cb45d707a608770954e56f7e3688c1e49a4372
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2356267
Commit-Queue: Robert Ogden <robertogden@chromium.org>
Reviewed-by: Ryan Sturm <ryansturm@chromium.org>
Cr-Commit-Position: refs/heads/master@{#798375}

chrome/browser/prerender/isolated/isolated_prerender_origin_prober.cc

@@ -26,8 +26,10 @@ namespace {
 
 class DNSProber : public network::mojom::ResolveHostClient {
  public:
-  explicit DNSProber(
-      IsolatedPrerenderOriginProber::OnProbeResultCallback callback)
+  using OnDNSResultsCallback = base::OnceCallback<
+      void(int, const base::Optional<net::AddressList>& resolved_addresses)>;
+
+  explicit DNSProber(OnDNSResultsCallback callback)
       : callback_(std::move(callback)) {
     DCHECK(callback_);
   }
@@ -35,7 +37,7 @@ class DNSProber : public network::mojom::ResolveHostClient {
   ~DNSProber() override {
     if (callback_) {
       // Indicates some kind of mojo error. Play it safe and return no success.
-      std::move(callback_).Run(false);
+      std::move(callback_).Run(net::ERR_FAILED, base::nullopt);
     }
   }
 
@@ -47,12 +49,12 @@ class DNSProber : public network::mojom::ResolveHostClient {
       const net::ResolveErrorInfo& resolve_error_info,
       const base::Optional<net::AddressList>& resolved_addresses) override {
     if (callback_) {
-      std::move(callback_).Run(error == net::OK);
+      std::move(callback_).Run(error, resolved_addresses);
     }
   }
 
  private:
-  IsolatedPrerenderOriginProber::OnProbeResultCallback callback_;
+  OnDNSResultsCallback callback_;
 };
 
 void HTTPProbeHelper(
@@ -223,6 +225,13 @@ void IsolatedPrerenderOriginProber::Probe(const GURL& url,
 
 void IsolatedPrerenderOriginProber::DNSProbe(const GURL& url,
                                              OnProbeResultCallback callback) {
+  StartDNSResolution(url, std::move(callback), /*also_do_tls_connect=*/false);
+}
+
+void IsolatedPrerenderOriginProber::StartDNSResolution(
+    const GURL& url,
+    OnProbeResultCallback callback,
+    bool also_do_tls_connect) {
   net::NetworkIsolationKey nik =
       net::IsolationInfo::CreateForInternalRequest(url::Origin::Create(url))
           .network_isolation_key();
@@ -233,7 +242,10 @@ void IsolatedPrerenderOriginProber::DNSProbe(const GURL& url,
   resolve_host_parameters->initial_priority = net::RequestPriority::HIGHEST;
 
   mojo::PendingRemote<network::mojom::ResolveHostClient> client_remote;
-  mojo::MakeSelfOwnedReceiver(std::make_unique<DNSProber>(std::move(callback)),
+  mojo::MakeSelfOwnedReceiver(std::make_unique<DNSProber>(base::BindOnce(
+                                  &IsolatedPrerenderOriginProber::OnDNSResolved,
+                                  weak_factory_.GetWeakPtr(), url,
+                                  std::move(callback), also_do_tls_connect)),
                               client_remote.InitWithNewPipeAndPassReceiver());
 
   content::BrowserContext::GetDefaultStoragePartition(profile_)
@@ -299,3 +311,27 @@ void IsolatedPrerenderOriginProber::HTTPProbe(const GURL& url,
 
   prober_ptr->SendNowIfInactive(false /* send_only_in_foreground */);
 }
+
+void IsolatedPrerenderOriginProber::OnDNSResolved(
+    const GURL& url,
+    OnProbeResultCallback callback,
+    bool also_do_tls_connect,
+    int net_error,
+    const base::Optional<net::AddressList>& resolved_addresses) {
+  bool successful = net_error == net::OK && resolved_addresses &&
+                    !resolved_addresses->empty();
+
+  // A TLS connection needs the resolved addresses, so it also fails here.
+  if (!successful) {
+    std::move(callback).Run(false);
+    return;
+  }
+
+  if (!also_do_tls_connect) {
+    std::move(callback).Run(true);
+    return;
+  }
+
+  // TODO(robertogden): Handle also_do_tls_connect.
+  NOTREACHED();
+}

chrome/browser/prerender/isolated/isolated_prerender_origin_prober.h

@@ -6,6 +6,9 @@
 #define CHROME_BROWSER_PRERENDER_ISOLATED_ISOLATED_PRERENDER_ORIGIN_PROBER_H_
 
 #include "base/callback.h"
+#include "base/memory/weak_ptr.h"
+#include "base/optional.h"
+#include "net/base/address_list.h"
 #include "url/gurl.h"
 
 class AvailabilityProber;
@@ -46,6 +49,23 @@ class IsolatedPrerenderOriginProber {
   void DNSProbe(const GURL& url, OnProbeResultCallback callback);
   void HTTPProbe(const GURL& url, OnProbeResultCallback callback);
 
+  // Does a DNS resolution for a DNS or TLS probe, passing all the arguments to
+  // |OnDNSResolved|.
+  void StartDNSResolution(const GURL& url,
+                          OnProbeResultCallback callback,
+                          bool also_do_tls_connect);
+
+  // If the DNS resolution was successful, this will either run |callback| for a
+  // DNS probe, or start the TLS socket for a TLS probe. This is determined by
+  // |also_do_tls_connect|. If the DNS resolution failed, |callback| is run with
+  // failure.
+  void OnDNSResolved(
+      const GURL& url,
+      OnProbeResultCallback callback,
+      bool also_do_tls_connect,
+      int net_error,
+      const base::Optional<net::AddressList>& resolved_addresses);
+
   // The current profile, not owned.
   Profile* profile_;
 
@@ -54,6 +74,8 @@ class IsolatedPrerenderOriginProber {
 
   // The canary url checker.
   std::unique_ptr<AvailabilityProber> canary_check_;
+
+  base::WeakPtrFactory<IsolatedPrerenderOriginProber> weak_factory_{this};
 };
 
 #endif  // CHROME_BROWSER_PRERENDER_ISOLATED_ISOLATED_PRERENDER_ORIGIN_PROBER_H_