From 820c2d72b98f76bc0d302f83b91220839a6f62bc Mon Sep 17 00:00:00 2001
From: Quentin Pubert <qpubert@google.com>
Date: Fri, 26 Jun 2026 01:33:11 -0700
Subject: [PATCH] [iOS] Enforce DownloadRestrictions policy uniformly across MIME types

Centralizes enterprise download restriction checking at the beginning of
BrowserDownloadService::OnDownloadCreated before routing tasks to
specialized tab helpers or the standard download manager.

Static restriction checking methods ShouldRestrictAllDownloads and
ShouldRestrictLocalDownloads are moved to BrowserDownloadService.
Incoming tasks are categorized via DownloadRoutingCategory enum:
- Standard downloads evaluate ShouldRestrictAllDownloads, which permits
  the download if Save to Google Drive is available even when local file
  downloading is restricted (kDownloadRestrictions == ALL_FILES).
- Specialized downloads (AR Quick Look .usdz, PassKit, MobileConfig,
  Calendar, Apple Wallet Order, VCard) evaluate
  ShouldRestrictLocalDownloads since they open in specialized system
  handlers and do not support Save to Google Drive.

When restricted, BrowserDownloadService directly presents the
restriction snackbar via DownloadManagerTabHelper on visible web states
and discards the task.

Fixed: 511816897
Change-Id: I84462216af619f0876aa211ee883598a3113ccda
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/7992457
Reviewed-by: Olivier Robin <olivierrobin@chromium.org>
Auto-Submit: Quentin Pubert <qpubert@google.com>
Commit-Queue: Quentin Pubert <qpubert@google.com>
Reviewed-by: Gauthier Ambard <gambard@chromium.org>
Cr-Commit-Position: refs/heads/main@{#1652982}
---

diff --git a/ios/chrome/browser/download/model/BUILD.gn b/ios/chrome/browser/download/model/BUILD.gn
index 82f9230..9a22f63 100644
--- a/ios/chrome/browser/download/model/BUILD.gn
+++ b/ios/chrome/browser/download/model/BUILD.gn
@@ -80,6 +80,7 @@
     "//ios/chrome/browser/download/ui:constants",
     "//ios/chrome/browser/download/ui:features",
     "//ios/chrome/browser/drive/model:drive_availability",
+    "//ios/chrome/browser/drive/model:drive_service",
     "//ios/chrome/browser/drive/model:drive_service_factory",
     "//ios/chrome/browser/drive/model:drive_tab_helper",
     "//ios/chrome/browser/drive/model:policy",
diff --git a/ios/chrome/browser/download/model/browser_download_service.h b/ios/chrome/browser/download/model/browser_download_service.h
index 2ee5acda..db288b2a 100644
--- a/ios/chrome/browser/download/model/browser_download_service.h
+++ b/ios/chrome/browser/download/model/browser_download_service.h
@@ -29,6 +29,17 @@
 
   ~BrowserDownloadService() override;
 
+  // Returns whether all downloads (both to the local filesystem and to Google
+  // Drive) should be restricted. This is more permissive than
+  // `ShouldRestrictLocalDownloads` because a download might still be allowed
+  // if it can be saved to Google Drive, even if local downloads are restricted.
+  static bool ShouldRestrictAllDownloads(web::WebState* web_state);
+
+  // Returns whether downloading to the local filesystem is restricted by
+  // policy. Unlike `ShouldRestrictAllDownloads`, this does not check if the
+  // download can be saved to Google Drive.
+  static bool ShouldRestrictLocalDownloads(web::WebState* web_state);
+
  private:
   // web::DownloadControllerDelegate overrides:
   void OnDownloadCreated(web::DownloadController*,
diff --git a/ios/chrome/browser/download/model/browser_download_service.mm b/ios/chrome/browser/download/model/browser_download_service.mm
index cdd0eb0..ebc1cb9 100644
--- a/ios/chrome/browser/download/model/browser_download_service.mm
+++ b/ios/chrome/browser/download/model/browser_download_service.mm
@@ -6,6 +6,8 @@
 
 #import "base/metrics/histogram_functions.h"
 #import "base/metrics/histogram_macros.h"
+#import "components/policy/core/common/policy_pref_names.h"
+#import "components/prefs/pref_service.h"
 #import "ios/chrome/browser/download/model/ar_quick_look_tab_helper.h"
 #import "ios/chrome/browser/download/model/download_manager_metric_names.h"
 #import "ios/chrome/browser/download/model/download_manager_tab_helper.h"
@@ -14,13 +16,72 @@
 #import "ios/chrome/browser/download/model/safari_download_tab_helper.h"
 #import "ios/chrome/browser/download/model/vcard_tab_helper.h"
 #import "ios/chrome/browser/download/ui/features.h"
+#import "ios/chrome/browser/drive/model/drive_availability.h"
+#import "ios/chrome/browser/drive/model/drive_policy.h"
+#import "ios/chrome/browser/drive/model/drive_service.h"
+#import "ios/chrome/browser/drive/model/drive_service_factory.h"
 #import "ios/chrome/browser/prerender/model/prerender_tab_helper.h"
 #import "ios/chrome/browser/shared/model/profile/profile_ios.h"
 #import "ios/chrome/browser/shared/model/utils/mime_type_util.h"
+#import "ios/chrome/browser/signin/model/authentication_service.h"
+#import "ios/chrome/browser/signin/model/authentication_service_factory.h"
+#import "ios/chrome/browser/signin/model/identity_manager_factory.h"
 #import "ios/web/public/download/download_controller.h"
 #import "ios/web/public/download/download_task.h"
 #import "net/base/url_util.h"
 
+namespace {
+
+// Categories representing the specialized tab helpers (or standard manager)
+// responsible for handling specific download tasks based on MIME type, file
+// format, feature flags, and URL scheme.
+enum class DownloadRoutingCategory {
+  kPassKit,
+  kARQuickLook,
+  kMobileConfig,
+  kCalendar,
+  kAppleWalletOrder,
+  kVCard,
+  kStandard,
+};
+
+// Evaluates the incoming download task's MIME type, filename, active feature
+// kill switches, and cryptographic schemes to determine which specialized tab
+// helper or standard download manager should handle the download.
+DownloadRoutingCategory GetDownloadRoutingCategory(
+    const web::DownloadTask* task) {
+  const std::string& mime_type = task->GetMimeType();
+  if ((mime_type == kPkPassMimeType || mime_type == kPkBundledPassMimeType) &&
+      !base::FeatureList::IsEnabled(kPassKitKillSwitch)) {
+    return DownloadRoutingCategory::kPassKit;
+  }
+  if (IsUsdzFileFormat(mime_type, task->GenerateFileName()) &&
+      !base::FeatureList::IsEnabled(kARKillSwitch)) {
+    return DownloadRoutingCategory::kARQuickLook;
+  }
+  if (mime_type == kMobileConfigurationType &&
+      (task->GetOriginalUrl().SchemeIsCryptographic() ||
+       net::IsLocalhost(task->GetOriginalUrl()))) {
+    return DownloadRoutingCategory::kMobileConfig;
+  }
+  if (mime_type == kCalendarMimeType &&
+      !base::FeatureList::IsEnabled(kCalendarKillSwitch) &&
+      task->GetOriginalUrl().SchemeIsHTTPOrHTTPS()) {
+    return DownloadRoutingCategory::kCalendar;
+  }
+  if (mime_type == kAppleWalletOrderMimeType &&
+      task->GetOriginalUrl().SchemeIsHTTPOrHTTPS()) {
+    return DownloadRoutingCategory::kAppleWalletOrder;
+  }
+  if ((mime_type == kVcardMimeType || mime_type == kXVcardMimeType) &&
+      !base::FeatureList::IsEnabled(kVCardKillSwitch)) {
+    return DownloadRoutingCategory::kVCard;
+  }
+  return DownloadRoutingCategory::kStandard;
+}
+
+}  // namespace
+
 BrowserDownloadService::BrowserDownloadService(
     web::DownloadController* download_controller)
     : download_controller_(download_controller) {
@@ -35,6 +96,42 @@
   }
 }
 
+// static
+bool BrowserDownloadService::ShouldRestrictLocalDownloads(
+    web::WebState* web_state) {
+  ProfileIOS* profile =
+      ProfileIOS::FromBrowserState(web_state->GetBrowserState());
+  PrefService* pref_service = profile->GetPrefs();
+  return static_cast<policy::DownloadRestriction>(pref_service->GetInteger(
+             policy::policy_prefs::kDownloadRestrictions)) ==
+         policy::DownloadRestriction::ALL_FILES;
+}
+
+// static
+bool BrowserDownloadService::ShouldRestrictAllDownloads(
+    web::WebState* web_state) {
+  if (!ShouldRestrictLocalDownloads(web_state)) {
+    return false;
+  }
+  ProfileIOS* profile =
+      ProfileIOS::FromBrowserState(web_state->GetBrowserState());
+  PrefService* pref_service = profile->GetPrefs();
+  drive::DriveService* drive_service =
+      drive::DriveServiceFactory::GetForProfile(profile);
+  if (!drive_service || !drive_service->IsSupported()) {
+    return true;
+  }
+  AuthenticationService* auth_service =
+      AuthenticationServiceFactory::GetForProfile(profile);
+  if (!auth_service) {
+    return true;
+  }
+  bool is_save_to_drive_available = drive::IsSaveToDriveAvailable(
+      profile->IsOffTheRecord(), IdentityManagerFactory::GetForProfile(profile),
+      drive_service, pref_service, auth_service);
+  return !is_save_to_drive_available;
+}
+
 void BrowserDownloadService::OnDownloadCreated(
     web::DownloadController* download_controller,
     web::WebState* web_state,
@@ -51,59 +148,83 @@
                                 DownloadFileUI::DownloadFilePresented,
                                 DownloadFileUI::Count);

diff --git a/ios/chrome/browser/download/model/browser_download_service_unittest.mm b/ios/chrome/browser/download/model/browser_download_service_unittest.mm
index fa50fa7..4d17e3c4 100644
--- a/ios/chrome/browser/download/model/browser_download_service_unittest.mm
+++ b/ios/chrome/browser/download/model/browser_download_service_unittest.mm
@@ -10,6 +10,8 @@
 #import "base/strings/utf_string_conversions.h"
 #import "base/test/metrics/histogram_tester.h"
 #import "base/test/scoped_feature_list.h"
+#import "components/policy/core/common/policy_pref_names.h"
+#import "components/prefs/pref_service.h"
 #import "ios/chrome/browser/download/model/ar_quick_look_tab_helper.h"
 #import "ios/chrome/browser/download/model/browser_download_service_factory.h"
 #import "ios/chrome/browser/download/model/download_manager_tab_helper.h"
@@ -17,6 +19,8 @@
 #import "ios/chrome/browser/download/model/pass_kit_tab_helper.h"
 #import "ios/chrome/browser/download/model/vcard_tab_helper.h"
 #import "ios/chrome/browser/download/ui/features.h"
+#import "ios/chrome/browser/drive/model/drive_policy.h"
+#import "ios/chrome/browser/shared/model/prefs/pref_names.h"
 #import "ios/chrome/browser/shared/model/profile/test/test_profile_ios.h"
 #import "ios/chrome/browser/shared/model/utils/mime_type_util.h"
 #import "ios/web/public/download/download_controller.h"
@@ -80,6 +84,14 @@
     tasks_.push_back(std::move(task));
   }
 
+  void ShowRestrictDownloadSnackbar() override {
+    show_restrict_download_snackbar_called_ = true;
+  }
+
+  bool show_restrict_download_snackbar_called() const {
+    return show_restrict_download_snackbar_called_;
+  }
+
   // Tasks added via Download() call.
   using DownloadTasks = std::vector<std::unique_ptr<web::DownloadTask>>;
   const DownloadTasks& tasks() const { return tasks_; }
@@ -89,6 +101,7 @@
 
  private:
   DownloadTasks tasks_;
+  bool show_restrict_download_snackbar_called_ = false;
 };
 
 }  // namespace
@@ -176,6 +189,98 @@
       1);
 }
 
+// Tests that when downloads are restricted by enterprise policy, specialized
+// downloads (such as .USDZ files) are blocked entirely and the restriction
+// snackbar is shown, instead of being routed to specialized tab helpers.
+TEST_F(BrowserDownloadServiceTest,
+       RestrictedDownloadBypassesSpecializedHelpers) {
+  profile_->GetPrefs()->SetInteger(policy::policy_prefs::kDownloadRestrictions,
+                                   3 /* ALL_FILES */);
+
+  ASSERT_TRUE(download_controller()->GetDelegate());
+  auto task = std::make_unique<web::FakeDownloadTask>(GURL(kUrl), "other");
+  task->SetGeneratedFileName(base::FilePath(kUsdzFileName));
+  web_state_.WasShown();
+
+  download_controller()->GetDelegate()->OnDownloadCreated(
+      download_controller(), &web_state_, std::move(task));
+
+  EXPECT_TRUE(ar_quick_look_tab_helper()->tasks().empty());
+  EXPECT_TRUE(pass_kit_tab_helper()->tasks().empty());
+  EXPECT_TRUE(vcard_tab_helper()->tasks().empty());
+  EXPECT_TRUE(download_manager_tab_helper()->tasks().empty());
+  EXPECT_TRUE(
+      download_manager_tab_helper()->show_restrict_download_snackbar_called());
+}
+
+// Tests that download is restricted for a visible web state when the download
+// restrictions policy is enabled.
+TEST_F(BrowserDownloadServiceTest, DownloadRestrictedForVisibleWebState) {
+  profile_->GetPrefs()->SetInteger(policy::policy_prefs::kDownloadRestrictions,
+                                   3 /* ALL_FILES */);
+  profile_->GetPrefs()->SetInteger(
+      prefs::kIosSaveToDriveDownloadManagerPolicySettings, 2 /* kDisabled */);
+
+  ASSERT_TRUE(download_controller()->GetDelegate());
+  auto task =
+      std::make_unique<web::FakeDownloadTask>(GURL(kUrl), "application/zip");
+  web_state_.WasShown();
+
+  download_controller()->GetDelegate()->OnDownloadCreated(
+      download_controller(), &web_state_, std::move(task));
+
+  EXPECT_TRUE(download_manager_tab_helper()->tasks().empty());
+  EXPECT_TRUE(
+      download_manager_tab_helper()->show_restrict_download_snackbar_called());
+}
+
+// Tests that download is restricted for a visible web state when the download
+// restrictions policy is enabled and browser is incognito.
+TEST_F(BrowserDownloadServiceTest,
+       DownloadRestrictedAndIncognitoForVisibleWebState) {
+  ProfileIOS* otr_profile = profile_->GetOffTheRecordProfile();
+  web_state_.SetBrowserState(otr_profile);
+  otr_profile->GetPrefs()->SetInteger(
+      policy::policy_prefs::kDownloadRestrictions, 3 /* ALL_FILES */);
+  otr_profile->GetPrefs()->SetInteger(
+      prefs::kIosSaveToDriveDownloadManagerPolicySettings, 1 /* kEnabled */);
+
+  ASSERT_TRUE(download_controller()->GetDelegate());
+  auto task =
+      std::make_unique<web::FakeDownloadTask>(GURL(kUrl), "application/zip");
+  web_state_.WasShown();
+
+  download_controller()->GetDelegate()->OnDownloadCreated(
+      download_controller(), &web_state_, std::move(task));
+
+  EXPECT_TRUE(download_manager_tab_helper()->tasks().empty());
+  EXPECT_TRUE(
+      download_manager_tab_helper()->show_restrict_download_snackbar_called());
+}
+
+// Tests that standard download is not restricted when download restrictions
+// policy is set to no restriction.
+TEST_F(BrowserDownloadServiceTest, NoDownloadRestrictionForVisibleWebState) {
+  profile_->GetPrefs()->SetInteger(policy::policy_prefs::kDownloadRestrictions,
+                                   0 /* NO_RESTRICTION */);
+  profile_->GetPrefs()->SetInteger(
+      prefs::kIosSaveToDriveDownloadManagerPolicySettings, 1 /* kEnabled */);
+
+  ASSERT_TRUE(download_controller()->GetDelegate());
+  auto task =
+      std::make_unique<web::FakeDownloadTask>(GURL(kUrl), "application/zip");
+  web::DownloadTask* task_ptr = task.get();
+  web_state_.WasShown();
+
+  download_controller()->GetDelegate()->OnDownloadCreated(
+      download_controller(), &web_state_, std::move(task));
+
+  ASSERT_EQ(1U, download_manager_tab_helper()->tasks().size());
+  EXPECT_EQ(task_ptr, download_manager_tab_helper()->tasks()[0].get());
+  EXPECT_FALSE(
+      download_manager_tab_helper()->show_restrict_download_snackbar_called());
+}
+
 // Tests that BrowserDownloadService uses ARQuickLookTabHelper for .REALITY
 // extension.
 TEST_F(BrowserDownloadServiceTest, RealityExtension) {
diff --git a/ios/chrome/browser/download/model/download_manager_tab_helper_unittest.mm b/ios/chrome/browser/download/model/download_manager_tab_helper_unittest.mm
index 4484b0e3..c5171748 100644
--- a/ios/chrome/browser/download/model/download_manager_tab_helper_unittest.mm
+++ b/ios/chrome/browser/download/model/download_manager_tab_helper_unittest.mm
@@ -310,99 +310,6 @@
   EXPECT_FALSE(tab_helper()->has_download_task());
 }
 
-// Tests that download is restricted for a visible web state when the download
-// restrictions policy is enabled and the Save to Drive policy is disabled. The
-// test verifies that the delegate state remains nil. Additionally, the test
-// checks that a snackbar is displayed to the user.
-TEST_F(DownloadManagerTabHelperTest, DownloadRestrictedForVisibleWebState) {
-  SignIn();
-  PrefService* pref_service = profile_.get()->GetPrefs();
-  pref_service->SetInteger(
-      policy::policy_prefs::kDownloadRestrictions,
-      static_cast<int>(policy::DownloadRestriction::ALL_FILES));
-  pref_service->SetInteger(
-      prefs::kIosSaveToDriveDownloadManagerPolicySettings,
-      static_cast<int>(SaveToDrivePolicySettings::kDisabled));
-
-  web_state_->WasShown();
-  id mock_snackbar_command_handler_ =
-      OCMProtocolMock(@protocol(SnackbarCommands));
-
-  OCMExpect([mock_snackbar_command_handler_ showSnackbarWithMessage:[OCMArg any]
-                                                         buttonText:[OCMArg any]
-                                                      messageAction:nil
-                                                   completionAction:nil]);
-  ASSERT_FALSE(delegate_.state);
-  std::unique_ptr<web::FakeDownloadTask> task =
-      CreateFakeDownloadTask(GURL(kUrl), kMimeType);
-  tab_helper()->SetSnackbarHandler(mock_snackbar_command_handler_);
-  tab_helper()->SetCurrentDownload(std::move(task));
-  ASSERT_FALSE(delegate_.state);
-  EXPECT_OCMOCK_VERIFY(mock_snackbar_command_handler_);
-}
-
-// Tests that download is restricted for a visible web state when the download
-// restrictions policy is enabled and browser is incognito. The test verifies
-// that the delegate state remains nil. Additionally, the test checks
-// that a snackbar is displayed to the user.
-TEST_F(DownloadManagerTabHelperTest,
-       DownloadRestrictedAndIncognitoForVisibleWebState) {
-  web_state_->SetBrowserState(profile_->GetOffTheRecordProfile());
-  SignIn();
-  PrefService* pref_service = profile_.get()->GetPrefs();
-  pref_service->SetInteger(
-      policy::policy_prefs::kDownloadRestrictions,
-      static_cast<int>(policy::DownloadRestriction::ALL_FILES));
-  pref_service->SetInteger(
-      prefs::kIosSaveToDriveDownloadManagerPolicySettings,
-      static_cast<int>(SaveToDrivePolicySettings::kEnabled));
-  web_state_->WasShown();
-  id mock_snackbar_command_handler_ =
-      OCMProtocolMock(@protocol(SnackbarCommands));
-
-  OCMExpect([mock_snackbar_command_handler_ showSnackbarWithMessage:[OCMArg any]
-                                                         buttonText:[OCMArg any]
-                                                      messageAction:nil
-                                                   completionAction:nil]);
-  ASSERT_FALSE(delegate_.state);
-  std::unique_ptr<web::FakeDownloadTask> task =
-      CreateFakeDownloadTask(GURL(kUrl), kMimeType);
-  tab_helper()->SetSnackbarHandler(mock_snackbar_command_handler_);
-  tab_helper()->SetCurrentDownload(std::move(task));
-  ASSERT_FALSE(delegate_.state);
-  EXPECT_OCMOCK_VERIFY(mock_snackbar_command_handler_);
-}
-
-// Tests that download is not restricted for a visible web state when the
-// download restrictions policy is enabled but the Save to Drive policy is also
-// enable. The test verifies that the delegate state is set. Additionally, the
-// test checks that a snackbar is not displayed to the user.
-TEST_F(DownloadManagerTabHelperTest, NoDownloadRestrictionForVisibleWebState) {
-  SignIn();
-  PrefService* pref_service = profile_.get()->GetPrefs();
-  pref_service->SetInteger(
-      policy::policy_prefs::kDownloadRestrictions,
-      static_cast<int>(policy::DownloadRestriction::ALL_FILES));
-  pref_service->SetInteger(
-      prefs::kIosSaveToDriveDownloadManagerPolicySettings,
-      static_cast<int>(SaveToDrivePolicySettings::kEnabled));
-  web_state_->WasShown();
-  id mock_snackbar_command_handler_ =
-      OCMProtocolMock(@protocol(SnackbarCommands));
-
-  OCMReject([mock_snackbar_command_handler_ showSnackbarWithMessage:[OCMArg any]
-                                                         buttonText:[OCMArg any]
-                                                      messageAction:nil
-                                                   completionAction:nil]);
-  ASSERT_FALSE(delegate_.state);
-  std::unique_ptr<web::FakeDownloadTask> task =
-      CreateFakeDownloadTask(GURL(kUrl), kMimeType);
-  tab_helper()->SetSnackbarHandler(mock_snackbar_command_handler_);
-  tab_helper()->SetCurrentDownload(std::move(task));
-  ASSERT_TRUE(delegate_.state);
-  EXPECT_OCMOCK_VERIFY(mock_snackbar_command_handler_);
-}
-
 // Tests that after a download task is complete, it finishes by moving the file.
 // This verifies that when the feature is disabled, the scan result is SUCCESS
 // and it proceeds without a warning dialog.