update_engine: Allow `cros flash` on base images.

A goal of the upcoming debugd dev tools (crbug.com/403170), is to
enable a path to modify a base image such that a developer could run
`cros flash` on it.

Currently update_engine disallows custom omaha URLs and forces a hash
check for base builds, which breaks `cros flash`. This CL relaxes the
restriction slightly to allow use on a base build as long as the system
is in dev mode and the debugd dev tools are also enabled (dev tools are
currently enabled only in dev mode when there is no owner).

The check is done in update_attempter.cc, which only allows an unofficial
Omaha URL if these conditions hold true (unofficial meaning not the main
AU server or the AU test server). The other main change is
AreHashChecksMandatory() in omaha_response_handler_action.cc, which now
allows skipping hash checks for unofficial Omaha URLs.

BUG=chromium:428053
TEST=Ran unit tests, `cros flash` on base images in various states.
CQ-DEPEND=CL:227431

Change-Id: I8583ce6aa70feac8fe74b7a3992e8a4e761833c3
Reviewed-on: https://chromium-review.googlesource.com/228293
Reviewed-by: Alex Deymo <deymo@chromium.org>
Trybot-Ready: David Pursell <dpursell@chromium.org>
Commit-Queue: David Pursell <dpursell@chromium.org>
Tested-by: David Pursell <dpursell@chromium.org>
diff --git a/connection_manager_unittest.cc b/connection_manager_unittest.cc
index 576d24d..8ca1a00 100644
--- a/connection_manager_unittest.cc
+++ b/connection_manager_unittest.cc
@@ -18,6 +18,7 @@
 
 using std::set;
 using std::string;
+using testing::A;
 using testing::AnyNumber;
 using testing::Return;
 using testing::SetArgumentPointee;
@@ -112,7 +113,7 @@
   // Plumb return value into mock object.
   EXPECT_CALL(dbus_iface_, ProxyCall_0_1(kMockFlimFlamManagerProxy_,
                                          StrEq(kGetPropertiesMethod),
-                                         _, _))
+                                         _, A<GHashTable**>()))
       .WillOnce(DoAll(SetArgumentPointee<3>(manager_hash_table), Return(TRUE)));
 
   // Set other expectations.
@@ -160,7 +161,7 @@
   // Plumb return value into mock object.
   EXPECT_CALL(dbus_iface_, ProxyCall_0_1(kMockFlimFlamServiceProxy_,
                                          StrEq(kGetPropertiesMethod),
-                                         _, _))
+                                         _, A<GHashTable**>()))
       .WillOnce(DoAll(SetArgumentPointee<3>(service_hash_table), Return(TRUE)));
 
   // Set other expectations.
diff --git a/dbus_service.cc b/dbus_service.cc
index ba99586..a7ec259 100644
--- a/dbus_service.cc
+++ b/dbus_service.cc
@@ -65,15 +65,6 @@
   return etype;
 }
 
-static const char kAUTestURLRequest[] = "autest";
-// By default autest bypasses scattering. If we want to test scattering,
-// we should use autest-scheduled. The Url used is same in both cases, but
-// different params are passed to CheckForUpdate method.
-static const char kScheduledAUTestURLRequest[] = "autest-scheduled";
-
-static const char kAUTestURL[] =
-    "https://omaha.sandbox.google.com/service/update2";
-
 G_DEFINE_TYPE(UpdateEngineService, update_engine_service, G_TYPE_OBJECT)
 
 static void update_engine_service_finalize(GObject* object) {
@@ -140,41 +131,21 @@
     gchar* omaha_url,
     gint flags_as_int,
     GError **error) {
-  string update_app_version;
-  string update_omaha_url;
+  string app_version_string, omaha_url_string;
   AttemptUpdateFlags flags = static_cast<AttemptUpdateFlags>(flags_as_int);
-  bool interactive = true;
+  bool interactive = !(flags & kAttemptUpdateFlagNonInteractive);
 
-  // Only non-official (e.g., dev and test) builds can override the current
-  // version and update server URL over D-Bus. However, pointing to the
-  // hardcoded test update server URL is always allowed.
-  if (!self->system_state_->hardware()->IsOfficialBuild()) {
-    if (app_version) {
-      update_app_version = app_version;
-    }
-    if (omaha_url) {
-      update_omaha_url = omaha_url;
-    }
-  }
-  if (omaha_url) {
-    if (strcmp(omaha_url, kScheduledAUTestURLRequest) == 0) {
-      update_omaha_url = kAUTestURL;
-      // pretend that it's not user-initiated even though it is,
-      // so as to test scattering logic, etc. which get kicked off
-      // only in scheduled update checks.
-      interactive = false;
-    } else if (strcmp(omaha_url, kAUTestURLRequest) == 0) {
-      update_omaha_url = kAUTestURL;
-    }
-  }
-  if (flags & kAttemptUpdateFlagNonInteractive)
-    interactive = false;
-  LOG(INFO) << "Attempt update: app_version=\"" << update_app_version << "\" "
-            << "omaha_url=\"" << update_omaha_url << "\" "
+  if (app_version)
+    app_version_string = app_version;
+  if (omaha_url)
+    omaha_url_string = omaha_url;
+
+  LOG(INFO) << "Attempt update: app_version=\"" << app_version_string << "\" "
+            << "omaha_url=\"" << omaha_url_string << "\" "
             << "flags=0x" << std::hex << flags << " "
             << "interactive=" << (interactive? "yes" : "no");
-  self->system_state_->update_attempter()->CheckForUpdate(update_app_version,
-                                                          update_omaha_url,
+  self->system_state_->update_attempter()->CheckForUpdate(app_version_string,
+                                                          omaha_url_string,
                                                           interactive);
   return TRUE;
 }
diff --git a/dbus_wrapper_interface.h b/dbus_wrapper_interface.h
index b765600..568d4fa 100644
--- a/dbus_wrapper_interface.h
+++ b/dbus_wrapper_interface.h
@@ -48,6 +48,11 @@
                                  GError** error,
                                  GHashTable** out1) = 0;
 
+  virtual gboolean ProxyCall_0_1(DBusGProxy* proxy,
+                                 const char* method,
+                                 GError** error,
+                                 gint* out1) = 0;
+
   virtual gboolean ProxyCall_1_0(DBusGProxy* proxy,
                                  const char* method,
                                  GError** error,
diff --git a/mock_dbus_wrapper.h b/mock_dbus_wrapper.h
index a9dd0d8..9c1f7c0 100644
--- a/mock_dbus_wrapper.h
+++ b/mock_dbus_wrapper.h
@@ -26,6 +26,10 @@
                                        const char *method,
                                        GError **error,
                                        GHashTable** out1));
+  MOCK_METHOD4(ProxyCall_0_1, gboolean(DBusGProxy *proxy,
+                                       const char *method,
+                                       GError **error,
+                                       gint* out1));
   MOCK_METHOD4(ProxyCall_1_0, gboolean(DBusGProxy *proxy,
                                        const char *method,
                                        GError **error,
diff --git a/mock_omaha_request_params.h b/mock_omaha_request_params.h
index 3f1c608..f04d361 100644
--- a/mock_omaha_request_params.h
+++ b/mock_omaha_request_params.h
@@ -43,6 +43,7 @@
                                       bool is_powerwash_allowed));
   MOCK_METHOD0(UpdateDownloadChannel, void(void));
   MOCK_CONST_METHOD0(is_powerwash_allowed, bool(void));
+  MOCK_CONST_METHOD0(IsUpdateUrlOfficial, bool(void));
 
  private:
   // Wrappers to call the parent class and behave like the real object by
diff --git a/mock_update_attempter.h b/mock_update_attempter.h
index 698962a..0c20311 100644
--- a/mock_update_attempter.h
+++ b/mock_update_attempter.h
@@ -43,6 +43,8 @@
   MOCK_CONST_METHOD0(consecutive_failed_update_checks, unsigned int(void));
 
   MOCK_CONST_METHOD0(server_dictated_poll_interval, unsigned int(void));
+
+  MOCK_METHOD0(IsAnyUpdateSourceAllowed, bool(void));
 };
 
 }  // namespace chromeos_update_engine
diff --git a/omaha_request_params.cc b/omaha_request_params.cc
index 7cef85a..7a3087c 100644
--- a/omaha_request_params.cc
+++ b/omaha_request_params.cc
@@ -30,17 +30,19 @@
 
 namespace chromeos_update_engine {
 
-const char* const OmahaRequestParams::kAppId(
-    "{87efface-864d-49a5-9bb3-4b050a7c227a}");
-const char* const OmahaRequestParams::kOsPlatform("Chrome OS");
-const char* const OmahaRequestParams::kOsVersion("Indy");
-const char* const kProductionOmahaUrl(
-    "https://tools.google.com/service/update2");
+const char kProductionOmahaUrl[] =
+    "https://tools.google.com/service/update2";
+const char kAUTestOmahaUrl[] =
+    "https://omaha.sandbox.google.com/service/update2";
 
-const char* const OmahaRequestParams::kUpdateChannelKey(
-    "CHROMEOS_RELEASE_TRACK");
-const char* const OmahaRequestParams::kIsPowerwashAllowedKey(
-    "CHROMEOS_IS_POWERWASH_ALLOWED");
+const char OmahaRequestParams::kAppId[] =
+    "{87efface-864d-49a5-9bb3-4b050a7c227a}";
+const char OmahaRequestParams::kOsPlatform[] = "Chrome OS";
+const char OmahaRequestParams::kOsVersion[] = "Indy";
+const char OmahaRequestParams::kUpdateChannelKey[] = "CHROMEOS_RELEASE_TRACK";
+const char OmahaRequestParams::kIsPowerwashAllowedKey[] =
+    "CHROMEOS_IS_POWERWASH_ALLOWED";
+const char OmahaRequestParams::kAutoUpdateServerKey[] = "CHROMEOS_AUSERVER";
 
 const char* kChannelsByStability[] = {
     // This list has to be sorted from least stable to most stable channel.
@@ -103,8 +105,8 @@
   }
 
   if (in_update_url.empty())
-    update_url_ = GetLsbValue("CHROMEOS_AUSERVER", kProductionOmahaUrl, nullptr,
-                              stateful_override);
+    update_url_ = GetLsbValue(kAutoUpdateServerKey, kProductionOmahaUrl,
+                              nullptr, stateful_override);
   else
     update_url_ = in_update_url;
 
@@ -113,6 +115,12 @@
   return true;
 }
 
+bool OmahaRequestParams::IsUpdateUrlOfficial() const {
+  return (update_url_ == kAUTestOmahaUrl ||
+          update_url_ == GetLsbValue(kAutoUpdateServerKey, kProductionOmahaUrl,
+                                     nullptr, !ShouldLockDown()));
+}
+
 bool OmahaRequestParams::CollectECFWVersions() const {
   return {
       StartsWithASCII(hwid_, string("SAMS ALEX"), true) ||
diff --git a/omaha_request_params.h b/omaha_request_params.h
index 0657649..345f5c2 100644
--- a/omaha_request_params.h
+++ b/omaha_request_params.h
@@ -19,7 +19,10 @@
 namespace chromeos_update_engine {
 
 // The default "official" Omaha update URL.
-extern const char* const kProductionOmahaUrl;
+extern const char kProductionOmahaUrl[];
+
+// The autoupdate test Omaha update URL.
+extern const char kAUTestOmahaUrl[];
 
 class SystemState;
 
@@ -176,12 +179,13 @@
   virtual std::string GetAppId() const;
 
   // Suggested defaults
-  static const char* const kAppId;
-  static const char* const kOsPlatform;
-  static const char* const kOsVersion;
-  static const char* const kUpdateUrl;
-  static const char* const kUpdateChannelKey;
-  static const char* const kIsPowerwashAllowedKey;
+  static const char kAppId[];
+  static const char kOsPlatform[];
+  static const char kOsVersion[];
+  static const char kUpdateUrl[];
+  static const char kUpdateChannelKey[];
+  static const char kIsPowerwashAllowedKey[];
+  static const char kAutoUpdateServerKey[];
   static const int64_t kDefaultMinUpdateChecks = 0;
   static const int64_t kDefaultMaxUpdateChecks = 8;
 
@@ -212,6 +216,10 @@
 
   virtual bool is_powerwash_allowed() const { return is_powerwash_allowed_; }
 
+  // Check if the provided update URL is official, meaning either the default
+  // autoupdate server or the autoupdate autotest server.
+  virtual bool IsUpdateUrlOfficial() const;
+
   // For unit-tests.
   void set_root(const std::string& root);
   void set_current_channel(const std::string& channel) {
diff --git a/omaha_response_handler_action.cc b/omaha_response_handler_action.cc
index fcb312a..8458db2 100644
--- a/omaha_response_handler_action.cc
+++ b/omaha_response_handler_action.cc
@@ -134,18 +134,11 @@
 bool OmahaResponseHandlerAction::AreHashChecksMandatory(
     const OmahaResponse& response) {
   // All our internal testing uses dev server which doesn't generate
-  // metadata signatures by default. So, in order not to break
-  // image_to_live or other AU tools, we should waive the hash checks
-  // for those cases, except if the response indicates that the
-  // payload is signed.
-  //
-  // Since all internal testing is done using a dev_image or
-  // test_image, we can use that as a criteria for waiving. This
-  // criteria reduces the attack surface as opposed to waiving the
-  // checks when we're in dev mode, because we do want to enforce the
-  // hash checks when our end customers run in dev mode if they are
-  // using an official build, so that they are protected more.
-  if (!system_state_->hardware()->IsOfficialBuild()) {
+  // metadata signatures by default, so we should waive hash checks for
+  // unofficial URLs. dbus_service.cc does the security enforcement by not
+  // allowing unofficial update URLs though except in specific cases.
+  if (!system_state_->request_params()->IsUpdateUrlOfficial()) {
+    // Still do a hash check if a public key is included.
     if (!response.public_key_rsa.empty()) {
       // The autoupdate_CatchBadSignatures test checks for this string
       // in log-files. Keep in sync.
@@ -153,7 +146,7 @@
                 << "for unofficial build includes public RSA key.";
       return true;
     } else {
-      LOG(INFO) << "Waiving payload hash checks for unofficial builds";
+      LOG(INFO) << "Waiving payload hash checks for unofficial update URL.";
       return false;
     }
   }
diff --git a/omaha_response_handler_action_unittest.cc b/omaha_response_handler_action_unittest.cc
index fb99ddb..81a583a 100644
--- a/omaha_response_handler_action_unittest.cc
+++ b/omaha_response_handler_action_unittest.cc
@@ -207,14 +207,41 @@
   in.more_info_url = "http://more/info";
   in.hash = "HASHj+";
   in.size = 12;
+  FakeSystemState fake_system_state;
+  // Hash checks are always skipped for non-official update URLs.
+  EXPECT_CALL(*(fake_system_state.mock_request_params()),
+              IsUpdateUrlOfficial())
+      .WillRepeatedly(Return(true));
   InstallPlan install_plan;
-  EXPECT_TRUE(DoTest(in, "/dev/sda5", "", &install_plan));
+  EXPECT_TRUE(DoTestCommon(&fake_system_state, in, "/dev/sda5", "",
+                           &install_plan));
   EXPECT_EQ(in.payload_urls[0], install_plan.download_url);
   EXPECT_EQ(in.hash, install_plan.payload_hash);
   EXPECT_TRUE(install_plan.hash_checks_mandatory);
   EXPECT_EQ(in.version, install_plan.version);
 }
 
+TEST_F(OmahaResponseHandlerActionTest, HashChecksForUnofficialUpdateUrl) {
+  OmahaResponse in;
+  in.update_exists = true;
+  in.version = "a.b.c.d";
+  in.payload_urls.push_back("http://url.normally/needs/hash.checks.signed");
+  in.more_info_url = "http://more/info";
+  in.hash = "HASHj+";
+  in.size = 12;
+  FakeSystemState fake_system_state;
+  EXPECT_CALL(*(fake_system_state.mock_request_params()),
+              IsUpdateUrlOfficial())
+      .WillRepeatedly(Return(false));
+  InstallPlan install_plan;
+  EXPECT_TRUE(DoTestCommon(&fake_system_state, in, "/dev/sda5", "",
+                           &install_plan));
+  EXPECT_EQ(in.payload_urls[0], install_plan.download_url);
+  EXPECT_EQ(in.hash, install_plan.payload_hash);
+  EXPECT_FALSE(install_plan.hash_checks_mandatory);
+  EXPECT_EQ(in.version, install_plan.version);
+}
+
 TEST_F(OmahaResponseHandlerActionTest, HashChecksForHttpsTest) {
   OmahaResponse in;
   in.update_exists = true;
@@ -223,8 +250,13 @@
   in.more_info_url = "http://more/info";
   in.hash = "HASHj+";
   in.size = 12;
+  FakeSystemState fake_system_state;
+  EXPECT_CALL(*(fake_system_state.mock_request_params()),
+              IsUpdateUrlOfficial())
+      .WillRepeatedly(Return(true));
   InstallPlan install_plan;
-  EXPECT_TRUE(DoTest(in, "/dev/sda5", "", &install_plan));
+  EXPECT_TRUE(DoTestCommon(&fake_system_state, in, "/dev/sda5", "",
+                           &install_plan));
   EXPECT_EQ(in.payload_urls[0], install_plan.download_url);
   EXPECT_EQ(in.hash, install_plan.payload_hash);
   EXPECT_FALSE(install_plan.hash_checks_mandatory);
@@ -240,8 +272,13 @@
   in.more_info_url = "http://more/info";
   in.hash = "HASHj+";
   in.size = 12;
+  FakeSystemState fake_system_state;
+  EXPECT_CALL(*(fake_system_state.mock_request_params()),
+              IsUpdateUrlOfficial())
+      .WillRepeatedly(Return(true));
   InstallPlan install_plan;
-  EXPECT_TRUE(DoTest(in, "/dev/sda5", "", &install_plan));
+  EXPECT_TRUE(DoTestCommon(&fake_system_state, in, "/dev/sda5", "",
+                           &install_plan));
   EXPECT_EQ(in.payload_urls[0], install_plan.download_url);
   EXPECT_EQ(in.hash, install_plan.payload_hash);
   EXPECT_TRUE(install_plan.hash_checks_mandatory);
@@ -347,6 +384,10 @@
 
   FakeSystemState fake_system_state;
   OmahaRequestParams params(&fake_system_state);
+  // We're using a real OmahaRequestParams object here so we can't mock
+  // IsUpdateUrlOfficial(), but setting the update URL to the AutoUpdate test
+  // server will cause IsUpdateUrlOfficial() to return true.
+  params.set_update_url(kAUTestOmahaUrl);
   fake_system_state.set_request_params(&params);
 
   EXPECT_CALL(*fake_system_state.mock_payload_state(),
diff --git a/real_dbus_wrapper.h b/real_dbus_wrapper.h
index 1535bef..f1ef8b7 100644
--- a/real_dbus_wrapper.h
+++ b/real_dbus_wrapper.h
@@ -45,6 +45,14 @@
                              out1, G_TYPE_INVALID);
   }
 
+  gboolean ProxyCall_0_1(DBusGProxy* proxy,
+                         const char* method,
+                         GError** error,
+                         gint* out1) override {
+    return dbus_g_proxy_call(proxy, method, error, G_TYPE_INVALID,
+                             G_TYPE_INT, out1, G_TYPE_INVALID);
+  }
+
   gboolean ProxyCall_1_0(DBusGProxy* proxy,
                          const char* method,
                          GError** error,
diff --git a/update_attempter.cc b/update_attempter.cc
index f523839..a6aa2b5 100644
--- a/update_attempter.cc
+++ b/update_attempter.cc
@@ -72,8 +72,14 @@
 namespace {
 const int kMaxConsecutiveObeyProxyRequests = 20;
 
-const char* kUpdateCompletedMarker =
+const char kUpdateCompletedMarker[] =
     "/var/run/update_engine_autoupdate_completed";
+
+// By default autest bypasses scattering. If we want to test scattering,
+// use kScheduledAUTestURLRequest. The URL used is same in both cases, but
+// different params are passed to CheckForUpdate().
+const char kAUTestURLRequest[] = "autest";
+const char kScheduledAUTestURLRequest[] = "autest-scheduled";
 }  // namespace
 
 const char* UpdateStatusToString(UpdateStatus status) {
@@ -810,8 +816,26 @@
                                      const string& omaha_url,
                                      bool interactive) {
   LOG(INFO) << "Forced update check requested.";
-  forced_app_version_ = app_version;
-  forced_omaha_url_ = omaha_url;
+  forced_app_version_.clear();
+  forced_omaha_url_.clear();
+
+  // Certain conditions must be met to allow setting custom version and update
+  // server URLs. However, kScheduledAUTestURLRequest and kAUTestURLRequest are
+  // always allowed regardless of device state.
+  if (IsAnyUpdateSourceAllowed()) {
+    forced_app_version_ = app_version;
+    forced_omaha_url_ = omaha_url;
+  }
+  if (omaha_url == kScheduledAUTestURLRequest) {
+    forced_omaha_url_ = chromeos_update_engine::kAUTestOmahaUrl;
+    // Pretend that it's not user-initiated even though it is,
+    // so as to test scattering logic, etc. which get kicked off
+    // only in scheduled update checks.
+    interactive = false;
+  } else if (omaha_url == kAUTestURLRequest) {
+    forced_omaha_url_ = chromeos_update_engine::kAUTestOmahaUrl;
+  }
+
   if (forced_update_pending_callback_.get()) {
     // Make sure that a scheduling request is made prior to calling the forced
     // update pending callback.
@@ -1603,4 +1627,50 @@
           waiting_for_scheduled_check_);
 }
 
+bool UpdateAttempter::IsAnyUpdateSourceAllowed() {
+  // Non-official (dev or test) builds can always use a custom update source.
+  if (!system_state_->hardware()->IsOfficialBuild()) {
+    LOG(INFO) << "Non-official build; allowing any update source.";
+    return true;
+  }
+
+  // Official images not in devmode are never allowed a custom update source.
+  if (system_state_->hardware()->IsNormalBootMode()) {
+    LOG(INFO) << "Not in devmode; disallowing custom update sources.";
+    return false;
+  }
+
+  // Official images in devmode are allowed a custom update source iff the
+  // debugd dev tools are enabled.
+  GError* error = nullptr;
+  DBusGConnection* bus = dbus_iface_->BusGet(DBUS_BUS_SYSTEM, &error);
+  if (!bus) {
+    LOG(ERROR) << "Failed to get system bus: "
+               << utils::GetAndFreeGError(&error);
+    return false;
+  }
+
+  gint dev_features = debugd::DEV_FEATURES_DISABLED;
+  DBusGProxy* proxy = dbus_iface_->ProxyNewForName(
+      bus,
+      debugd::kDebugdServiceName,
+      debugd::kDebugdServicePath,
+      debugd::kDebugdInterface);
+  const gboolean success = dbus_iface_->ProxyCall_0_1(
+      proxy,
+      debugd::kQueryDevFeatures,
+      &error,
+      &dev_features);
+  dbus_iface_->ProxyUnref(proxy);
+
+  // Some boards may not include debugd so it's expected that this may fail,
+  // in which case we default to disallowing custom update sources.
+  if (success && !(dev_features & debugd::DEV_FEATURES_DISABLED)) {
+    LOG(INFO) << "Debugd dev tools enabled; allowing any update source.";
+    return true;
+  }
+  LOG(INFO) << "Debugd dev tools disabled; disallowing custom update sources.";
+  return false;
+}
+
 }  // namespace chromeos_update_engine
diff --git a/update_attempter.h b/update_attempter.h
index 7e13df3..7bdcab4 100644
--- a/update_attempter.h
+++ b/update_attempter.h
@@ -214,6 +214,12 @@
     forced_update_pending_callback_.reset(callback);
   }
 
+  // Returns true if we should allow updates from any source. In official builds
+  // we want to restrict updates to known safe sources, but under certain
+  // conditions it's useful to allow updating from anywhere (e.g. to allow
+  // 'cros flash' to function properly).
+  virtual bool IsAnyUpdateSourceAllowed();
+
  private:
   // Update server URL for automated lab test.
   static const char* const kTestUpdateUrl;
diff --git a/update_attempter_unittest.cc b/update_attempter_unittest.cc
index 8eec668..5d62cbe 100644
--- a/update_attempter_unittest.cc
+++ b/update_attempter_unittest.cc
@@ -9,6 +9,7 @@
 #include <memory>
 
 #include <base/files/file_util.h>
+#include <chromeos/dbus/service_constants.h>
 #include <gtest/gtest.h>
 #include <policy/libpolicy.h>
 #include <policy/mock_device_policy.h>
@@ -34,6 +35,7 @@
 using base::TimeDelta;
 using std::string;
 using std::unique_ptr;
+using testing::A;
 using testing::DoAll;
 using testing::InSequence;
 using testing::Ne;
@@ -43,6 +45,7 @@
 using testing::ReturnPointee;
 using testing::SaveArg;
 using testing::SetArgumentPointee;
+using testing::StrEq;
 using testing::_;
 
 namespace chromeos_update_engine {
@@ -80,6 +83,9 @@
   // Indicates whether ScheduleUpdates() was called.
   bool schedule_updates_called() const { return schedule_updates_called_; }
 
+  // Need to expose forced_omaha_url_ so we can test it.
+  const std::string& forced_omaha_url() const { return forced_omaha_url_; }
+
  private:
   bool schedule_updates_called_ = false;
   bool do_schedule_updates_ = true;
@@ -90,7 +96,9 @@
   UpdateAttempterTest()
       : attempter_(&fake_system_state_, &dbus_),
         mock_connection_manager(&fake_system_state_),
-        loop_(nullptr) {
+        loop_(nullptr),
+        fake_dbus_system_bus_(reinterpret_cast<DBusGConnection*>(1)),
+        fake_dbus_debugd_proxy_(reinterpret_cast<DBusGProxy*>(2)) {
     // Override system state members.
     fake_system_state_.set_connection_manager(&mock_connection_manager);
     fake_system_state_.set_update_attempter(&attempter_);
@@ -138,6 +146,16 @@
     EXPECT_CALL(*fake_system_state_.mock_payload_state(),
                 GetUsingP2PForDownloading())
         .WillRepeatedly(ReturnPointee(&actual_using_p2p_for_sharing_));
+
+    // Set up mock debugd access over the system D-Bus. ProxyCall_0_1() also
+    // needs to be mocked in any test using debugd to provide the desired value.
+    ON_CALL(dbus_, BusGet(DBUS_BUS_SYSTEM, _))
+        .WillByDefault(Return(fake_dbus_system_bus_));
+    ON_CALL(dbus_, ProxyNewForName(fake_dbus_system_bus_,
+                                   StrEq(debugd::kDebugdServiceName),
+                                   StrEq(debugd::kDebugdServicePath),
+                                   StrEq(debugd::kDebugdInterface)))
+        .WillByDefault(Return(fake_dbus_debugd_proxy_));
   }
 
   void TearDown() override {
@@ -203,6 +221,10 @@
   NiceMock<MockPrefs>* prefs_;  // Shortcut to fake_system_state_->mock_prefs().
   NiceMock<MockConnectionManager> mock_connection_manager;
   GMainLoop* loop_;
+  // fake_dbus_xxx pointers will be non-null for comparison purposes, but won't
+  // be valid objects so don't try to use them.
+  DBusGConnection* fake_dbus_system_bus_;
+  DBusGProxy* fake_dbus_debugd_proxy_;
 
   string test_dir_;
 
@@ -1040,4 +1062,69 @@
   EXPECT_EQ(boot_time.ToTimeT(), 42);
 }
 
+TEST_F(UpdateAttempterTest, AnyUpdateSourceAllowedUnofficial) {
+  fake_system_state_.fake_hardware()->SetIsOfficialBuild(false);
+  EXPECT_TRUE(attempter_.IsAnyUpdateSourceAllowed());
+}
+
+TEST_F(UpdateAttempterTest, AnyUpdateSourceAllowedOfficialDevmode) {
+  fake_system_state_.fake_hardware()->SetIsOfficialBuild(true);
+  fake_system_state_.fake_hardware()->SetIsNormalBootMode(false);
+  EXPECT_CALL(dbus_, ProxyCall_0_1(fake_dbus_debugd_proxy_,
+                                   StrEq(debugd::kQueryDevFeatures),
+                                   _, A<gint*>()))
+      .WillRepeatedly(DoAll(SetArgumentPointee<3>(0),
+                            Return(true)));
+  EXPECT_TRUE(attempter_.IsAnyUpdateSourceAllowed());
+}
+
+TEST_F(UpdateAttempterTest, AnyUpdateSourceDisallowedOfficialNormal) {
+  fake_system_state_.fake_hardware()->SetIsOfficialBuild(true);
+  fake_system_state_.fake_hardware()->SetIsNormalBootMode(true);
+  // debugd should not be queried in this case.
+  EXPECT_CALL(dbus_, ProxyCall_0_1(fake_dbus_debugd_proxy_,
+                                   StrEq(debugd::kQueryDevFeatures),
+                                   _, A<gint*>()))
+      .Times(0);
+  EXPECT_FALSE(attempter_.IsAnyUpdateSourceAllowed());
+}
+
+TEST_F(UpdateAttempterTest, AnyUpdateSourceDisallowedDebugdDisabled) {
+  using debugd::DEV_FEATURES_DISABLED;
+  fake_system_state_.fake_hardware()->SetIsOfficialBuild(true);
+  fake_system_state_.fake_hardware()->SetIsNormalBootMode(false);
+  EXPECT_CALL(dbus_, ProxyCall_0_1(fake_dbus_debugd_proxy_,
+                                   StrEq(debugd::kQueryDevFeatures),
+                                   _, A<gint*>()))
+      .WillRepeatedly(DoAll(SetArgumentPointee<3>(DEV_FEATURES_DISABLED),
+                            Return(true)));
+  EXPECT_FALSE(attempter_.IsAnyUpdateSourceAllowed());
+}
+
+TEST_F(UpdateAttempterTest, AnyUpdateSourceDisallowedDebugdFailure) {
+  fake_system_state_.fake_hardware()->SetIsOfficialBuild(true);
+  fake_system_state_.fake_hardware()->SetIsNormalBootMode(false);
+  EXPECT_CALL(dbus_, ProxyCall_0_1(fake_dbus_debugd_proxy_,
+                                   StrEq(debugd::kQueryDevFeatures),
+                                   _, A<gint*>()))
+      .WillRepeatedly(Return(false));
+  EXPECT_FALSE(attempter_.IsAnyUpdateSourceAllowed());
+}
+
+TEST_F(UpdateAttempterTest, CheckForUpdateAUTest) {
+  fake_system_state_.fake_hardware()->SetIsOfficialBuild(true);
+  fake_system_state_.fake_hardware()->SetIsNormalBootMode(true);
+  attempter_.CheckForUpdate("", "autest", true);
+  EXPECT_EQ(chromeos_update_engine::kAUTestOmahaUrl,
+            attempter_.forced_omaha_url());
+}
+
+TEST_F(UpdateAttempterTest, CheckForUpdateScheduledAUTest) {
+  fake_system_state_.fake_hardware()->SetIsOfficialBuild(true);
+  fake_system_state_.fake_hardware()->SetIsNormalBootMode(true);
+  attempter_.CheckForUpdate("", "autest-scheduled", true);
+  EXPECT_EQ(chromeos_update_engine::kAUTestOmahaUrl,
+            attempter_.forced_omaha_url());
+}
+
 }  // namespace chromeos_update_engine
diff --git a/update_manager/real_shill_provider_unittest.cc b/update_manager/real_shill_provider_unittest.cc
index bca45e2..78bdd93 100644
--- a/update_manager/real_shill_provider_unittest.cc
+++ b/update_manager/real_shill_provider_unittest.cc
@@ -25,6 +25,7 @@
 using chromeos_update_engine::test_utils::GValueNewString;
 using std::pair;
 using std::unique_ptr;
+using testing::A;
 using testing::Eq;
 using testing::Mock;
 using testing::Return;
@@ -186,7 +187,7 @@
     // Set mock expectations.
     EXPECT_CALL(mock_dbus_,
                 ProxyCall_0_1(proxy, StrEq(shill::kGetPropertiesFunction),
-                              _, _))
+                              _, A<GHashTable**>()))
         .WillOnce(DoAll(SetArgPointee<3>(g_hash_table_ref(properties)),
                         Return(true)));
 
@@ -198,7 +199,7 @@
   void SetupGetPropertiesFail(DBusGProxy* proxy) {
     EXPECT_CALL(mock_dbus_,
                 ProxyCall_0_1(proxy, StrEq(shill::kGetPropertiesFunction),
-                              _, _))
+                              _, A<GHashTable**>()))
       .WillOnce(Return(false));
   }