File size: 12,772 Bytes
fc93158
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
import CryptoKit
import DeviceCheck
import Foundation
import StoreKit

enum PushRelayError: LocalizedError {
    case relayBaseURLMissing
    case relayMisconfigured(String)
    case invalidResponse(String)
    case requestFailed(status: Int, message: String)
    case unsupportedAppAttest
    case missingReceipt

    var errorDescription: String? {
        switch self {
        case .relayBaseURLMissing:
            "Push relay base URL missing"
        case let .relayMisconfigured(message):
            message
        case let .invalidResponse(message):
            message
        case let .requestFailed(status, message):
            "Push relay request failed (\(status)): \(message)"
        case .unsupportedAppAttest:
            "App Attest unavailable on this device"
        case .missingReceipt:
            "App Store receipt missing after refresh"
        }
    }
}

private struct PushRelayChallengeResponse: Decodable {
    var challengeId: String
    var challenge: String
    var expiresAtMs: Int64
}

private struct PushRelayRegisterSignedPayload: Encodable {
    var challengeId: String
    var installationId: String
    var bundleId: String
    var environment: String
    var distribution: String
    var gateway: PushRelayGatewayIdentity
    var appVersion: String
    var apnsToken: String
}

private struct PushRelayAppAttestPayload: Encodable {
    var keyId: String
    var attestationObject: String?
    var assertion: String
    var clientDataHash: String
    var signedPayloadBase64: String
}

private struct PushRelayReceiptPayload: Encodable {
    var base64: String
}

private struct PushRelayRegisterRequest: Encodable {
    var challengeId: String
    var installationId: String
    var bundleId: String
    var environment: String
    var distribution: String
    var gateway: PushRelayGatewayIdentity
    var appVersion: String
    var apnsToken: String
    var appAttest: PushRelayAppAttestPayload
    var receipt: PushRelayReceiptPayload
}

struct PushRelayRegisterResponse: Decodable {
    var relayHandle: String
    var sendGrant: String
    var expiresAtMs: Int64?
    var tokenSuffix: String?
    var status: String
}

private struct RelayErrorResponse: Decodable {
    var error: String?
    var message: String?
    var reason: String?
}

private final class PushRelayReceiptRefreshCoordinator: NSObject, SKRequestDelegate {
    private var continuation: CheckedContinuation<Void, Error>?
    private var activeRequest: SKReceiptRefreshRequest?

    func refresh() async throws {
        try await withCheckedThrowingContinuation { continuation in
            self.continuation = continuation
            let request = SKReceiptRefreshRequest()
            self.activeRequest = request
            request.delegate = self
            request.start()
        }
    }

    func requestDidFinish(_ request: SKRequest) {
        self.continuation?.resume(returning: ())
        self.continuation = nil
        self.activeRequest = nil
    }

    func request(_ request: SKRequest, didFailWithError error: Error) {
        self.continuation?.resume(throwing: error)
        self.continuation = nil
        self.activeRequest = nil
    }
}

private struct PushRelayAppAttestProof {
    var keyId: String
    var attestationObject: String?
    var assertion: String
    var clientDataHash: String
    var signedPayloadBase64: String
}

private final class PushRelayAppAttestService {
    func createProof(challenge: String, signedPayload: Data) async throws -> PushRelayAppAttestProof {
        let service = DCAppAttestService.shared
        guard service.isSupported else {
            throw PushRelayError.unsupportedAppAttest
        }

        let keyID = try await self.loadOrCreateKeyID(using: service)
        let attestationObject = try await self.attestKeyIfNeeded(
            service: service,
            keyID: keyID,
            challenge: challenge)
        let signedPayloadHash = Data(SHA256.hash(data: signedPayload))
        let assertion = try await self.generateAssertion(
            service: service,
            keyID: keyID,
            signedPayloadHash: signedPayloadHash)

        return PushRelayAppAttestProof(
            keyId: keyID,
            attestationObject: attestationObject,
            assertion: assertion.base64EncodedString(),
            clientDataHash: Self.base64URL(signedPayloadHash),
            signedPayloadBase64: signedPayload.base64EncodedString())
    }

    private func loadOrCreateKeyID(using service: DCAppAttestService) async throws -> String {
        if let existing = PushRelayRegistrationStore.loadAppAttestKeyID(), !existing.isEmpty {
            return existing
        }
        let keyID = try await service.generateKey()
        _ = PushRelayRegistrationStore.saveAppAttestKeyID(keyID)
        return keyID
    }

    private func attestKeyIfNeeded(
        service: DCAppAttestService,
        keyID: String,
        challenge: String)
    async throws -> String? {
        if PushRelayRegistrationStore.loadAttestedKeyID() == keyID {
            return nil
        }
        let challengeData = Data(challenge.utf8)
        let clientDataHash = Data(SHA256.hash(data: challengeData))
        let attestation = try await service.attestKey(keyID, clientDataHash: clientDataHash)
        // Apple treats App Attest key attestation as a one-time operation. Save the
        // attested marker immediately so later receipt/network failures do not cause a
        // permanently broken re-attestation loop on the same key.
        _ = PushRelayRegistrationStore.saveAttestedKeyID(keyID)
        return attestation.base64EncodedString()
    }

    private func generateAssertion(
        service: DCAppAttestService,
        keyID: String,
        signedPayloadHash: Data)
    async throws -> Data {
        do {
            return try await service.generateAssertion(keyID, clientDataHash: signedPayloadHash)
        } catch {
            _ = PushRelayRegistrationStore.clearAppAttestKeyID()
            _ = PushRelayRegistrationStore.clearAttestedKeyID()
            throw error
        }
    }

    private static func base64URL(_ data: Data) -> String {
        data.base64EncodedString()
            .replacingOccurrences(of: "+", with: "-")
            .replacingOccurrences(of: "/", with: "_")
            .replacingOccurrences(of: "=", with: "")
    }
}

private final class PushRelayReceiptProvider {
    func loadReceiptBase64() async throws -> String {
        if let receipt = self.readReceiptData() {
            return receipt.base64EncodedString()
        }
        let refreshCoordinator = PushRelayReceiptRefreshCoordinator()
        try await refreshCoordinator.refresh()
        if let refreshed = self.readReceiptData() {
            return refreshed.base64EncodedString()
        }
        throw PushRelayError.missingReceipt
    }

    private func readReceiptData() -> Data? {
        guard let url = Bundle.main.appStoreReceiptURL else { return nil }
        guard let data = try? Data(contentsOf: url), !data.isEmpty else { return nil }
        return data
    }
}

// The client is constructed once and used behind PushRegistrationManager actor isolation.
final class PushRelayClient: @unchecked Sendable {
    private let baseURL: URL
    private let session: URLSession
    private let jsonDecoder = JSONDecoder()
    private let jsonEncoder = JSONEncoder()
    private let appAttest = PushRelayAppAttestService()
    private let receiptProvider = PushRelayReceiptProvider()

    init(baseURL: URL, session: URLSession = .shared) {
        self.baseURL = baseURL
        self.session = session
    }

    var normalizedBaseURLString: String {
        Self.normalizeBaseURLString(self.baseURL)
    }

    func register(
        installationId: String,
        bundleId: String,
        appVersion: String,
        environment: PushAPNsEnvironment,
        distribution: PushDistributionMode,
        apnsTokenHex: String,
        gatewayIdentity: PushRelayGatewayIdentity)
    async throws -> PushRelayRegisterResponse {
        let challenge = try await self.fetchChallenge()
        let signedPayload = PushRelayRegisterSignedPayload(
            challengeId: challenge.challengeId,
            installationId: installationId,
            bundleId: bundleId,
            environment: environment.rawValue,
            distribution: distribution.rawValue,
            gateway: gatewayIdentity,
            appVersion: appVersion,
            apnsToken: apnsTokenHex)
        let signedPayloadData = try self.jsonEncoder.encode(signedPayload)
        let appAttest = try await self.appAttest.createProof(
            challenge: challenge.challenge,
            signedPayload: signedPayloadData)
        let receiptBase64 = try await self.receiptProvider.loadReceiptBase64()
        let requestBody = PushRelayRegisterRequest(
            challengeId: signedPayload.challengeId,
            installationId: signedPayload.installationId,
            bundleId: signedPayload.bundleId,
            environment: signedPayload.environment,
            distribution: signedPayload.distribution,
            gateway: signedPayload.gateway,
            appVersion: signedPayload.appVersion,
            apnsToken: signedPayload.apnsToken,
            appAttest: PushRelayAppAttestPayload(
                keyId: appAttest.keyId,
                attestationObject: appAttest.attestationObject,
                assertion: appAttest.assertion,
                clientDataHash: appAttest.clientDataHash,
                signedPayloadBase64: appAttest.signedPayloadBase64),
            receipt: PushRelayReceiptPayload(base64: receiptBase64))

        let endpoint = self.baseURL.appending(path: "v1/push/register")
        var request = URLRequest(url: endpoint)
        request.httpMethod = "POST"
        request.timeoutInterval = 20
        request.setValue("application/json", forHTTPHeaderField: "Content-Type")
        request.httpBody = try self.jsonEncoder.encode(requestBody)

        let (data, response) = try await self.session.data(for: request)
        let status = Self.statusCode(from: response)
        guard (200..<300).contains(status) else {
            if status == 401 {
                // If the relay rejects registration, drop local App Attest state so the next
                // attempt re-attests instead of getting stuck without an attestation object.
                _ = PushRelayRegistrationStore.clearAppAttestKeyID()
                _ = PushRelayRegistrationStore.clearAttestedKeyID()
            }
            throw PushRelayError.requestFailed(
                status: status,
                message: Self.decodeErrorMessage(data: data))
        }
        let decoded = try self.decode(PushRelayRegisterResponse.self, from: data)
        return decoded
    }

    private func fetchChallenge() async throws -> PushRelayChallengeResponse {
        let endpoint = self.baseURL.appending(path: "v1/push/challenge")
        var request = URLRequest(url: endpoint)
        request.httpMethod = "POST"
        request.timeoutInterval = 10
        request.setValue("application/json", forHTTPHeaderField: "Content-Type")
        request.httpBody = Data("{}".utf8)

        let (data, response) = try await self.session.data(for: request)
        let status = Self.statusCode(from: response)
        guard (200..<300).contains(status) else {
            throw PushRelayError.requestFailed(
                status: status,
                message: Self.decodeErrorMessage(data: data))
        }
        return try self.decode(PushRelayChallengeResponse.self, from: data)
    }

    private func decode<T: Decodable>(_ type: T.Type, from data: Data) throws -> T {
        do {
            return try self.jsonDecoder.decode(type, from: data)
        } catch {
            throw PushRelayError.invalidResponse(error.localizedDescription)
        }
    }

    private static func statusCode(from response: URLResponse) -> Int {
        (response as? HTTPURLResponse)?.statusCode ?? 0
    }

    private static func normalizeBaseURLString(_ url: URL) -> String {
        var absolute = url.absoluteString
        while absolute.hasSuffix("/") {
            absolute.removeLast()
        }
        return absolute
    }

    private static func decodeErrorMessage(data: Data) -> String {
        if let decoded = try? JSONDecoder().decode(RelayErrorResponse.self, from: data) {
            let message = decoded.message ?? decoded.reason ?? decoded.error ?? ""
            if !message.isEmpty {
                return message
            }
        }
        let raw = String(data: data, encoding: .utf8)?.trimmingCharacters(in: .whitespacesAndNewlines) ?? ""
        return raw.isEmpty ? "unknown relay error" : raw
    }
}