From bc1be2a050c8c76a310608e6dd61cf8fdfaf2336 Mon Sep 17 00:00:00 2001
From: Shane Osbourne <sosbourne@duckduckgo.com>
Date: Wed, 18 Oct 2023 14:52:13 +0100
Subject: [PATCH] accept more data from UI

---
 .../PrivacyDashboardController.swift          |  6 ++---
 .../PrivacyDashboardUserScript.swift          | 24 +++++++++++++++----
 2 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/Sources/PrivacyDashboard/PrivacyDashboardController.swift b/Sources/PrivacyDashboard/PrivacyDashboardController.swift
index aefd7452d..41cb6e3eb 100644
--- a/Sources/PrivacyDashboard/PrivacyDashboardController.swift
+++ b/Sources/PrivacyDashboard/PrivacyDashboardController.swift
@@ -27,7 +27,7 @@ public enum PrivacyDashboardOpenSettingsTarget: String {
 }
 
 public protocol PrivacyDashboardControllerDelegate: AnyObject {
-    func privacyDashboardController(_ privacyDashboardController: PrivacyDashboardController, didChangeProtectionSwitch isEnabled: Bool)
+    func privacyDashboardController(_ privacyDashboardController: PrivacyDashboardController, didChangeProtectionSwitch protectionState: ProtectionState)
     func privacyDashboardController(_ privacyDashboardController: PrivacyDashboardController, didRequestOpenUrlInNewTab url: URL)
     func privacyDashboardController(_ privacyDashboardController: PrivacyDashboardController,
                                     didRequestOpenSettings target: PrivacyDashboardOpenSettingsTarget)
@@ -239,8 +239,8 @@ extension PrivacyDashboardController: PrivacyDashboardUserScriptDelegate {
         delegate?.privacyDashboardController(self, didRequestOpenSettings: settingsTarget)
     }
 
-    func userScript(_ userScript: PrivacyDashboardUserScript, didChangeProtectionStateTo isProtected: Bool) {
-        delegate?.privacyDashboardController(self, didChangeProtectionSwitch: isProtected)
+    func userScript(_ userScript: PrivacyDashboardUserScript, didChangeProtectionState protectionState: ProtectionState) {
+        delegate?.privacyDashboardController(self, didChangeProtectionSwitch: protectionState)
     }
     
     func userScript(_ userScript: PrivacyDashboardUserScript, didRequestOpenUrlInNewTab url: URL) {
diff --git a/Sources/PrivacyDashboard/UserScript/PrivacyDashboardUserScript.swift b/Sources/PrivacyDashboard/UserScript/PrivacyDashboardUserScript.swift
index d062bf373..bdf391969 100644
--- a/Sources/PrivacyDashboard/UserScript/PrivacyDashboardUserScript.swift
+++ b/Sources/PrivacyDashboard/UserScript/PrivacyDashboardUserScript.swift
@@ -20,9 +20,10 @@ import Foundation
 import WebKit
 import TrackerRadarKit
 import UserScript
+import Common
 
 protocol PrivacyDashboardUserScriptDelegate: AnyObject {
-    func userScript(_ userScript: PrivacyDashboardUserScript, didChangeProtectionStateTo protectionState: Bool)
+    func userScript(_ userScript: PrivacyDashboardUserScript, didChangeProtectionState protectionState: ProtectionState)
     func userScript(_ userScript: PrivacyDashboardUserScript, setHeight height: Int)
     func userScriptDidRequestClosing(_ userScript: PrivacyDashboardUserScript)
     func userScriptDidRequestShowReportBrokenSite(_ userScript: PrivacyDashboardUserScript)
@@ -38,6 +39,20 @@ public enum PrivacyDashboardTheme: String, Encodable {
     case dark
 }
 
+public struct ProtectionState: Decodable {
+    public let isProtected: Bool
+    public let eventOrigin: EventOrigin
+
+    public struct EventOrigin: Decodable {
+        public let screen: EventOriginScreen
+    }
+
+    public enum EventOriginScreen: String, Decodable {
+        case primaryScreen;
+        case breakageForm;
+    }
+}
+
 final class PrivacyDashboardUserScript: NSObject, StaticUserScript {
 
     enum MessageNames: String, CaseIterable {
@@ -91,12 +106,13 @@ final class PrivacyDashboardUserScript: NSObject, StaticUserScript {
     // MARK: - JS message handlers
 
     private func handleSetProtection(message: WKScriptMessage) {
-        guard let isProtected = message.body as? Bool else {
-            assertionFailure("privacyDashboardSetProtection: expected Bool")
+
+        guard let protectionState: ProtectionState = DecodableHelper.decode(from: message.messageBody) else {
+            assertionFailure("privacyDashboardSetProtection: expected ProtectionState")
             return
         }
 
-        delegate?.userScript(self, didChangeProtectionStateTo: isProtected)
+        delegate?.userScript(self, didChangeProtectionState: protectionState)
     }
     
     private func handleSetSize(message: WKScriptMessage) {