Skip to content

Commit

Permalink
Merge remote-tracking branch 'origin/main' into alex/malware-protecti…
Browse files Browse the repository at this point in the history 
…on-1
  • Loading branch information
@mallexxx
mallexxx committed Nov 28, 2024
2 parents 49f022e + 5a24a88 commit e2ef57c
Show file tree
Hide file tree
Showing 5 changed files with 157 additions and 36 deletions.
4 changes: 4 additions & 0 deletions Sources/BrowserServicesKit/ContentBlocking/ContentBlockerDebugEvents.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -48,4 +48,8 @@ public enum ContentBlockerDebugEvents {

case contentBlockingCompilationTime

case contentBlockingLookupRulesSucceeded
case contentBlockingFetchLRCSucceeded
case contentBlockingLRCMissing
case contentBlockingNoMatchInLRC
}
19 changes: 15 additions & 4 deletions Sources/BrowserServicesKit/ContentBlocking/ContentBlockerRulesManager.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -152,10 +152,17 @@ public class ContentBlockerRulesManager: CompiledRuleListsSource {

if !self.lookupCompiledRules() {
if let lastCompiledRules = lastCompiledRulesStore?.rules, !lastCompiledRules.isEmpty {
self.fetchLastCompiledRules(with: lastCompiledRules)
if self.fetchLastCompiledRules(with: lastCompiledRules) {
self.errorReporting?.fire(.contentBlockingFetchLRCSucceeded)
} else {
self.errorReporting?.fire(.contentBlockingNoMatchInLRC)
}
} else {
self.errorReporting?.fire(.contentBlockingLRCMissing)
self.startCompilationProcess()
}
} else {
self.errorReporting?.fire(.contentBlockingLookupRulesSucceeded)
}
}
}
Expand Down Expand Up @@ -261,7 +268,7 @@ public class ContentBlockerRulesManager: CompiledRuleListsSource {
Go through source managers and check if there are already compiled rules in the WebKit rule cache.
Returns true if rules were found, false otherwise.
*/
private func fetchLastCompiledRules(with lastCompiledRules: [LastCompiledRules]) {
private func fetchLastCompiledRules(with lastCompiledRules: [LastCompiledRules]) -> Bool {
Logger.contentBlocking.debug("Fetch last compiled rules: \(lastCompiledRules.count, privacy: .public)")

let initialCompilationTask = LastCompiledRulesLookupTask(sourceRules: rulesSource.contentBlockerRulesLists,
Expand All @@ -274,8 +281,10 @@ public class ContentBlockerRulesManager: CompiledRuleListsSource {
// We want to confine Compilation work to WorkQueue, so we wait to come back from async Task
mutex.wait()

if let rules = initialCompilationTask.getFetchedRules() {
applyRules(rules)
let rulesFound = initialCompilationTask.getFetchedRules()

if let rulesFound {
applyRules(rulesFound)
} else {
lock.lock()
state = .idle
Expand All @@ -284,6 +293,8 @@ public class ContentBlockerRulesManager: CompiledRuleListsSource {

// No matter if rules were found or not, we need to schedule recompilation, after all
scheduleCompilation()

return rulesFound != nil
}

private func prepareSourceManagers() {
Expand Down
34 changes: 33 additions & 1 deletion Sources/Subscription/SubscriptionCookie/HTTPCookieStore.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,4 +25,36 @@ public protocol HTTPCookieStore {
func deleteCookie(_ cookie: HTTPCookie) async
}

extension WKHTTPCookieStore: HTTPCookieStore {}
@MainActor
public struct WKHTTPCookieStoreWrapper: HTTPCookieStore {

let store: WKHTTPCookieStore

public init(store: WKHTTPCookieStore) {
self.store = store
}

public func allCookies() async -> [HTTPCookie] {
return await withCheckedContinuation { continuation in
store.getAllCookies { cookies in
continuation.resume(returning: cookies)
}
}
}

public func setCookie(_ cookie: HTTPCookie) async {
await withCheckedContinuation { continuation in
store.setCookie(cookie) {
continuation.resume()
}
}
}

public func deleteCookie(_ cookie: HTTPCookie) async {
await withCheckedContinuation { continuation in
store.delete(cookie) {
continuation.resume()
}
}
}
}
107 changes: 81 additions & 26 deletions ...erServicesKitTests/ContentBlocker/ContentBlockerRulesManagerInitialCompilationTests.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ import TrackerRadarKit
import BrowserServicesKit
import WebKit
import XCTest
import Common

final class CountedFulfillmentTestExpectation: XCTestExpectation {
private(set) var currentFulfillmentCount: Int = 0
Expand Down Expand Up @@ -57,12 +58,24 @@ final class ContentBlockerRulesManagerInitialCompilationTests: XCTestCase {
expStore.fulfill()
}

let lookupAndFetchExp = self.expectation(description: "LRC should be missing")
let errorHandler = EventMapping<ContentBlockerDebugEvents> { event, _, params, _ in
if case .contentBlockingLRCMissing = event {
lookupAndFetchExp.fulfill()
} else if case .contentBlockingCompilationTime = event {
XCTAssertNotNil(params?["compilationTime"])
} else {
XCTFail("Unexpected event: \(event)")
}
}

let cbrm = ContentBlockerRulesManager(rulesSource: mockRulesSource,
exceptionsSource: mockExceptionsSource,
lastCompiledRulesStore: mockLastCompiledRulesStore,
updateListener: rulesUpdateListener)
updateListener: rulesUpdateListener,
errorReporting: errorHandler)

wait(for: [exp, expStore], timeout: 15.0)
wait(for: [exp, expStore, lookupAndFetchExp], timeout: 15.0)

XCTAssertNotNil(mockLastCompiledRulesStore.rules)
XCTAssertEqual(mockLastCompiledRulesStore.rules.first?.etag, mockRulesSource.trackerData?.etag)
Expand Down Expand Up @@ -93,6 +106,8 @@ final class ContentBlockerRulesManagerInitialCompilationTests: XCTestCase {
XCTFail("Should use rules cached by WebKit")
}

let lookupAndFetchExp = self.expectation(description: "Should not fetch LRC")

// simulate the rules have been compiled in the past so the WKContentRuleListStore contains it
_ = ContentBlockerRulesManager(rulesSource: mockRulesSource,
exceptionsSource: mockExceptionsSource,
Expand All @@ -103,15 +118,28 @@ final class ContentBlockerRulesManagerInitialCompilationTests: XCTestCase {
rulesUpdateListener.onRulesUpdated = { rules in
exp.fulfill()
if exp.currentFulfillmentCount == 1 { // finished compilation after first installation
let errorHandler = EventMapping<ContentBlockerDebugEvents> { event, _, params, _ in
if case .contentBlockingFetchLRCSucceeded = event {
XCTFail("Should not fetch LRC")
} else if case .contentBlockingLookupRulesSucceeded = event {
lookupAndFetchExp.fulfill()
} else if case .contentBlockingCompilationTime = event {
XCTAssertNotNil(params?["compilationTime"])
} else {
XCTFail("Unexpected event: \(event)")
}
}

_ = ContentBlockerRulesManager(rulesSource: mockRulesSource,
exceptionsSource: mockExceptionsSource,
lastCompiledRulesStore: mockLastCompiledRulesStore,
updateListener: self.rulesUpdateListener)
updateListener: self.rulesUpdateListener,
errorReporting: errorHandler)
}
assertRules(rules)
}

wait(for: [exp], timeout: 15.0)
wait(for: [exp, lookupAndFetchExp], timeout: 15.0)

func assertRules(_ rules: [ContentBlockerRulesManager.Rules]) {
guard let rules = rules.first else { XCTFail("Couldn't get rules"); return }
Expand Down Expand Up @@ -178,12 +206,27 @@ final class ContentBlockerRulesManagerInitialCompilationTests: XCTestCase {
XCTAssertEqual(newListName, rules.name)
}

let lookupAndFetchExp = self.expectation(description: "Should not fetch LRC")

let errorHandler = EventMapping<ContentBlockerDebugEvents> { event, _, params, _ in
if case .contentBlockingFetchLRCSucceeded = event {
XCTFail("Should not fetch LRC")
} else if case .contentBlockingNoMatchInLRC = event {
lookupAndFetchExp.fulfill()
} else if case .contentBlockingCompilationTime = event {
XCTAssertNotNil(params?["compilationTime"])
} else {
XCTFail("Unexpected event: \(event)")
}
}

_ = ContentBlockerRulesManager(rulesSource: mockRulesSource,
exceptionsSource: mockExceptionsSource,
lastCompiledRulesStore: mockLastCompiledRulesStore,
updateListener: rulesUpdateListener)
updateListener: rulesUpdateListener,
errorReporting: errorHandler)

wait(for: [expCacheLookup, expNext], timeout: 15.0)
wait(for: [expCacheLookup, expNext, lookupAndFetchExp], timeout: 15.0)
}

func testInitialCompilation_WhenThereAreChangesToTDS_ShouldBuildRulesUsingLastCompiledRulesAndScheduleRecompilationWithNewSource() {
Expand Down Expand Up @@ -220,14 +263,26 @@ final class ContentBlockerRulesManagerInitialCompilationTests: XCTestCase {
exceptionsSource: mockExceptionsSource,
updateListener: rulesUpdateListener)

let lookupAndFetchExp = self.expectation(description: "Fetch LRC succeeded")
let expOld = CountedFulfillmentTestExpectation(description: "Old Rules Compiled")
rulesUpdateListener.onRulesUpdated = { _ in
expOld.fulfill()

let errorHandler = EventMapping<ContentBlockerDebugEvents> { event, _, params, _ in
if case .contentBlockingFetchLRCSucceeded = event {
lookupAndFetchExp.fulfill()
} else if case .contentBlockingCompilationTime = event {
XCTAssertNotNil(params?["compilationTime"])
} else {
XCTFail("Unexpected event: \(event)")
}
}

_ = ContentBlockerRulesManager(rulesSource: mockUpdatedRulesSource,
exceptionsSource: mockExceptionsSource,
lastCompiledRulesStore: mockLastCompiledRulesStore,
updateListener: self.rulesUpdateListener)
updateListener: self.rulesUpdateListener,
errorReporting: errorHandler)
}

wait(for: [expOld], timeout: 15.0)
Expand All @@ -237,27 +292,27 @@ final class ContentBlockerRulesManagerInitialCompilationTests: XCTestCase {
expLastCompiledFetched.fulfill()
}

let expRecompiled = CountedFulfillmentTestExpectation(description: "New Rules Compiled")
rulesUpdateListener.onRulesUpdated = { _ in
expRecompiled.fulfill()
}
let expRecompiled = CountedFulfillmentTestExpectation(description: "New Rules Compiled")
rulesUpdateListener.onRulesUpdated = { _ in
expRecompiled.fulfill()

if expRecompiled.currentFulfillmentCount == 1 { // finished compilation after cold start (using last compiled rules)
mockLastCompiledRulesStore.onRulesGet = {}
XCTAssertEqual(mockLastCompiledRulesStore.rules.first?.etag, oldEtag)
XCTAssertEqual(mockLastCompiledRulesStore.rules.first?.name, mockRulesSource.ruleListName)
XCTAssertEqual(mockLastCompiledRulesStore.rules.first?.trackerData, mockRulesSource.trackerData?.tds)
XCTAssertEqual(mockLastCompiledRulesStore.rules.first?.identifier, oldIdentifier)
} else if expRecompiled.currentFulfillmentCount == 2 { // finished recompilation of rules due to changed tds
XCTAssertEqual(mockLastCompiledRulesStore.rules.first?.etag, updatedEtag)
XCTAssertEqual(mockLastCompiledRulesStore.rules.first?.name, mockRulesSource.ruleListName)
XCTAssertEqual(mockLastCompiledRulesStore.rules.first?.trackerData, mockRulesSource.trackerData?.tds)
XCTAssertEqual(mockLastCompiledRulesStore.rules.first?.identifier, newIdentifier)
}
}

if expRecompiled.currentFulfillmentCount == 1 { // finished compilation after cold start (using last compiled rules)

mockLastCompiledRulesStore.onRulesGet = {}
XCTAssertEqual(mockLastCompiledRulesStore.rules.first?.etag, oldEtag)
XCTAssertEqual(mockLastCompiledRulesStore.rules.first?.name, mockRulesSource.ruleListName)
XCTAssertEqual(mockLastCompiledRulesStore.rules.first?.trackerData, mockRulesSource.trackerData?.tds)
XCTAssertEqual(mockLastCompiledRulesStore.rules.first?.identifier, oldIdentifier)
} else if expRecompiled.currentFulfillmentCount == 2 { // finished recompilation of rules due to changed tds
XCTAssertEqual(mockLastCompiledRulesStore.rules.first?.etag, updatedEtag)
XCTAssertEqual(mockLastCompiledRulesStore.rules.first?.name, mockRulesSource.ruleListName)
XCTAssertEqual(mockLastCompiledRulesStore.rules.first?.trackerData, mockRulesSource.trackerData?.tds)
XCTAssertEqual(mockLastCompiledRulesStore.rules.first?.identifier, newIdentifier)
}
wait(for: [expLastCompiledFetched, expRecompiled, lookupAndFetchExp], timeout: 15.0)

wait(for: [expLastCompiledFetched, expRecompiled], timeout: 15.0)
}
}

struct MockLastCompiledRules: LastCompiledRules {

Expand Down
29 changes: 24 additions & 5 deletions Tests/BrowserServicesKitTests/ContentBlocker/ContentBlockerRulesManagerTests.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -203,6 +203,8 @@ class ContentBlockerRulesManagerLoadingTests: ContentBlockerRulesManagerTests {

let errorExp = expectation(description: "No error reported")
errorExp.isInverted = true

let lookupAndFetchExp = expectation(description: "Look and Fetch rules failed")
let compilationTimeExp = expectation(description: "Compilation Time reported")
let errorHandler = EventMapping<ContentBlockerDebugEvents> { event, _, params, _ in
if case .contentBlockingCompilationFailed(let listName, let component) = event {
Expand All @@ -217,6 +219,8 @@ class ContentBlockerRulesManagerLoadingTests: ContentBlockerRulesManagerTests {
} else if case .contentBlockingCompilationTime = event {
XCTAssertNotNil(params?["compilationTime"])
compilationTimeExp.fulfill()
} else if case .contentBlockingLRCMissing = event {
lookupAndFetchExp.fulfill()
} else {
XCTFail("Unexpected event: \(event)")
}
Expand All @@ -227,7 +231,7 @@ class ContentBlockerRulesManagerLoadingTests: ContentBlockerRulesManagerTests {
updateListener: rulesUpdateListener,
errorReporting: errorHandler)

wait(for: [exp, errorExp, compilationTimeExp], timeout: 15.0)
wait(for: [exp, errorExp, compilationTimeExp, lookupAndFetchExp], timeout: 15.0)

XCTAssertNotNil(cbrm.currentRules)
XCTAssertEqual(cbrm.currentRules.first?.etag, mockRulesSource.trackerData?.etag)
Expand All @@ -254,6 +258,8 @@ class ContentBlockerRulesManagerLoadingTests: ContentBlockerRulesManagerTests {
}

let errorExp = expectation(description: "Error reported")
let lookupAndFetchExp = expectation(description: "Look and Fetch rules failed")

let errorHandler = EventMapping<ContentBlockerDebugEvents> { event, _, params, _ in
if case .contentBlockingCompilationFailed(let listName, let component) = event {
XCTAssertEqual(listName, DefaultContentBlockerRulesListsSource.Constants.trackerDataSetRulesListName)
Expand All @@ -266,6 +272,8 @@ class ContentBlockerRulesManagerLoadingTests: ContentBlockerRulesManagerTests {

} else if case .contentBlockingCompilationTime = event {
XCTAssertNotNil(params?["compilationTime"])
} else if case .contentBlockingLRCMissing = event {
lookupAndFetchExp.fulfill()
} else {
XCTFail("Unexpected event: \(event)")
}
Expand All @@ -276,7 +284,7 @@ class ContentBlockerRulesManagerLoadingTests: ContentBlockerRulesManagerTests {
updateListener: rulesUpdateListener,
errorReporting: errorHandler)

wait(for: [exp, errorExp], timeout: 15.0)
wait(for: [exp, errorExp, lookupAndFetchExp], timeout: 15.0)

XCTAssertNotNil(cbrm.currentRules)
XCTAssertEqual(cbrm.currentRules.first?.etag, mockRulesSource.embeddedTrackerData.etag)
Expand Down Expand Up @@ -539,6 +547,9 @@ class ContentBlockerRulesManagerLoadingTests: ContentBlockerRulesManagerTests {

let errorExp = expectation(description: "Error reported")
errorExp.expectedFulfillmentCount = 5

let lookupAndFetchExp = expectation(description: "Look and Fetch rules failed")

var errorEvents = [ContentBlockerDebugEvents.Component]()
let errorHandler = EventMapping<ContentBlockerDebugEvents> { event, _, params, _ in
if case .contentBlockingCompilationFailed(let listName, let component) = event {
Expand All @@ -554,6 +565,8 @@ class ContentBlockerRulesManagerLoadingTests: ContentBlockerRulesManagerTests {
} else if case .contentBlockingCompilationTime = event {
XCTAssertNotNil(params?["compilationTime"])
errorExp.fulfill()
} else if case .contentBlockingLRCMissing = event {
lookupAndFetchExp.fulfill()
} else {
XCTFail("Unexpected event: \(event)")
}
Expand All @@ -564,7 +577,7 @@ class ContentBlockerRulesManagerLoadingTests: ContentBlockerRulesManagerTests {
updateListener: rulesUpdateListener,
errorReporting: errorHandler)

wait(for: [exp, errorExp], timeout: 15.0)
wait(for: [exp, errorExp, lookupAndFetchExp], timeout: 15.0)

XCTAssertEqual(Set(errorEvents), Set([.tds,
.tempUnprotected,
Expand Down Expand Up @@ -619,6 +632,9 @@ class ContentBlockerRulesManagerLoadingTests: ContentBlockerRulesManagerTests {

let errorExp = expectation(description: "Error reported")
errorExp.expectedFulfillmentCount = 4

let lookupAndFetchExp = expectation(description: "Look and Fetch rules failed")

var errorEvents = [ContentBlockerDebugEvents.Component]()
let errorHandler = EventMapping<ContentBlockerDebugEvents> { event, _, params, _ in
if case .contentBlockingCompilationFailed(let listName, let component) = event {
Expand All @@ -634,7 +650,10 @@ class ContentBlockerRulesManagerLoadingTests: ContentBlockerRulesManagerTests {
} else if case .contentBlockingCompilationTime = event {
XCTAssertNotNil(params?["compilationTime"])
errorExp.fulfill()
} else {
} else if case .contentBlockingLRCMissing = event {
lookupAndFetchExp.fulfill()
} else
{
XCTFail("Unexpected event: \(event)")
}
}
Expand All @@ -644,7 +663,7 @@ class ContentBlockerRulesManagerLoadingTests: ContentBlockerRulesManagerTests {
updateListener: rulesUpdateListener,
errorReporting: errorHandler)

wait(for: [exp, errorExp], timeout: 15.0)
wait(for: [exp, errorExp, lookupAndFetchExp], timeout: 15.0)

XCTAssertEqual(Set(errorEvents), Set([.tempUnprotected,
.allowlist,
Expand Down

0 comments on commit e2ef57c

Please sign in to comment.