From aa689ad625cb97b87e8cddc2bc4ec5564c9904b9 Mon Sep 17 00:00:00 2001 From: Brad Slayter Date: Mon, 11 Dec 2023 13:56:38 -0600 Subject: [PATCH 1/3] Use URLRequest instead of URL on link rewrites to perserve headers --- .../LinkProtection/LinkProtection.swift | 29 ++++++++++++------- 1 file changed, 18 insertions(+), 11 deletions(-) diff --git a/Sources/BrowserServicesKit/LinkProtection/LinkProtection.swift b/Sources/BrowserServicesKit/LinkProtection/LinkProtection.swift index 7dd8b0986..d5cac7780 100644 --- a/Sources/BrowserServicesKit/LinkProtection/LinkProtection.swift +++ b/Sources/BrowserServicesKit/LinkProtection/LinkProtection.swift @@ -35,7 +35,13 @@ public struct LinkProtection { contentBlockingManager: contentBlockingManager, errorReporting: errorReporting) } - + + private func getNewRequest(changingUrl url: URL, inRequest request: URLRequest) -> URLRequest { + var newRequest = request + newRequest.url = url + return newRequest + } + public mutating func setMainFrameUrl(_ url: URL?) { mainFrameUrl = url } @@ -77,11 +83,12 @@ public struct LinkProtection { // swiftlint:disable function_parameter_count public func requestTrackingLinkRewrite(initiatingURL: URL?, - destinationURL: URL, + destinationRequest: URLRequest, onStartExtracting: () -> Void, onFinishExtracting: @escaping () -> Void, - onLinkRewrite: @escaping (URL) -> Void, + onLinkRewrite: @escaping (URLRequest) -> Void, policyDecisionHandler: @escaping (Bool) -> Void) -> Bool { + let destinationURL = destinationRequest.url if let mainFrameUrl = mainFrameUrl, destinationURL != mainFrameUrl { // If mainFrameUrl is set and is different from destinationURL we will assume this is a redirect // We do not rewrite redirects due to breakage concerns @@ -91,7 +98,7 @@ public struct LinkProtection { var didRewriteLink = false if let newURL = linkCleaner.extractCanonicalFromAMPLink(initiator: initiatingURL, destination: destinationURL) { policyDecisionHandler(false) - onLinkRewrite(newURL) + onLinkRewrite(getNewRequest(changingUrl: newURL, inRequest: destinationRequest)) didRewriteLink = true } else if ampExtractor.urlContainsAMPKeyword(destinationURL) { onStartExtracting() @@ -103,13 +110,13 @@ public struct LinkProtection { } policyDecisionHandler(false) - onLinkRewrite(canonical) + onLinkRewrite(getNewRequest(changingUrl: canonical, inRequest: destinationRequest)) } didRewriteLink = true } else if let newURL = linkCleaner.cleanTrackingParameters(initiator: initiatingURL, url: destinationURL) { if newURL != destinationURL { policyDecisionHandler(false) - onLinkRewrite(newURL) + onLinkRewrite(getNewRequest(changingUrl: newURL, inRequest: destinationRequest)) didRewriteLink = true } } @@ -121,10 +128,10 @@ public struct LinkProtection { navigationAction: WKNavigationAction, onStartExtracting: () -> Void, onFinishExtracting: @escaping () -> Void, - onLinkRewrite: @escaping (URL, WKNavigationAction) -> Void, + onLinkRewrite: @escaping (URLRequest, WKNavigationAction) -> Void, policyDecisionHandler: @escaping (WKNavigationActionPolicy) -> Void) -> Bool { requestTrackingLinkRewrite(initiatingURL: initiatingURL, - destinationURL: navigationAction.request.url!, + destinationRequest: navigationAction.request, onStartExtracting: onStartExtracting, onFinishExtracting: onFinishExtracting, onLinkRewrite: { onLinkRewrite($0, navigationAction) }, @@ -134,13 +141,13 @@ public struct LinkProtection { @MainActor public func requestTrackingLinkRewrite(initiatingURL: URL?, - destinationURL: URL, + destinationRequest: URLRequest, onStartExtracting: () -> Void, onFinishExtracting: @escaping () -> Void, - onLinkRewrite: @escaping (URL) -> Void) async -> Bool? { + onLinkRewrite: @escaping (URLRequest) -> Void) async -> Bool? { await withCheckedContinuation { continuation in let didRewriteLink = requestTrackingLinkRewrite(initiatingURL: initiatingURL, - destinationURL: destinationURL, + destinationRequest: destinationRequest, onStartExtracting: onStartExtracting, onFinishExtracting: onFinishExtracting, onLinkRewrite: onLinkRewrite) { navigationActionPolicy in From 7998cd1cd6e92281a89c7271dcdb521fb8e14f46 Mon Sep 17 00:00:00 2001 From: Brad Slayter Date: Thu, 21 Dec 2023 12:31:23 -0600 Subject: [PATCH 2/3] Method rename --- .../LinkProtection/LinkProtection.swift | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/Sources/BrowserServicesKit/LinkProtection/LinkProtection.swift b/Sources/BrowserServicesKit/LinkProtection/LinkProtection.swift index aebb2a3de..7d98a4a4b 100644 --- a/Sources/BrowserServicesKit/LinkProtection/LinkProtection.swift +++ b/Sources/BrowserServicesKit/LinkProtection/LinkProtection.swift @@ -36,7 +36,7 @@ public struct LinkProtection { errorReporting: errorReporting) } - private func getNewRequest(changingUrl url: URL, inRequest request: URLRequest) -> URLRequest { + private func makeNewRequest(changingUrl url: URL, inRequest request: URLRequest) -> URLRequest { var newRequest = request newRequest.url = url return newRequest @@ -98,7 +98,7 @@ public struct LinkProtection { var didRewriteLink = false if let newURL = linkCleaner.extractCanonicalFromAMPLink(initiator: initiatingURL, destination: destinationURL) { policyDecisionHandler(false) - onLinkRewrite(getNewRequest(changingUrl: newURL, inRequest: destinationRequest)) + onLinkRewrite(makeNewRequest(changingUrl: newURL, inRequest: destinationRequest)) didRewriteLink = true } else if ampExtractor.urlContainsAMPKeyword(destinationURL) { onStartExtracting() @@ -110,13 +110,13 @@ public struct LinkProtection { } policyDecisionHandler(false) - onLinkRewrite(getNewRequest(changingUrl: canonical, inRequest: destinationRequest)) + onLinkRewrite(makeNewRequest(changingUrl: canonical, inRequest: destinationRequest)) } didRewriteLink = true } else if let newURL = linkCleaner.cleanTrackingParameters(initiator: initiatingURL, url: destinationURL) { if newURL != destinationURL { policyDecisionHandler(false) - onLinkRewrite(getNewRequest(changingUrl: newURL, inRequest: destinationRequest)) + onLinkRewrite(makeNewRequest(changingUrl: newURL, inRequest: destinationRequest)) didRewriteLink = true } } From 38d7f2c679e1c4138b70ffb936dce90d4344d0e0 Mon Sep 17 00:00:00 2001 From: Brad Slayter Date: Thu, 21 Dec 2023 14:14:35 -0600 Subject: [PATCH 3/3] Add getURLRequest method --- .../LinkProtection/LinkProtection.swift | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/Sources/BrowserServicesKit/LinkProtection/LinkProtection.swift b/Sources/BrowserServicesKit/LinkProtection/LinkProtection.swift index 7d98a4a4b..fe23aeb1a 100644 --- a/Sources/BrowserServicesKit/LinkProtection/LinkProtection.swift +++ b/Sources/BrowserServicesKit/LinkProtection/LinkProtection.swift @@ -46,6 +46,16 @@ public struct LinkProtection { mainFrameUrl = url } + public func getCleanURLRequest(from urlRequest: URLRequest, + onStartExtracting: () -> Void, + onFinishExtracting: @escaping () -> Void, + completion: @escaping (URLRequest) -> Void) { + getCleanURL(from: urlRequest.url!, onStartExtracting: onStartExtracting, onFinishExtracting: onFinishExtracting) { newUrl in + let newRequest = makeNewRequest(changingUrl: newUrl, inRequest: urlRequest) + completion(newRequest) + } + } + public func getCleanURL(from url: URL, onStartExtracting: () -> Void, onFinishExtracting: @escaping () -> Void,