diff --git a/Sources/Brave/Frontend/Shields/ShieldsViewController.swift b/Sources/Brave/Frontend/Shields/ShieldsViewController.swift index 66188cd6d37..50ced0ae6b6 100644 --- a/Sources/Brave/Frontend/Shields/ShieldsViewController.swift +++ b/Sources/Brave/Frontend/Shields/ShieldsViewController.swift @@ -18,14 +18,9 @@ class ShieldsViewController: UIViewController, PopoverContentComponent { let tab: Tab private lazy var url: URL? = { guard let _url = tab.url else { return nil } - - if InternalURL.isValid(url: _url), - let internalURL = InternalURL(_url) { - if internalURL.isErrorPage { - return internalURL.originalURLFromErrorPage - } else if internalURL.isWeb3URL { - return internalURL.extractedUrlParam?.displayURL - } + + if let tabURL = _url.stippedInternalURL { + return tabURL } return _url diff --git a/Sources/BraveShared/Extensions/URLExtensions.swift b/Sources/BraveShared/Extensions/URLExtensions.swift index 879d64dbf48..a8d81717b07 100644 --- a/Sources/BraveShared/Extensions/URLExtensions.swift +++ b/Sources/BraveShared/Extensions/URLExtensions.swift @@ -5,6 +5,7 @@ import Foundation import BraveCore +import Shared extension URL { /// Obtains a URLOrigin given the current URL @@ -22,4 +23,60 @@ extension URL { public var origin: URLOrigin { .init(url: self) } + + /// Obtains a clean stripped url from the current Internal URL + /// + /// Returns the original url without internal parameters + public var stippedInternalURL: URL? { + if InternalURL.isValid(url: self), + let internalURL = InternalURL(self) { + + switch internalURL.urlType { + case .errorPage: + return internalURL.originalURLFromErrorPage + case .web3Page, .sessionRestorePage, .readerModePage, .aboutHomePage: + return internalURL.extractedUrlParam + default: + return nil + } + } + + return nil + } +} + +extension InternalURL { + + enum URLType { + case sessionRestorePage + case errorPage + case readerModePage + case aboutHomePage + case web3Page + case other + } + + var urlType: URLType { + if isErrorPage { + return .errorPage + } + + if isWeb3URL { + return .web3Page + } + + if isReaderModePage { + return .readerModePage + } + + if isSessionRestore { + return .sessionRestorePage + } + + if isAboutHomeURL { + return .aboutHomePage + } + + return .other + } } diff --git a/Tests/BraveSharedTests/URLExtensionTests.swift b/Tests/BraveSharedTests/URLExtensionTests.swift index ac84b7ee803..50c6ca34100 100644 --- a/Tests/BraveSharedTests/URLExtensionTests.swift +++ b/Tests/BraveSharedTests/URLExtensionTests.swift @@ -25,5 +25,15 @@ class URLExtensionTests: XCTestCase { urls.forEach { XCTAssertEqual(URL(string: $0.0)!.origin.serialized, $0.1) } badurls.forEach { XCTAssertTrue(URL(string: $0)!.origin.isOpaque) } } + + func testStrippedInternalURL() { + let urls = [ + ("internal://local/web3/ddns?service_id=ethereum&url=http%3A%2F%2Fvitalik%2Eeth%2F", URL(string: "http://vitalik.eth/")), + ("internal://local/sessionrestore?url=https://en.m.wikipedia.org/wiki/Main_Page", URL(string: "https://en.m.wikipedia.org/wiki/Main_Page")), + ("internal://local/reader-mode?url=https://en.m.wikipedia.org/wiki/Main_Page", URL(string: "https://en.m.wikipedia.org/wiki/Main_Page")) + ] + + urls.forEach { XCTAssertEqual(URL(string: $0.0)!.stippedInternalURL, $0.1) } + } }