Android private search fix

android/java/org/chromium/chrome/browser/app/BraveActivity.java

@@ -268,11 +268,11 @@ public void onResumeWithNative() {
 
             BraveVpnUtils.reportBackgroundUsageP3A();
         }
-        Tab tab = getActivityTab();
-        if (tab != null) {
+        Profile profile = getCurrentTabModel().getProfile();
+        if (profile != null) {
             // Set proper active DSE whenever brave returns to foreground.
             // If active tab is private, set private DSE as an active DSE.
-            BraveSearchEngineUtils.updateActiveDSE(tab.isIncognito());
+            BraveSearchEngineUtils.updateActiveDSE(profile);
         }
 
         // The check on mNativeInitialized is mostly to ensure that mojo
@@ -292,11 +292,11 @@ public void onPauseWithNative() {
         if (BraveVpnUtils.isBraveVpnFeatureEnable()) {
             BraveVpnNativeWorker.getInstance().removeObserver(this);
         }
-        Tab tab = getActivityTab();
-        if (tab != null && tab.isIncognito()) {
+        Profile profile = getCurrentTabModel().getProfile();
+        if (profile != null && profile.isOffTheRecord()) {
             // Set normal DSE as an active DSE when brave goes in background
             // because currently set DSE is used by outside of brave(ex, brave search widget).
-            BraveSearchEngineUtils.updateActiveDSE(false);
+            BraveSearchEngineUtils.updateActiveDSE(profile);
         }
         super.onPauseWithNative();
     }
@@ -1191,11 +1191,13 @@ public void openBraveWalletDAppsActivity(BraveWalletDAppsActivity.ActivityType a
     private void checkForYandexSE() {
         String countryCode = Locale.getDefault().getCountry();
         if (yandexRegions.contains(countryCode)) {
-            TemplateUrl yandexTemplateUrl =
-                    BraveSearchEngineUtils.getTemplateUrlByShortName(OnboardingPrefManager.YANDEX);
+            Profile lastUsedRegularProfile = Profile.getLastUsedRegularProfile();
+            TemplateUrl yandexTemplateUrl = BraveSearchEngineUtils.getTemplateUrlByShortName(
+                    lastUsedRegularProfile, OnboardingPrefManager.YANDEX);
             if (yandexTemplateUrl != null) {
-                BraveSearchEngineUtils.setDSEPrefs(yandexTemplateUrl, false);
-                BraveSearchEngineUtils.setDSEPrefs(yandexTemplateUrl, true);
+                BraveSearchEngineUtils.setDSEPrefs(yandexTemplateUrl, lastUsedRegularProfile);
+                BraveSearchEngineUtils.setDSEPrefs(yandexTemplateUrl,
+                        lastUsedRegularProfile.getPrimaryOTRProfile(/* createIfNeeded= */ true));
             }
         }
     }

android/java/org/chromium/chrome/browser/onboarding/SearchEngineOnboardingFragment.java

@@ -29,7 +29,8 @@
 import org.chromium.chrome.R;
 import org.chromium.chrome.browser.onboarding.OnboardingPrefManager;
 import org.chromium.chrome.browser.onboarding.SearchEngineEnum;
-import org.chromium.chrome.browser.search_engines.TemplateUrlServiceFactory;
+import org.chromium.chrome.browser.profiles.Profile;
+import org.chromium.chrome.browser.search_engines.BraveTemplateUrlServiceFactory;
 import org.chromium.chrome.browser.settings.BraveSearchEngineUtils;
 import org.chromium.components.search_engines.TemplateUrl;
 import org.chromium.components.search_engines.TemplateUrlService;
@@ -45,12 +46,19 @@ public class SearchEngineOnboardingFragment extends Fragment {
 
     private Button btnSave;
 
-    private TemplateUrl selectedSearchEngine;
+    private Profile mProfile;
+    private TemplateUrl mSelectedSearchEngine;
 
     public SearchEngineOnboardingFragment() {
         // Required empty public constructor
     }
 
+    @Override
+    public void onCreate(Bundle savedInstanceState) {
+        super.onCreate(savedInstanceState);
+        mProfile = Profile.getLastUsedRegularProfile();
+    }
+
     @Override
     public View onCreateView(
         LayoutInflater inflater, ViewGroup container, Bundle savedInstanceState) {
@@ -67,10 +75,12 @@ public View onCreateView(
     }
 
     private void refreshData() {
-        TemplateUrlService templateUrlService = TemplateUrlServiceFactory.get();
+        TemplateUrlService templateUrlService =
+                BraveTemplateUrlServiceFactory.getForProfile(mProfile);
         List<TemplateUrl> templateUrls = templateUrlService.getTemplateUrls();
         TemplateUrl defaultSearchEngineTemplateUrl =
-            BraveSearchEngineUtils.getTemplateUrlByShortName(BraveSearchEngineUtils.getDSEShortName(false));
+                BraveSearchEngineUtils.getTemplateUrlByShortName(
+                        mProfile, BraveSearchEngineUtils.getDSEShortName(mProfile, false));
 
         Iterator<TemplateUrl> iterator = templateUrls.iterator();
         Set<String> templateUrlSet = new HashSet<String>();
@@ -153,19 +163,21 @@ private void setActions() {
         btnSave.setOnClickListener(new View.OnClickListener() {
             @Override
             public void onClick(View view) {
-                if (selectedSearchEngine == null) {
-                    selectedSearchEngine = BraveSearchEngineUtils.getTemplateUrlByShortName(BraveSearchEngineUtils.getDSEShortName(false));
+                if (mSelectedSearchEngine == null) {
+                    mSelectedSearchEngine = BraveSearchEngineUtils.getTemplateUrlByShortName(
+                            mProfile, BraveSearchEngineUtils.getDSEShortName(mProfile, false));
                 }
-                if (selectedSearchEngine != null) {
-                    BraveSearchEngineUtils.setDSEPrefs(selectedSearchEngine, false);
-                    BraveSearchEngineUtils.setDSEPrefs(selectedSearchEngine, true);
+                if (mSelectedSearchEngine != null) {
+                    BraveSearchEngineUtils.setDSEPrefs(mSelectedSearchEngine, mProfile);
+                    BraveSearchEngineUtils.setDSEPrefs(mSelectedSearchEngine,
+                            mProfile.getPrimaryOTRProfile(/* createIfNeeded= */ true));
                 }
                 getActivity().finish();
             }
         });
     }
 
     private void searchEngineSelected(int position, List<TemplateUrl> templateUrls) {
-        selectedSearchEngine = templateUrls.get(position);
+        mSelectedSearchEngine = templateUrls.get(position);
     }
-}
+}

android/java/org/chromium/chrome/browser/search_engines/settings/BUILD.gn → android/java/org/chromium/chrome/browser/search_engines/BUILD.gn

@@ -7,18 +7,20 @@ import("//build/config/android/rules.gni")
 
 android_library("java") {
   sources = [
-    "BraveBaseSearchEngineAdapter.java",
-    "BravePrivateSearchEnginePreference.java",
-    "BraveSearchEngineAdapter.java",
-    "BraveSearchEnginePrefHelper.java",
-    "BraveSearchEnginePreference.java",
-    "BraveStandardSearchEnginePreference.java",
+    "BraveTemplateUrlServiceFactory.java",
+    "settings/BraveBaseSearchEngineAdapter.java",
+    "settings/BravePrivateSearchEnginePreference.java",
+    "settings/BraveSearchEngineAdapter.java",
+    "settings/BraveSearchEnginePrefHelper.java",
+    "settings/BraveSearchEnginePreference.java",
+    "settings/BraveStandardSearchEnginePreference.java",
   ]
   deps = [
     ":java_resources",
     "//base:base_java",
     "//base:jni_java",
     "//build/android:build_java",
+    "//chrome/browser/profiles/android:java",
     "//chrome/browser/search_engines/android:java",
     "//components/search_engines/android:java",
     "//third_party/androidx:androidx_annotation_annotation_java",
@@ -29,7 +31,10 @@ android_library("java") {
 }
 
 generate_jni("jni_headers") {
-  sources = [ "BraveSearchEnginePrefHelper.java" ]
+  sources = [
+    "BraveTemplateUrlServiceFactory.java",
+    "settings/BraveSearchEnginePrefHelper.java",
+  ]
 }
 
 android_resources("java_resources") {

android/java/org/chromium/chrome/browser/search_engines/BraveTemplateUrlServiceFactory.java

@@ -0,0 +1,78 @@
+/* Copyright (c) 2022 The Brave Authors. All rights reserved.
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this file,
+ * You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+package org.chromium.chrome.browser.search_engines;
+
+import androidx.annotation.Nullable;
+import androidx.annotation.VisibleForTesting;
+
+import org.chromium.base.Log;
+import org.chromium.base.ThreadUtils;
+import org.chromium.base.annotations.NativeMethods;
+import org.chromium.chrome.browser.profiles.Profile;
+import org.chromium.components.search_engines.TemplateUrlService;
+
+import java.lang.reflect.InvocationTargetException;
+
+public class BraveTemplateUrlServiceFactory {
+    private static final String TAG = "BraveTemplateUrlServiceFactory";
+
+    private static TemplateUrlService sNormalTemplateUrlService;
+    private static @Nullable Profile sProfile;
+
+    private BraveTemplateUrlServiceFactory() {}
+
+    /**
+     * @return The singleton instance of {@link TemplateUrlService} for profile, based on current
+     *         active tab model.
+     * All TemplateUrlServiceFactory.get() calls are redirected here.
+     */
+    public static TemplateUrlService get() {
+        ThreadUtils.assertOnUiThread();
+
+        if (sProfile != null) return getForProfile(sProfile);
+
+        return BraveTemplateUrlServiceFactoryJni.get().getTemplateUrlService();
+    }
+
+    /**
+     * @return The singleton instance of {@link TemplateUrlService} for each profile, creating it if
+     *         necessary.
+     */
+    public static TemplateUrlService getForProfile(Profile profile) {
+        ThreadUtils.assertOnUiThread();
+
+        if (!profile.isOffTheRecord()) {
+            if (sNormalTemplateUrlService == null) {
+                sNormalTemplateUrlService =
+                        BraveTemplateUrlServiceFactoryJni.get().getTemplateUrlServiceByProfile(
+                                profile);
+            }
+            return sNormalTemplateUrlService;
+        } else {
+            // Do not cache OTR service. It will be destroyed every time
+            // a new private tab model is created
+            return BraveTemplateUrlServiceFactoryJni.get().getTemplateUrlServiceByProfile(profile);
+        }
+    }
+
+    @VisibleForTesting
+    public static void setInstanceForTesting(
+            TemplateUrlService normalService, TemplateUrlService otrService) {
+        sNormalTemplateUrlService = normalService;
+    }
+
+    public static void setCurrentProfile(@Nullable Profile profile) {
+        sProfile = profile;
+    }
+
+    // Natives interface is public to allow mocking in tests outside of
+    // org.chromium.chrome.browser.search_engines package.
+    @NativeMethods
+    public interface Natives {
+        TemplateUrlService getTemplateUrlService();
+        TemplateUrlService getTemplateUrlServiceByProfile(Profile profile);
+    }
+}