NetP x Subscription Clean-up

Sources/NetworkProtection/FeatureActivation/NetworkProtectionCodeRedemptionCoordinator.swift

@@ -25,6 +25,7 @@ public protocol NetworkProtectionCodeRedeeming {
     func redeem(_ code: String) async throws
 
     /// Exchanges an access token for an auth token, and stores the resulting auth token
+    @available(*, deprecated, message: "[NetP Subscription] Use subscription access token instead")
     func exchange(accessToken: String) async throws
 
 }

Sources/NetworkProtection/KeyManagement/NetworkProtectionTokenStore.swift

@@ -20,9 +20,9 @@ import Foundation
 import Common
 
 public protocol NetworkProtectionTokenStore {
-
     /// Store an auth token.
     ///
+    @available(*, deprecated, message: "[NetP Subscription] Do not manually manage auth token")
     func store(_ token: String) throws
 
     /// Obtain the current auth token.
@@ -31,17 +31,8 @@ public protocol NetworkProtectionTokenStore {
 
     /// Delete the stored auth token.
     ///
+    @available(*, deprecated, message: "[NetP Subscription] Do not manually manage auth token")
     func deleteToken() throws
-
-    /// Convert DDG-access-token to NetP-auth-token
-    ///
-    /// todo - https://app.asana.com/0/0/1206541966681608/f
-    static func makeToken(from subscriptionAccessToken: String) -> String
-
-    /// Check if a given token is derived from DDG-access-token
-    ///
-    /// todo - https://app.asana.com/0/0/1206541966681608/f
-    static func isSubscriptionAccessToken(_ token: String) -> Bool
 }
 
 /// Store an auth token for NetworkProtection on behalf of the user. This key is then used to authenticate requests for registration and server fetches from the Network Protection backend servers.
@@ -50,6 +41,9 @@ public final class NetworkProtectionKeychainTokenStore: NetworkProtectionTokenSt
     private let keychainStore: NetworkProtectionKeychainStore
     private let errorEvents: EventMapping<NetworkProtectionError>?
     private let isSubscriptionEnabled: Bool
+    private let accessTokenProvider: () -> String?
+
+    private static var authTokenPrefix: String { "ddg:" }
 
     public struct Defaults {
         static let tokenStoreEntryLabel = "DuckDuckGo Network Protection Auth Token"
@@ -60,12 +54,14 @@ public final class NetworkProtectionKeychainTokenStore: NetworkProtectionTokenSt
     public init(keychainType: KeychainType,
                 serviceName: String = Defaults.tokenStoreService,
                 errorEvents: EventMapping<NetworkProtectionError>?,
-                isSubscriptionEnabled: Bool) {
+                isSubscriptionEnabled: Bool,
+                accessTokenProvider: @escaping () -> String?) {
         keychainStore = NetworkProtectionKeychainStore(label: Defaults.tokenStoreEntryLabel,
                                                        serviceName: serviceName,
                                                        keychainType: keychainType)
         self.errorEvents = errorEvents
         self.isSubscriptionEnabled = isSubscriptionEnabled
+        self.accessTokenProvider = accessTokenProvider
     }
 
     public func store(_ token: String) throws {
@@ -78,7 +74,15 @@ public final class NetworkProtectionKeychainTokenStore: NetworkProtectionTokenSt
         }
     }
 
+    public func makeToken(from subscriptionAccessToken: String) -> String {
+        Self.authTokenPrefix + subscriptionAccessToken
+    }
+
     public func fetchToken() throws -> String? {
+        if isSubscriptionEnabled, let authToken = accessTokenProvider() {
+            return makeToken(from: authToken)
+        }
+
         do {
             return try keychainStore.readData(named: Defaults.tokenStoreName).flatMap {
                 String(data: $0, encoding: .utf8)
@@ -91,12 +95,6 @@ public final class NetworkProtectionKeychainTokenStore: NetworkProtectionTokenSt
 
     public func deleteToken() throws {
         do {
-            // Skip deleting DDG-access-token as it's used for entitlement validity check
-            // todo - https://app.asana.com/0/0/1206541966681608/f
-            if let token = try? fetchToken(), Self.isSubscriptionAccessToken(token) {
-                return
-            }
-
             try keychainStore.deleteAll()
         } catch {
             handle(error)
@@ -118,13 +116,4 @@ public final class NetworkProtectionKeychainTokenStore: NetworkProtectionTokenSt
 }
 
 extension NetworkProtectionTokenStore {
-    private static var authTokenPrefix: String { "ddg:" }
-
-    public static func makeToken(from subscriptionAccessToken: String) -> String {
-        "\(authTokenPrefix)\(subscriptionAccessToken)"
-    }
-
-    public static func isSubscriptionAccessToken(_ token: String) -> Bool {
-        token.hasPrefix(authTokenPrefix)
-    }
 }

Sources/NetworkProtection/PacketTunnelProvider.swift

@@ -799,7 +799,7 @@ open class PacketTunnelProvider: NEPacketTunnelProvider {
             )
         } catch {
             if isSubscriptionEnabled, let error = error as? NetworkProtectionError, case .vpnAccessRevoked = error {
-                await handleInvalidEntitlement(attemptsShutdown: false)
+                await handleInvalidEntitlement(attemptsShutdown: true)
                 throw TunnelError.vpnAccessRevoked
             }
 

Sources/NetworkProtectionTestUtils/KeyManagement/MockNetworkProtectionTokenStore.swift

@@ -46,4 +46,9 @@ public final class MockNetworkProtectionTokenStorage: NetworkProtectionTokenStor
     public func deleteToken() throws {
         didCallDeleteToken = true
     }
+
+    public func fetchSubscriptionToken() throws -> String? {
+        try fetchToken()
+    }
+
 }

Sources/Subscription/AccountManager.swift

@@ -219,8 +219,8 @@ public class AccountManager: AccountManaging {
         case noCachedData
     }
 
-    public func hasEntitlement(for entitlement: Entitlement.ProductName) async -> Result<Bool, Error> {
-        switch await fetchEntitlements() {
+    public func hasEntitlement(for entitlement: Entitlement.ProductName, cachePolicy: CachePolicy = .returnCacheDataElseLoad) async -> Result<Bool, Error> {
+        switch await fetchEntitlements(policy: cachePolicy) {
         case .success(let entitlements):
             return .success(entitlements.compactMap { $0.product }.contains(entitlement))
         case .failure(let error):

Tests/NetworkProtectionTests/Mocks/NetworkProtectionTokenStoreMocks.swift

@@ -32,9 +32,10 @@ final class NetworkProtectionTokenStoreMock: NetworkProtectionTokenStore {
     }
 
     func deleteToken() {
-        if let token = fetchToken(), Self.isSubscriptionAccessToken(token) {
-            return
-        }
         self.token = nil
     }
+
+    func fetchSubscriptionToken() throws -> String? {
+        "ddg:accessToken"
+    }
 }

Tests/NetworkProtectionTests/NetworkProtectionDeviceManagerTests.swift

@@ -222,22 +222,6 @@ final class NetworkProtectionDeviceManagerTests: XCTestCase {
         XCTAssertEqual(try? serverListStore.storedNetworkProtectionServerList(), [registeredServer])
         XCTAssertNotNil(networkClient.spyRegister)
     }
-
-    func testStoringAccessToken() {
-        tokenStore.store(NetworkProtectionTokenStoreMock.makeToken(from: "access-token"))
-        XCTAssertEqual(tokenStore.fetchToken(), "ddg:access-token")
-
-        tokenStore.deleteToken()
-        XCTAssertEqual(tokenStore.fetchToken(), "ddg:access-token")
-    }
-
-    func testStoringAuthToken() {
-        tokenStore.store("auth-token")
-        XCTAssertEqual(tokenStore.fetchToken(), "auth-token")
-
-        tokenStore.deleteToken()
-        XCTAssertNil(tokenStore.fetchToken())
-    }
 }
 
 extension NetworkProtectionDeviceManager {