Add assigner support to Blimp

blimp/client/session/assignment_source.cc

Index: blimp/client/session/assignment_source.cc
diff --git a/blimp/client/session/assignment_source.cc b/blimp/client/session/assignment_source.cc
index a83c66f4ac1b7ec63fb5883b679806d269677a3b..242d7839122af23edbbeef36cf87feeedd5821d2 100644
--- a/blimp/client/session/assignment_source.cc
+++ b/blimp/client/session/assignment_source.cc
@@ -5,71 +5,308 @@
 #include "blimp/client/session/assignment_source.h"
 
 #include "base/bind.h"
+#include "base/callback_helpers.h"
 #include "base/command_line.h"
+#include "base/json/json_reader.h"
+#include "base/json/json_writer.h"
 #include "base/location.h"
 #include "base/numerics/safe_conversions.h"
 #include "base/strings/string_number_conversions.h"
+#include "base/values.h"
 #include "blimp/client/app/blimp_client_switches.h"
+#include "blimp/common/protocol_version.h"
 #include "net/base/ip_address.h"
 #include "net/base/ip_endpoint.h"
+#include "net/base/load_flags.h"
+#include "net/base/net_errors.h"
+#include "net/base/url_util.h"
+#include "net/http/http_status_code.h"
+#include "net/proxy/proxy_config_service.h"
+#include "net/proxy/proxy_service.h"
+#include "net/url_request/url_fetcher.h"
+#include "net/url_request/url_request_context.h"
+#include "net/url_request/url_request_context_builder.h"
+#include "net/url_request/url_request_context_getter.h"
 
 namespace blimp {
+namespace client {
+
 namespace {
 
-// TODO(kmarshall): Take values from configuration data.
-const char kDummyClientToken[] = "MyVoiceIsMyPassport";
-const std::string kDefaultBlimpletIPAddress = "127.0.0.1";
-const uint16_t kDefaultBlimpletTCPPort = 25467;
+// Assignment request JSON keys.
+const char kProtocolVersionKey[] = "protocol_version";
+
+// Assignment response JSON keys.
+const char kClientTokenKey[] = "clientToken";
+const char kHostKey[] = "host";
+const char kPortKey[] = "port";
+const char kCertificateFingerprintKey[] = "certificateFingerprint";
+const char kCertificateKey[] = "certificate";
+
+// URL scheme constants for custom assignments.  See the '--blimplet-endpoint'
+// documentation in blimp_client_switches.cc for details.
+const char kCustomSSLScheme[] = "ssl";
+const char kCustomTCPScheme[] = "tcp";
+const char kCustomQUICScheme[] = "quic";
+
+Assignment GetCustomBlimpletAssignment() {
+  GURL url(base::CommandLine::ForCurrentProcess()->GetSwitchValueASCII(
+      switches::kBlimpletEndpoint));
 
-net::IPAddress GetBlimpletIPAddress() {
   std::string host;
-  if (base::CommandLine::ForCurrentProcess()->HasSwitch(
-          switches::kBlimpletHost)) {
-    host = base::CommandLine::ForCurrentProcess()->GetSwitchValueASCII(
-        switches::kBlimpletHost);
-  } else {
-    host = kDefaultBlimpletIPAddress;
+  int port;
+  if (url.is_empty() || !url.is_valid() || !url.has_scheme() ||
+      !net::ParseHostAndPort(url.path(), &host, &port)) {
+    return Assignment();
   }
+
   net::IPAddress ip_address;
-  if (!ip_address.AssignFromIPLiteral(host))
+  if (!ip_address.AssignFromIPLiteral(host)) {
     CHECK(false) << "Invalid BlimpletAssignment host " << host;
-  return ip_address;
-}
+  }
 
-uint16_t GetBlimpletTCPPort() {
-  if (base::CommandLine::ForCurrentProcess()->HasSwitch(
-          switches::kBlimpletTCPPort)) {
-    std::string port_str =
-        base::CommandLine::ForCurrentProcess()->GetSwitchValueASCII(
-            switches::kBlimpletTCPPort);
-    uint port_64t;
-    if (!base::StringToUint(port_str, &port_64t) ||
-        !base::IsValueInRangeForNumericType<uint16_t>(port_64t)) {
-      CHECK(false) << "Invalid BlimpletAssignment port " << port_str;
+  if (!base::IsValueInRangeForNumericType<uint16_t>(port)) {
+    CHECK(false) << "Invalid BlimpletAssignment port " << port;
+  }
+
+  Assignment::TransportProtocol protocol =
+      Assignment::TransportProtocol::UNKNOWN;
+  if (url.has_scheme()) {
+    if (url.SchemeIs(kCustomSSLScheme)) {
+      protocol = Assignment::TransportProtocol::SSL;
+    } else if (url.SchemeIs(kCustomTCPScheme)) {
+      protocol = Assignment::TransportProtocol::TCP;
+    } else if (url.SchemeIs(kCustomQUICScheme)) {
+      protocol = Assignment::TransportProtocol::QUIC;
+    } else {
+      CHECK(false) << "Invalid BlimpletAssignment scheme " << url.scheme();
     }
-    return base::checked_cast<uint16_t>(port_64t);
-  } else {
-    return kDefaultBlimpletTCPPort;
   }
+
+  Assignment assignment;
+  assignment.transport_protocol = protocol;
+  assignment.ip_endpoint = net::IPEndPoint(ip_address, port);
+  assignment.client_token = kDummyClientToken;
+  return assignment;
+}
+
+GURL GetBlimpAssignerURL() {
+  // TODO(dtrainor): Add a way to specify another assigner.
+  return GURL(kDefaultAssignerURL);
 }
 
+class SimpleURLRequestContextGetter : public net::URLRequestContextGetter {
+ public:
+  SimpleURLRequestContextGetter(
+      const scoped_refptr<base::SingleThreadTaskRunner>& io_loop_task_runner)
+      : io_loop_task_runner_(io_loop_task_runner),
+        proxy_config_service_(net::ProxyService::CreateSystemProxyConfigService(
+            io_loop_task_runner_,
+            io_loop_task_runner_)) {}
+
+  // net::URLRequestContextGetter implementation.
+  net::URLRequestContext* GetURLRequestContext() override {
+    if (!url_request_context_) {
+      net::URLRequestContextBuilder builder;
+      builder.set_proxy_config_service(std::move(proxy_config_service_));
+      builder.DisableHttpCache();
+      url_request_context_ = builder.Build();
+    }
+
+    return url_request_context_.get();
+  }
+
+  scoped_refptr<base::SingleThreadTaskRunner> GetNetworkTaskRunner()
+      const override {
+    return io_loop_task_runner_;
+  }
+
+ private:
+  ~SimpleURLRequestContextGetter() override {}
+
+  scoped_refptr<base::SingleThreadTaskRunner> io_loop_task_runner_;
+  scoped_ptr<net::URLRequestContext> url_request_context_;
+
+  // Temporary storage for the ProxyConfigService, which needs to be created on
+  // the main thread but cleared on the IO thread.  This will be built in the
+  // constructor and cleared on the IO thread.  Due to the usage of this class
+  // this is safe.
+  scoped_ptr<net::ProxyConfigService> proxy_config_service_;
+
+  DISALLOW_COPY_AND_ASSIGN(SimpleURLRequestContextGetter);
+};
+
 }  // namespace
 
-namespace client {
+Assignment::Assignment() : transport_protocol(TransportProtocol::UNKNOWN) {}
+
+Assignment::~Assignment() {}
+
+bool Assignment::is_null() const {
+  return ip_endpoint.address().empty() || ip_endpoint.port() == 0 ||
+         transport_protocol == TransportProtocol::UNKNOWN;
+}
 
 AssignmentSource::AssignmentSource(
-    const scoped_refptr<base::SingleThreadTaskRunner>& main_task_runner)
-    : main_task_runner_(main_task_runner) {}
+    const scoped_refptr<base::SingleThreadTaskRunner>& main_task_runner,
+    const scoped_refptr<base::SingleThreadTaskRunner>& io_task_runner)
+    : main_task_runner_(main_task_runner),
+      url_request_context_(new SimpleURLRequestContextGetter(io_task_runner)) {}
 
 AssignmentSource::~AssignmentSource() {}
 
-void AssignmentSource::GetAssignment(const AssignmentCallback& callback) {
+void AssignmentSource::GetAssignment(const std::string& client_auth_token,
+                                     const AssignmentCallback& callback) {
   DCHECK(main_task_runner_->BelongsToCurrentThread());
+
+  // Cancel any outstanding callback.
+  if (!callback_.is_null()) {
+    base::ResetAndReturn(&callback_)
+        .Run(AssignmentSource::Result::RESULT_SERVER_INTERRUPTED, Assignment());
+  }
+  callback_ = AssignmentCallback(callback);
+
+  Assignment assignment = GetCustomBlimpletAssignment();
+  if (!assignment.is_null()) {
+    // Post the result so that the behavior of this function is consistent.
+    main_task_runner_->PostTask(
+        FROM_HERE, base::Bind(base::ResetAndReturn(&callback_),
+                              AssignmentSource::Result::RESULT_OK, assignment));
+    return;
+  }
+
+  // Call out to the network for a real assignment.  Build the network request
+  // to hit the assigner.
+  url_fetcher_ = net::URLFetcher::Create(GetBlimpAssignerURL(),
+                                         net::URLFetcher::POST, this);
+  url_fetcher_->SetRequestContext(url_request_context_.get());
+  url_fetcher_->SetLoadFlags(net::LOAD_DO_NOT_SAVE_COOKIES |
+                             net::LOAD_DO_NOT_SEND_COOKIES);
+  url_fetcher_->AddExtraRequestHeader("Authorization: Bearer " +
+                                      client_auth_token);
+
+  // Write the JSON for the request data.
+  base::DictionaryValue dictionary;
+  dictionary.SetString(kProtocolVersionKey, blimp::kEngineVersion);
+  std::string json;
+  base::JSONWriter::Write(dictionary, &json);
+  url_fetcher_->SetUploadData("application/json", json);
+
+  url_fetcher_->Start();
+}
+
+void AssignmentSource::OnURLFetchComplete(const net::URLFetcher* source) {
+  DCHECK(main_task_runner_->BelongsToCurrentThread());
+  DCHECK(!callback_.is_null());
+  DCHECK_EQ(url_fetcher_.get(), source);
+
+  if (!source->GetStatus().is_success()) {
+    DVLOG(1) << "Assignment request failed due to network error: "
+             << net::ErrorToString(source->GetStatus().error());
+    base::ResetAndReturn(&callback_)
+        .Run(AssignmentSource::Result::RESULT_NETWORK_FAILURE, Assignment());
+    return;
+  }
+
+  switch (source->GetResponseCode()) {
+    case net::HTTP_OK:
+      ParseAssignerResponse();
+      break;
+    case net::HTTP_BAD_REQUEST:
+      base::ResetAndReturn(&callback_)
+          .Run(AssignmentSource::Result::RESULT_BAD_REQUEST, Assignment());
+      break;
+    case net::HTTP_UNAUTHORIZED:
+      base::ResetAndReturn(&callback_)
+          .Run(AssignmentSource::Result::RESULT_EXPIRED_ACCESS_TOKEN,
+               Assignment());
+      break;
+    case net::HTTP_FORBIDDEN:
+      base::ResetAndReturn(&callback_)
+          .Run(AssignmentSource::Result::RESULT_USER_INVALID, Assignment());
+      break;
+    case 429:  // Too Many Requests
+      base::ResetAndReturn(&callback_)
+          .Run(AssignmentSource::Result::RESULT_OUT_OF_VMS, Assignment());
+      break;
+    case net::HTTP_INTERNAL_SERVER_ERROR:
+      base::ResetAndReturn(&callback_)
+          .Run(AssignmentSource::Result::RESULT_SERVER_ERROR, Assignment());
+      break;
+    default:
+      base::ResetAndReturn(&callback_)
+          .Run(AssignmentSource::Result::RESULT_BAD_RESPONSE, Assignment());
+      break;
+  }
+}
+
+void AssignmentSource::ParseAssignerResponse() {
+  DCHECK(url_fetcher_.get());
+  DCHECK(url_fetcher_->GetStatus().is_success());
+  DCHECK_EQ(net::HTTP_OK, url_fetcher_->GetResponseCode());
+
+  // Grab the response from the assigner request.
+  std::string response;
+  if (!url_fetcher_->GetResponseAsString(&response)) {
+    base::ResetAndReturn(&callback_)
+        .Run(AssignmentSource::Result::RESULT_BAD_RESPONSE, Assignment());
+    return;
+  }
+
+  // Attempt to interpret the response as JSON and treat it as a dictionary.
+  scoped_ptr<base::Value> json = base::JSONReader::Read(response);
+  if (!json) {
+    base::ResetAndReturn(&callback_)
+        .Run(AssignmentSource::Result::RESULT_BAD_RESPONSE, Assignment());
+    return;
+  }
+
+  const base::DictionaryValue* dict;
+  if (!json->GetAsDictionary(&dict)) {
+    base::ResetAndReturn(&callback_)
+        .Run(AssignmentSource::Result::RESULT_BAD_RESPONSE, Assignment());
+    return;
+  }
+
+  // Validate that all the expected fields are present.
+  std::string client_token;
+  std::string host;
+  int port;
+  std::string cert_fingerprint;
+  std::string cert;
+  if (!(dict->GetString(kClientTokenKey, &client_token) &&
+        dict->GetString(kHostKey, &host) && dict->GetInteger(kPortKey, &port) &&
+        dict->GetString(kCertificateFingerprintKey, &cert_fingerprint) &&
+        dict->GetString(kCertificateKey, &cert))) {
+    base::ResetAndReturn(&callback_)
+        .Run(AssignmentSource::Result::RESULT_BAD_RESPONSE, Assignment());
+    return;
+  }
+
+  net::IPAddress ip_address;
+  if (!ip_address.AssignFromIPLiteral(host)) {
+    base::ResetAndReturn(&callback_)
+        .Run(AssignmentSource::Result::RESULT_BAD_RESPONSE, Assignment());
+    return;
+  }
+
+  if (!base::IsValueInRangeForNumericType<uint16_t>(port)) {
+    base::ResetAndReturn(&callback_)
+        .Run(AssignmentSource::Result::RESULT_BAD_RESPONSE, Assignment());
+    return;
+  }
+
   Assignment assignment;
-  assignment.ip_endpoint =
-      net::IPEndPoint(GetBlimpletIPAddress(), GetBlimpletTCPPort());
-  assignment.client_token = kDummyClientToken;
-  main_task_runner_->PostTask(FROM_HERE, base::Bind(callback, assignment));
+  // The assigner assumes SSL-only and all engines it assigns only communicate
+  // over SSL.
+  assignment.transport_protocol = Assignment::TransportProtocol::SSL;
+  assignment.ip_endpoint = net::IPEndPoint(ip_address, port);
+  assignment.client_token = client_token;
+  assignment.certificate = cert;
+  assignment.certificate_fingerprint = cert_fingerprint;
+
+  base::ResetAndReturn(&callback_)
+      .Run(AssignmentSource::Result::RESULT_OK, assignment);
 }
 
 }  // namespace client