From 0091827450cf20461536406956e76460a8ba39f3 Mon Sep 17 00:00:00 2001
From: Soner Yuksel <grkmyksl@gmail.com>
Date: Wed, 28 Feb 2024 15:54:46 -0500
Subject: [PATCH] Problem ASA Custom Onboarding Linked Screen

---
 .../BrowserViewController.swift                  |  2 +-
 .../Welcome/WelcomeViewController.swift          | 16 +++++++++++-----
 2 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/Sources/Brave/Frontend/Browser/BrowserViewController/BrowserViewController.swift b/Sources/Brave/Frontend/Browser/BrowserViewController/BrowserViewController.swift
index 5ce33b5965c..f7c22e96cac 100644
--- a/Sources/Brave/Frontend/Browser/BrowserViewController/BrowserViewController.swift
+++ b/Sources/Brave/Frontend/Browser/BrowserViewController/BrowserViewController.swift
@@ -1994,7 +1994,7 @@ public class BrowserViewController: UIViewController {
   func logSecureContentState(tab: Tab, path: KVOConstants? = nil, details: String? = nil) {
     var text = """
     Tab URL: \(tab.url?.absoluteString ?? "Empty Tab URL")
-     Tab VebView URL: \(tab.webView?.url?.absoluteString ?? "Empty Webview URL")
+     Tab WebView URL: \(tab.webView?.url?.absoluteString ?? "Empty Webview URL")
      Secure State: \(tab.secureContentState.rawValue)
     """
     
diff --git a/Sources/Onboarding/Welcome/WelcomeViewController.swift b/Sources/Onboarding/Welcome/WelcomeViewController.swift
index 8e651e10872..381dbb93e88 100644
--- a/Sources/Onboarding/Welcome/WelcomeViewController.swift
+++ b/Sources/Onboarding/Welcome/WelcomeViewController.swift
@@ -407,16 +407,18 @@ public class WelcomeViewController: UIViewController {
     let attributionManager = controller.attributionManager
     
     Task { @MainActor in
+      var featureLinkType: FeatureLinkageType?
+
       do {
         if controller.p3aUtilities.isP3AEnabled {
           switch attributionManager.activeFetureLinkageLogic {
           case .campaingId:
             let featureType = try await attributionManager.handleSearchAdsFeatureLinkage()
-            attributionManager.adFeatureLinkage = featureType
+            featureLinkType = featureType
           case .reporting:
             // Handle API calls and send linkage type
             let featureType = try await controller.attributionManager.handleAdsReportingFeatureLinkage()
-            attributionManager.adFeatureLinkage = featureType
+            featureLinkType = featureType
           }
         } else {
           // p3a consent is not given
@@ -424,7 +426,11 @@ public class WelcomeViewController: UIViewController {
         }
         
         controller.calloutView.isLoading = false
-        close()
+        close() {
+          if let featureType = featureLinkType {
+            attributionManager.adFeatureLinkage = featureType
+          }
+        }
       } catch FeatureLinkageError.executionTimeout(let attributionData) {
         // Time out occurred while executing ad reports lookup
         // Ad Campaign Lookup is successful so dau server should be pinged
@@ -471,7 +477,7 @@ public class WelcomeViewController: UIViewController {
     animateToP3aState()
   }
 
-  private func close() {
+  private func close(completion: (() -> Void)? = nil) {
     var presenting: UIViewController = self
     while true {
       if let presentingController = presenting.presentingViewController {
@@ -488,7 +494,7 @@ public class WelcomeViewController: UIViewController {
     }
     
     Preferences.Onboarding.basicOnboardingProgress.value = OnboardingProgress.newTabPage.rawValue
-    presenting.dismiss(animated: false, completion: nil)
+    presenting.dismiss(animated: false, completion: completion)
   }
 }