Receiving web push notifications on iOS, using the toot-relay forwarder

DataObjects.swift

import Foundation

struct PushNotification: Codable {
	let accessToken: String
	let preferredLocale: String
	let notificationId: Int64
	let notificationType: Type
	let icon: URL
	let title: String
	let body: String

	enum CodingKeys: String, CodingKey {
		case accessToken = "access_token"
		case preferredLocale = "preferred_locale"
		case notificationId = "notification_id"
		case notificationType = "notification_type"
		case icon = "icon"
		case title = "title"
		case body = "body"
	}

	enum `Type`: String, Codable {
		case favourite = "favourite"
		case follow = "follow"
		case mention = "mention"
		case reblog = "reblog"
	}
}

struct PushNotificationSubscription: Codable {
	let id: Int
	let endpoint: URL
	let alerts: PushNotificationAlerts
}

struct PushNotificationAlerts: Codable {
	let favourite: Bool
	let follow: Bool
	let mention: Bool
	let reblog: Bool

	static var all = PushNotificationAlerts(favourite: true, follow: true, mention: true, reblog: true)

	var isActive: Bool {
		return favourite || follow || mention || reblog
	}
}

struct PushNotificationSubscriptionRequest: Codable {
	let subscription: Subscription?
	let data: Data

	struct Subscription: Codable {
		let endpoint: String
		let keys: Keys

		struct Keys: Codable {
			let p256dh: String
			let auth: String
		}
	}

	struct Data: Codable {
		let alerts: PushNotificationAlerts
	}
}

extension PushNotificationSubscriptionRequest {
	init(endpoint: String, receiver: PushNotificationReceiver, alerts: PushNotificationAlerts) {
		self.init(
			subscription: .init(
				endpoint: endpoint,
				keys: .init(
					p256dh: receiver.publicKeyData.base64UrlEncodedString(),
					auth: receiver.authentication.base64UrlEncodedString()
				)
			), data: .init(alerts: alerts)
		)
	}
}

extension Data {
	func base64UrlEncodedString() -> String {
		return base64EncodedString()
			.replacingOccurrences(of: "+", with: "-")
			.replacingOccurrences(of: "/", with: "_")
			.replacingOccurrences(of: "=", with: "")
	}
}
struct PushNotificationDeviceToken: Codable, Equatable {
	let deviceToken: Data
	let isProduction: Bool

	init(deviceToken: Data) {
		self.deviceToken = deviceToken

		let startData = "<?xml version=\"1.0\" encoding=\"UTF-8\"?>".data(using: .ascii)!
		let endData = "</plist>".data(using: .ascii)!

		if let url = Bundle.main.url(forResource: "embedded", withExtension: "mobileprovision"),
 		let data = try? Data(contentsOf: url),
 		let startIndex = data.range(of: startData)?.lowerBound,
		let endIndex = data.range(of: endData)?.upperBound,
		let plist = try? PropertyListSerialization.propertyList(from: data[startIndex ..< endIndex], options: [], format: nil),
 		let dict = plist as? [String: Any],
 		let entitlements = dict["Entitlements"] as? [String: Any],
		entitlements["aps-environment"] as? String == "development" {
			self.isProduction = false
 		} else {
			self.isProduction = true
    	}
	}

	func endpoint(service: URL, extra: String?) -> URL {
		var endpoint = service
		endpoint.appendPathComponent(isProduction ? "production" : "development")
		endpoint.appendPathComponent(deviceToken.hexString)
		if let extra = extra {
			endpoint.appendPathComponent(extra)
		}
		return endpoint
	}

	static func ==(lhs: PushNotificationDeviceToken, rhs: PushNotificationDeviceToken) -> Bool {
		return lhs.deviceToken == rhs.deviceToken && lhs.isProduction == rhs.isProduction
	}
}

extension Data {
	var hexString: String {
	 	return map { String(format: "%02x", $0) }.joined()
	}

	func range(of substring: Data) -> Range<Int>? {
		for i in 0 ..< count - substring.count {
			var match = true
			for j in 0 ..< substring.count {
				if self[i + j] != substring[j] {
					match = false
					break
				}
			}
			if match {
				return i ..< i + substring.count
			}
		}
		return nil
	}
}

Decode85.swift

import Foundation

private let DecodeTable: [UInt32] = [
	0xff, 0x44, 0xff, 0x54, 0x53, 0x52, 0x48, 0xff,
	0x4b, 0x4c, 0x46, 0x41, 0xff, 0x3f, 0x3e, 0x45,
	0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
	0x08, 0x09, 0x40, 0xff, 0x49, 0x42, 0x4a, 0x47,
	0x51, 0x24, 0x25, 0x26, 0x27, 0x28, 0x29, 0x2a,
	0x2b, 0x2c, 0x2d, 0x2e, 0x2f, 0x30, 0x31, 0x32,
	0x33, 0x34, 0x35, 0x36, 0x37, 0x38, 0x39, 0x3a,
	0x3b, 0x3c, 0x3d, 0x4d, 0xff, 0x4e, 0x43, 0xff,
	0xff, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f, 0x10,
	0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18,
	0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f, 0x20,
	0x21, 0x22, 0x23, 0x4f, 0xff, 0x50, 0xff, 0xff
]

extension String {
	func decode85() -> Data {
		var data = Data()
		var block: UInt32 = 0
		var n = 0
		for c in utf8 {
			if c >= 32, c < 128, DecodeTable[Int(c - 32)] != 0xff {
				let value = DecodeTable[Int(c - 32)]
				block = block * 85 + value
				n += 1
				if n == 5 {
					data.append(UInt8(block >> 24))
					data.append(UInt8((block >> 16) & 0xff))
					data.append(UInt8((block >> 8) & 0xff))
					data.append(UInt8(block & 0xff))
					block = 0
					n = 0
				}
			}
		}

		if n >= 4 { data.append(UInt8((block >> 16) & 0xff)) }
		if n >= 3 { data.append(UInt8((block >> 8) & 0xff)) }
		if n >= 2 { data.append(UInt8(block & 0xff)) }

		return data
	}
}

Example.swift

extension UNNotificationContent {
	func decrypt() throws -> PushNotification {
		// TODO: aes128gcm only uses p and x.
		guard let payload = (userInfo["p"] as? String)?.decode85(),
		let salt = (userInfo["s"] as? String)?.decode85(),
		let serverPublicKeyData = (userInfo["k"] as? String)?.decode85(),
		let identifier = userInfo["x"] as? String else {
			throw DecryptptNotificationErrorType.fieldsNotFound
		}

		// TODO: Find your receiver: PushNotificationReceiver object from somewhere!
		
		let decrypted = try receiver.decrypt(payload: payload, salt: salt, serverPublicKeyData: serverPublicKeyData)

		return try JSONDecoder().decode(PushNotification.self, from: decrypted)
	}
}

enum DecryptptNotificationErrorType: Error {
	case fieldsNotFound
}

PushNotificationReceiver.swift

import Foundation
import Security


struct PushNotificationReceiver: Codable {
	let privateKeyData: Data
	let publicKeyData: Data
	let authentication: Data
}

extension PushNotificationReceiver {
	init() throws {
		var error: Unmanaged<CFError>?

		guard let privateKey = SecKeyCreateRandomKey([
			kSecAttrKeyType as String: kSecAttrKeyTypeECSECPrimeRandom,
			kSecAttrKeySizeInBits as String: 256,
		] as CFDictionary, &error) else {
			throw PushNotificationReceiverErrorType.creatingKeyFailed(error?.takeRetainedValue())
		}

		guard let privateKeyData = SecKeyCopyExternalRepresentation(privateKey, &error) as Data? else {
			throw PushNotificationReceiverErrorType.extractingPrivateKeyFailed(error?.takeRetainedValue())
		}

		guard let publicKey = SecKeyCopyPublicKey(privateKey) else {
			throw PushNotificationReceiverErrorType.impossible
		}

		guard let publicKeyData = SecKeyCopyExternalRepresentation(publicKey, &error) as Data? else {
			throw PushNotificationReceiverErrorType.extractingPublicKeyFailed(error?.takeRetainedValue())
		}

		var authentication = Data(count: 16)
		try authentication.withUnsafeMutableBytes { (bytes: UnsafeMutablePointer<UInt8>) -> Void in
			guard SecRandomCopyBytes(kSecRandomDefault, 16, bytes) == errSecSuccess else {
				throw PushNotificationReceiverErrorType.creatingRandomDataFailed(error?.takeRetainedValue())
			}
		}

		self.init(
			privateKeyData: privateKeyData,
			publicKeyData: publicKeyData,
			authentication: authentication
		)
	}
}

extension PushNotificationReceiver {
	func decrypt(payload: Data, salt: Data, serverPublicKeyData: Data) throws -> Data {
		var error: Unmanaged<CFError>?

		guard let privateKey = SecKeyCreateWithData(privateKeyData as CFData,[
			kSecAttrKeyType as String: kSecAttrKeyTypeECSECPrimeRandom,
			kSecAttrKeyClass as String: kSecAttrKeyClassPrivate,
			kSecAttrKeySizeInBits as String: 256,
		] as CFDictionary, &error) else {
			throw PushNotificationReceiverErrorType.restoringKeyFailed(error?.takeRetainedValue())
		}

		guard let serverPublicKey = SecKeyCreateWithData(serverPublicKeyData as CFData,[
			kSecAttrKeyType as String: kSecAttrKeyTypeECSECPrimeRandom,
			kSecAttrKeyClass as String: kSecAttrKeyClassPublic,
			kSecAttrKeySizeInBits as String: 256,
		] as CFDictionary, &error) else {
			throw PushNotificationReceiverErrorType.creatingKeyFailed(error?.takeRetainedValue())
		}

		guard let sharedSecret = SecKeyCopyKeyExchangeResult(privateKey, .ecdhKeyExchangeStandard, serverPublicKey, [:] as CFDictionary, &error) as Data? else {
			throw PushNotificationReceiverErrorType.keyExhangedFailed(error?.takeRetainedValue())
		}

		// TODO: These steps are slightly different from aes128gcm
		let secondSaltInfo = "Content-Encoding: auth\0".data(using: .utf8)!
		let secondSalt = deriveKey(firstSalt: authentication, secondSalt: sharedSecret, info: secondSaltInfo, length: 32)

		let keyInfo = info(type: "aesgcm", clientPublicKey: publicKeyData, serverPublicKey: serverPublicKeyData)
		let key = deriveKey(firstSalt: salt, secondSalt: secondSalt, info: keyInfo, length: 16)

		let nonceInfo = info(type: "nonce", clientPublicKey: publicKeyData, serverPublicKey: serverPublicKeyData)
		let nonce = deriveKey(firstSalt: salt, secondSalt: secondSalt, info: nonceInfo, length: 12)

		let gcm = try SwiftGCM(key: key, nonce: nonce, tagSize: 16)
		let clearText = try gcm.decrypt(auth: nil, ciphertext: payload)

		guard clearText.count >= 2 else {
			throw PushNotificationReceiverErrorType.clearTextTooShort
		}

		let paddingLength = Int(clearText[0]) * 256 + Int(clearText[1])
		guard clearText.count >= 2 + paddingLength else {
			throw PushNotificationReceiverErrorType.clearTextTooShort
		}

		let unpadded = clearText.suffix(from: paddingLength + 2)

		return unpadded
	}

	private func deriveKey(firstSalt: Data, secondSalt: Data, info: Data, length: Int) -> Data {
		return firstSalt.withUnsafeBytes { (firstSaltBytes: UnsafePointer<UInt8>) -> Data in
			return secondSalt.withUnsafeBytes { (secondSaltBytes: UnsafePointer<UInt8>) -> Data in
				return info.withUnsafeBytes { (infoBytes: UnsafePointer<UInt8>) -> Data in
					// RFC5869 Extract
					var context = CCHmacContext()
					CCHmacInit(&context, CCHmacAlgorithm(kCCHmacAlgSHA256), firstSaltBytes, firstSalt.count)
					CCHmacUpdate(&context, secondSaltBytes, secondSalt.count)

					var hmac: [UInt8] = .init(repeating: 0, count: 32)
					CCHmacFinal(&context, &hmac)

					// RFC5869 Expand
					CCHmacInit(&context, CCHmacAlgorithm(kCCHmacAlgSHA256), &hmac, hmac.count)
					CCHmacUpdate(&context, infoBytes, info.count)

					var one: [UInt8] = [1] // Add sequence byte. We only support short keys so this is always just 1.
					CCHmacUpdate(&context, &one, 1)
					CCHmacFinal(&context, &hmac)

					return Data(bytes: hmac.prefix(upTo: length))
				}
			}
		}
	}

	private func info(type: String, clientPublicKey: Data, serverPublicKey: Data) -> Data {
		var info = Data()

		info.append("Content-Encoding: ".data(using: .utf8)!)
		info.append(type.data(using: .utf8)!)
		info.append(0)
		info.append("P-256".data(using: .utf8)!)
		info.append(0)
		info.append(0)
		info.append(65)
		info.append(clientPublicKey)
		info.append(0)
		info.append(65)
		info.append(serverPublicKey)

		return info
	}
}

enum PushNotificationReceiverErrorType: Error {
	case invalidKey
	case impossible
	case creatingKeyFailed(Error?)
	case restoringKeyFailed(Error?)
	case extractingPrivateKeyFailed(Error?)
	case extractingPublicKeyFailed(Error?)
	case creatingRandomDataFailed(Error?)
	case keyExhangedFailed(Error?)
	case clearTextTooShort
}