From 9ff4c47e376d578524bfcd115deb38a5cba465d4 Mon Sep 17 00:00:00 2001 From: Karl Tarbe Date: Wed, 30 Oct 2024 14:37:33 -0700 Subject: [PATCH] Implement client side for testing --- Package.resolved | 6 +- Package.swift | 14 +- .../Controllers/PrivacyPassController.swift | 2 +- .../PIRClient+KeyRotation.swift | 111 ++++++++++ .../PIRClient+PrivacyPass.swift | 90 ++++++++ Sources/PIRServiceTesting/PIRClient.swift | 207 ++++++++++++++++++ .../PIRServiceTesting/PIRClientError.swift | 28 +++ .../PIRConfig+ShardConfig.swift | 55 +++++ .../TestClientProtocol+Protobuf.swift | 48 ++++ Tests/PIRServiceTests/ExampleUsecase.swift | 41 ++++ .../PIRServiceControllerTests.swift | 128 +---------- Tests/PIRServiceTests/PIRServiceTests.swift | 56 +++++ .../TestClientProtocol+Protobuf.swift | 2 + 13 files changed, 655 insertions(+), 133 deletions(-) create mode 100644 Sources/PIRServiceTesting/PIRClient+KeyRotation.swift create mode 100644 Sources/PIRServiceTesting/PIRClient+PrivacyPass.swift create mode 100644 Sources/PIRServiceTesting/PIRClient.swift create mode 100644 Sources/PIRServiceTesting/PIRClientError.swift create mode 100644 Sources/PIRServiceTesting/PIRConfig+ShardConfig.swift create mode 100644 Sources/PIRServiceTesting/TestClientProtocol+Protobuf.swift create mode 100644 Tests/PIRServiceTests/ExampleUsecase.swift create mode 100644 Tests/PIRServiceTests/PIRServiceTests.swift diff --git a/Package.resolved b/Package.resolved index 244db9c..93a45ba 100644 --- a/Package.resolved +++ b/Package.resolved @@ -1,5 +1,5 @@ { - "originHash" : "e8f3c77635a38c01853b4b0ecb7ce6dc069868602c5669bc662ad2c5148b86b2", + "originHash" : "0f6313dce696999312478d28b401a6688986a9f24af45143c83d80689a7952b9", "pins" : [ { "identity" : "async-http-client", @@ -132,7 +132,7 @@ "kind" : "remoteSourceControl", "location" : "https://github.com/apple/swift-homomorphic-encryption", "state" : { - "revision" : "7091583923d9e25ec760e8479b56f6aefd7fa6d5" + "revision" : "b73daaca802e16c9f6a31da76f26375c34896c15" } }, { @@ -210,7 +210,7 @@ { "identity" : "swift-numerics", "kind" : "remoteSourceControl", - "location" : "https://github.com/apple/swift-numerics.git", + "location" : "https://github.com/apple/swift-numerics", "state" : { "revision" : "0a5bc04095a675662cf24757cc0640aa2204253b", "version" : "1.0.2" diff --git a/Package.swift b/Package.swift index d7f3b8c..7796c7a 100644 --- a/Package.swift +++ b/Package.swift @@ -28,6 +28,7 @@ let package = Package( .executable(name: "PIRService", targets: ["PIRService"]), .executable(name: "ConstructDatabase", targets: ["ConstructDatabase"]), .library(name: "PrivacyPass", targets: ["PrivacyPass"]), + .library(name: "PIRServiceTesting", targets: ["PIRServiceTesting"]), ], dependencies: [ .package(url: "https://github.com/apple/swift-argument-parser.git", from: "1.2.0"), @@ -35,7 +36,7 @@ let package = Package( .package(url: "https://github.com/apple/swift-crypto.git", from: "3.5.0"), .package( url: "https://github.com/apple/swift-homomorphic-encryption", - revision: "7091583923d9e25ec760e8479b56f6aefd7fa6d5"), + revision: "b73daaca802e16c9f6a31da76f26375c34896c15"), .package(url: "https://github.com/apple/swift-protobuf.git", from: "1.27.0"), .package(url: "https://github.com/hummingbird-project/hummingbird", from: "2.0.0"), .package(url: "https://github.com/hummingbird-project/hummingbird-compression", from: "2.0.0-rc.2"), @@ -52,13 +53,13 @@ let package = Package( .product(name: "Hummingbird", package: "hummingbird"), .product(name: "HummingbirdCompression", package: "hummingbird-compression"), .product(name: "PrivateInformationRetrievalProtobuf", package: "swift-homomorphic-encryption"), - .product(name: "SwiftASN1", package: "swift-asn1"), ], swiftSettings: swiftSettings), .testTarget( name: "PIRServiceTests", dependencies: [ "PIRService", + "PIRServiceTesting", .product(name: "HummingbirdTesting", package: "hummingbird"), ], swiftSettings: swiftSettings), @@ -91,6 +92,15 @@ let package = Package( "TestVectors/PrivacyPassChallengeAndRedemptionStructure.json", ], swiftSettings: swiftSettings), + .target( + name: "PIRServiceTesting", + dependencies: [ + "PrivacyPass", + .product(name: "HomomorphicEncryptionProtobuf", package: "swift-homomorphic-encryption"), + .product(name: "HummingbirdTesting", package: "hummingbird"), + .product(name: "PrivateInformationRetrievalProtobuf", package: "swift-homomorphic-encryption"), + ], + swiftSettings: swiftSettings), ]) #if canImport(Darwin) diff --git a/Sources/PIRService/Controllers/PrivacyPassController.swift b/Sources/PIRService/Controllers/PrivacyPassController.swift index f9921c5..816a6bb 100644 --- a/Sources/PIRService/Controllers/PrivacyPassController.swift +++ b/Sources/PIRService/Controllers/PrivacyPassController.swift @@ -29,7 +29,7 @@ struct PrivacyPassController { guard let userToken = request.headers.bearerToken, let userTier = try await state.userAuthenticator.authenticate(userToken: userToken) else { - throw HTTPError(.unauthorized) + throw HTTPError(.unauthorized, message: "User token is unauthorized") } return userTier } diff --git a/Sources/PIRServiceTesting/PIRClient+KeyRotation.swift b/Sources/PIRServiceTesting/PIRClient+KeyRotation.swift new file mode 100644 index 0000000..3b3b1ab --- /dev/null +++ b/Sources/PIRServiceTesting/PIRClient+KeyRotation.swift @@ -0,0 +1,111 @@ +// Copyright 2024 Apple Inc. and the Swift Homomorphic Encryption project authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +import Foundation +import HomomorphicEncryption +import PrivateInformationRetrieval +import PrivateInformationRetrievalProtobuf +import SwiftProtobuf + +extension PIRClient { + mutating func fetchKeyStatus(for usecase: String) async throws + -> Apple_SwiftHomomorphicEncryption_Api_Shared_V1_KeyStatus + { + let configRequest = Apple_SwiftHomomorphicEncryption_Api_Pir_V1_ConfigRequest.with { configRequest in + configRequest.usecases = [usecase] + configRequest.existingConfigIds = [configCache[usecase]?.configurationId ?? Data()] + } + + let configResponse: Apple_SwiftHomomorphicEncryption_Api_Pir_V1_ConfigResponse = try await post( + path: "/config", + body: configRequest) + guard let config = configResponse.configs[usecase] else { + throw PIRClientError.missingConfiguration + } + configCache[usecase] = Configuration(config: config.pirConfig, configurationId: config.configID) + + guard let keyStatus = configResponse.keyInfo.first else { + throw PIRClientError.missingKeyStatus + } + return keyStatus + } + + mutating func rotateKey(for usecase: String) async throws { + let keyStatus = try await fetchKeyStatus(for: usecase) + guard let config = configCache[usecase]?.config else { + throw PIRClientError.missingConfiguration + } + + if let storedSecretKey = secretKeys[config.evaluationKeyConfigHash], + storedSecretKey.timestamp == keyStatus.timestamp + { + // we do not need to rotate + return + } + + let context = try Context(encryptionParameters: config.encryptionParameters.native()) + let secretKey = try context.generateSecretKey() + let storedSecretKey = StoredSecretKey(secretKey: secretKey.serialize()) + let evaluationKey = try context.generateEvaluationKey(config: keyStatus.keyConfig.native(), using: secretKey) + secretKeys[config.evaluationKeyConfigHash] = storedSecretKey + + let evaluationKeyWithMetadata = Apple_SwiftHomomorphicEncryption_Api_Shared_V1_EvaluationKey.with { evalKey in + evalKey.metadata = .with { metadata in + metadata.timestamp = storedSecretKey.timestamp + metadata.identifier = config.evaluationKeyConfigHash + } + evalKey.evaluationKey = evaluationKey.serialize().proto() + } + + try await uploadKey(evaluationKeyWithMetadata) + } + + mutating func uploadKey(_ key: Apple_SwiftHomomorphicEncryption_Api_Shared_V1_EvaluationKey) async throws { + let keys = Apple_SwiftHomomorphicEncryption_Api_Shared_V1_EvaluationKeys.with { keyRequest in + keyRequest.keys = [key] + } + + let _: EmptyProtobufMessage = try await post(path: "/key", body: keys) + } +} + +/// Empty message +private struct EmptyProtobufMessage: Message { + static let protoMessageName = "Empty" + + var unknownFields = SwiftProtobuf.UnknownStorage() + + static func == (lhs: EmptyProtobufMessage, rhs: EmptyProtobufMessage) -> Bool { + if lhs.unknownFields != rhs.unknownFields { + return false + } + return true + } + + mutating func decodeMessage(decoder: inout some SwiftProtobuf.Decoder) throws { + // Load everything into unknown fields + while try decoder.nextFieldNumber() != nil {} + } + + func traverse(visitor: inout some SwiftProtobuf.Visitor) throws { + try unknownFields.traverse(visitor: &visitor) + } + + func isEqualTo(message: any SwiftProtobuf.Message) -> Bool { + guard let other = message as? EmptyProtobufMessage else { + return false + } + return self == other + } +} diff --git a/Sources/PIRServiceTesting/PIRClient+PrivacyPass.swift b/Sources/PIRServiceTesting/PIRClient+PrivacyPass.swift new file mode 100644 index 0000000..9d9e55b --- /dev/null +++ b/Sources/PIRServiceTesting/PIRClient+PrivacyPass.swift @@ -0,0 +1,90 @@ +// Copyright 2024 Apple Inc. and the Swift Homomorphic Encryption project authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +import Foundation +import PrivacyPass + +extension PIRClient { + func fetchTokenDirectory() async throws -> TokenIssuerDirectory { + let response = try await connection.get( + path: "/.well-known/private-token-issuer-directory", + body: [], + headers: [:]) + + let body = Data(buffer: response.body) + guard response.status == .ok else { + throw PIRClientError.failedToFetchTokenIssuerDirectory( + status: response.status, + message: String(data: body, encoding: .utf8) ?? "<\(body.count) bytes of binary response>") + } + return try JSONDecoder().decode(TokenIssuerDirectory.self, from: body) + } + + func fetchPublicKeyForUserToken(authenticationToken: String) async throws -> PublicKey { + let response = try await connection.get( + path: "/token-key-for-user-token", + body: [], + headers: [.authorization: "Bearer \(authenticationToken)"]) + + let body = Array(buffer: response.body) + guard response.status == .ok else { + throw PIRClientError.failedToFetchTokenPublicKey( + status: response.status, + message: String(data: Data(body), encoding: .utf8) ?? "<\(body.count) bytes of binary response>") + } + + return try PublicKey(fromSPKI: body) + } + + mutating func fetchTokens(count: Int) async throws { + guard let userToken else { + throw PIRClientError.missingAuthenticationToken + } + + let tokenIssuerDirectory = try await fetchTokenDirectory() + let publicKey = try await fetchPublicKeyForUserToken(authenticationToken: userToken) + + guard try tokenIssuerDirectory.isValid(tokenKey: publicKey.spki()) else { + throw PIRClientError.invalidTokenIssuerPublicKey + } + + let connection = connection + let challenge = try TokenChallenge(tokenType: TokenTypeBlindRSA, issuer: "test") + + try await withThrowingTaskGroup(of: Token.self) { group in + for _ in 0..") + } + let tokenResponse = try TokenResponse(from: body) + return try preparedRequest.finalize(response: tokenResponse) + } + } + + for try await token in group { + tokens.append(token) + } + } + } +} diff --git a/Sources/PIRServiceTesting/PIRClient.swift b/Sources/PIRServiceTesting/PIRClient.swift new file mode 100644 index 0000000..08dc8c3 --- /dev/null +++ b/Sources/PIRServiceTesting/PIRClient.swift @@ -0,0 +1,207 @@ +// Copyright 2024 Apple Inc. and the Swift Homomorphic Encryption project authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +import Foundation +import HomomorphicEncryption +import HTTPTypes +import HummingbirdTesting +import PrivacyPass +import PrivateInformationRetrieval +import PrivateInformationRetrievalProtobuf +import SwiftProtobuf + +/// PIRClient useful for testing `PIRService`. +public struct PIRClient { + public typealias EvaluationKeyConfigHash = Data + + /// Configuration that will cached by the client. + public struct Configuration: Hashable, Sendable { + /// PIR configuration. + public let config: Apple_SwiftHomomorphicEncryption_Api_Pir_V1_PIRConfig + /// Configuration identifier, typically a hash of the PIR configuration. + public let configurationId: Data + + /// Iniitialize a new configuration. + /// - Parameters: + /// - config: PIR configuration. + /// - configurationId: Configuration identifier, typically a hash of the PIR configuration. + public init(config: Apple_SwiftHomomorphicEncryption_Api_Pir_V1_PIRConfig, configurationId: Data) { + self.config = config + self.configurationId = configurationId + } + } + + /// Secret key that will be stored on the client. + public struct StoredSecretKey { + /// Secret key. + public let secretKey: SerializedSecretKey + /// Key generation timestamp, seconds from UNIX epoch. + public let timestamp: UInt64 + + /// Initialize a new stored secret key. + /// - Parameters: + /// - secretKey: Secret key. + /// - timestamp: Key generation timestamp, seconds from UNIX epoch. + public init(secretKey: SerializedSecretKey, timestamp: Date = Date.now) { + self.secretKey = secretKey + self.timestamp = UInt64(timestamp.timeIntervalSince1970) + } + } + + let connection: TestClientProtocol + + /// User identifier. + public var userID = UUID() + /// Configuration cache. + public var configCache: [String: Configuration] + /// Stored secret keys. + public var secretKeys: [EvaluationKeyConfigHash: StoredSecretKey] + /// Privacy pass tokens. + public var tokens: [Token] + /// User token for requesting privacy pass tokens. + public var userToken: String? + + /// Initialize a new testing client. + /// - Parameters: + /// - connection: Connection to the service under test. + /// - userID: User identifier. + /// - configCache: Configuration cache. + /// - secretKeys: Stored secret keys. + /// - tokens: Privacy pass tokens. + /// - userToken: User token for requesting privacy pass tokens. + public init( + connection: TestClientProtocol, + userID: UUID = UUID(), + configCache: [String: Configuration] = [:], + secretKeys: [EvaluationKeyConfigHash: StoredSecretKey] = [:], + tokens: [Token] = [], + userToken: String? = nil) + { + self.connection = connection + self.userID = userID + self.configCache = configCache + self.secretKeys = secretKeys + self.tokens = tokens + self.userToken = userToken + } + + /// Request a value from the service. + /// + /// When `allowKeyRotation` is `true`, this will also: + /// - Fetch the configuration if there’s none. + /// - Generate a new secret key and evaluation key and upload the evaluation key to the server if there’s none. + /// - Parameters: + /// - keywords: Keywords to request values for. + /// - usecase: Name of the use case to query. + /// - allowKeyRotation: Allow fetching missing configuration and uploading a new evaluation key when needed. + /// - Returns: An array with the same length as `keywords`, where each element is either the corresponding value or + /// `nil` to indicate the absence of a value. + public mutating func request( + keywords: [KeywordValuePair.Keyword], + usecase: String, + allowKeyRotation: Bool = true) async throws -> [KeywordValuePair.Value?] + { + guard let configuration = configCache[usecase] else { + if allowKeyRotation { + try await rotateKey(for: usecase) + return try await request(keywords: keywords, usecase: usecase, allowKeyRotation: false) + } + throw PIRClientError.missingConfiguration + } + + let config = configuration.config + guard let storedSecretKey = secretKeys[config.evaluationKeyConfigHash] else { + if allowKeyRotation { + try await rotateKey(for: usecase) + return try await request(keywords: keywords, usecase: usecase, allowKeyRotation: false) + } + throw PIRClientError.missingSecretKey(evaluationKeyConfigHash: Array(config.evaluationKeyConfigHash)) + } + + let context = try Context(encryptionParameters: config.encryptionParameters.native()) + let secretKey = try SecretKey(deserialize: storedSecretKey.secretKey, context: context) + + let pirRequests: [Apple_SwiftHomomorphicEncryption_Api_Pir_V1_PIRRequest] = try keywords.map { keyword in + let client = try keywordPIRClient(for: keyword, config: config, context: context) + let query = try client.generateQuery(at: keyword, using: secretKey) + return try Apple_SwiftHomomorphicEncryption_Api_Pir_V1_PIRRequest.with { pirRequest in + pirRequest.shardIndex = try UInt32(config.shardindex(for: keyword)) + pirRequest.query = try query.proto() + pirRequest.evaluationKeyMetadata = .with { evaluationKeyMetadata in + evaluationKeyMetadata.timestamp = storedSecretKey.timestamp + evaluationKeyMetadata.identifier = config.evaluationKeyConfigHash + } + pirRequest.configurationHash = configuration.configurationId + } + } + + let requests = Apple_SwiftHomomorphicEncryption_Api_Pir_V1_Requests.with { requests in + requests.requests = pirRequests.map { pirRequest in + Apple_SwiftHomomorphicEncryption_Api_Pir_V1_Request.with { request in + request.usecase = usecase + request.pirRequest = pirRequest + } + } + } + + let responses: Apple_SwiftHomomorphicEncryption_Api_Pir_V1_Responses = try await post( + path: "/queries", + body: requests) + + return try zip(keywords, responses.responses).map { keyword, response in + let client = try keywordPIRClient(for: keyword, config: config, context: context) + return try client.decrypt( + response: response.pirResponse.native(context: context), + at: keyword, + using: secretKey) + } + } + + private func keywordPIRClient( + for keyword: KeywordValuePair.Keyword, + config: Apple_SwiftHomomorphicEncryption_Api_Pir_V1_PIRConfig, + context: Context) throws -> KeywordPirClient + { + let shardIndex = try config.shardindex(for: keyword) + let shardConfig = config.shardConfig(shardIndex: shardIndex) + let evaluationKeyConfig = EvaluationKeyConfig() + return KeywordPirClient( + keywordParameter: config.keywordPirParams.native(), + pirParameter: shardConfig.native( + batchSize: Int(config.batchSize), + evaluationKeyConfig: evaluationKeyConfig), + context: context) + } + + mutating func post(path: String, body: some Message) async throws -> Response { + var headers = HTTPFields() + headers[.userIdentifier] = userID.uuidString + if userToken != nil { + if tokens.isEmpty { + // Note: actual device behaviour is more complex than just fetching 4 tokens. + // Device implementation is subject to change, but as of iOS 18.0 the implementation is like this: + // If there are no tokens, fetch 1 token plus 3 extra tokens. + // If after using a token, there are fewer than 5 tokens left, schedule a background task to fetch + // enough tokens to have 10 tokens cached. The backgound task should run in `5 + random(in: 0…60)` + // seconds. + // Tokens cached for more than 24 hours are considered expired and removed from the cache. + try await fetchTokens(count: 4) + } + + let token = tokens.removeFirst() + headers[.authorization] = "PrivateToken token=\(token.bytes().base64URLEncodedString())" + } + return try await connection.post(path: path, body: body, headers: headers) + } +} diff --git a/Sources/PIRServiceTesting/PIRClientError.swift b/Sources/PIRServiceTesting/PIRClientError.swift new file mode 100644 index 0000000..1391c0a --- /dev/null +++ b/Sources/PIRServiceTesting/PIRClientError.swift @@ -0,0 +1,28 @@ +// Copyright 2024 Apple Inc. and the Swift Homomorphic Encryption project authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +import HTTPTypes + +enum PIRClientError: Error { + case failedToFetchToken(status: HTTPResponse.Status, message: String) + case failedToFetchTokenIssuerDirectory(status: HTTPResponse.Status, message: String) + case failedToFetchTokenPublicKey(status: HTTPResponse.Status, message: String) + case invalidTokenIssuerPublicKey + case missingAuthenticationToken + case missingConfiguration + case missingKeyStatus + case missingSecretKey(evaluationKeyConfigHash: [UInt8]) + case serverError(status: HTTPResponse.Status, message: String) + case unknownShardingFunction(shardingFunction: String) +} diff --git a/Sources/PIRServiceTesting/PIRConfig+ShardConfig.swift b/Sources/PIRServiceTesting/PIRConfig+ShardConfig.swift new file mode 100644 index 0000000..d93c7ab --- /dev/null +++ b/Sources/PIRServiceTesting/PIRConfig+ShardConfig.swift @@ -0,0 +1,55 @@ +// Copyright 2024 Apple Inc. and the Swift Homomorphic Encryption project authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +import PrivateInformationRetrieval +import PrivateInformationRetrievalProtobuf + +extension Apple_SwiftHomomorphicEncryption_Api_Pir_V1_PIRConfig { + var shardCount: Int { + if let pirShardConfigs = pirShardConfigs.shardConfigs { + switch pirShardConfigs { + case let .repeatedShardConfig(repeatedConfig): + return Int(repeatedConfig.shardCount) + } + } + return shardConfigs.count + } + + func shardConfig(shardIndex: Int) -> Apple_SwiftHomomorphicEncryption_Api_Pir_V1_PIRShardConfig { + if let pirShardConfigs = pirShardConfigs.shardConfigs { + switch pirShardConfigs { + case let .repeatedShardConfig(repeatedConfig): + return repeatedConfig.shardConfig + } + } + return shardConfigs[shardIndex] + } + + func shardindex(for keyword: KeywordValuePair.Keyword) throws -> Int { + if keywordPirParams.hasShardingFunction { + switch keywordPirParams.shardingFunction.function { + case .sha256: + return keyword.shardIndex(shardCount: shardCount) + case let .doubleMod(doubleMod): + let otherShardIndex = keyword.shardIndex(shardCount: Int(doubleMod.otherShardCount)) + return otherShardIndex % shardCount + default: + throw PIRClientError + .unknownShardingFunction(shardingFunction: keywordPirParams.shardingFunction.textFormatString()) + } + } + + return keyword.shardIndex(shardCount: shardCount) + } +} diff --git a/Sources/PIRServiceTesting/TestClientProtocol+Protobuf.swift b/Sources/PIRServiceTesting/TestClientProtocol+Protobuf.swift new file mode 100644 index 0000000..fca02be --- /dev/null +++ b/Sources/PIRServiceTesting/TestClientProtocol+Protobuf.swift @@ -0,0 +1,48 @@ +// Copyright 2024 Apple Inc. and the Swift Homomorphic Encryption project authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +import Foundation +import HTTPTypes +import HummingbirdTesting +import SwiftProtobuf + +extension TestClientProtocol { + func post(path: String, body: [UInt8], headers: HTTPFields) async throws -> TestResponse { + try await executeRequest(uri: path, method: .post, headers: headers, body: .init(bytes: body)) + } + + func get(path: String, body: [UInt8], headers: HTTPFields) async throws -> TestResponse { + try await executeRequest(uri: path, method: .get, headers: headers, body: .init(bytes: body)) + } + + func post(path: String, body: some Message, headers: HTTPFields) async throws -> Response { + let response = try await executeRequest( + uri: path, + method: .post, + headers: headers, + body: .init(data: body.serializedBytes())) + guard response.status == .ok else { + throw PIRClientError.serverError( + status: response.status, + message: String(data: Data(buffer: response.body), encoding: .utf8) ?? + "<\(response.body.readableBytes) bytes of binary response>") + } + return try Response(serializedBytes: Array(buffer: response.body)) + } +} + +extension HTTPField.Name { + // swiftlint:disable:next force_unwrapping + static var userIdentifier: Self { Self("User-Identifier")! } +} diff --git a/Tests/PIRServiceTests/ExampleUsecase.swift b/Tests/PIRServiceTests/ExampleUsecase.swift new file mode 100644 index 0000000..e2b1730 --- /dev/null +++ b/Tests/PIRServiceTests/ExampleUsecase.swift @@ -0,0 +1,41 @@ +// Copyright 2024 Apple Inc. and the Swift Homomorphic Encryption project authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +import HomomorphicEncryption +@testable import PIRService +import PrivateInformationRetrieval + +enum ExampleUsecase { + /// Usecase where there are keys in the range `0..<100` and the values are equal to keys. + static let hundred: Usecase = // swiftlint:disable:next force_try + try! buildExampleUsecase() + + private static func buildExampleUsecase() throws -> Usecase { + typealias ServerType = KeywordPirServer>> + let databaseRows = (0..<100) + .map { KeywordValuePair(keyword: [UInt8](String($0).utf8), value: [UInt8](String($0).utf8)) } + let context: Context = + try .init(encryptionParameters: .init(from: .n_4096_logq_27_28_28_logt_4)) + let config = try KeywordPirConfig( + dimensionCount: 2, + cuckooTableConfig: .defaultKeywordPir(maxSerializedBucketSize: context.bytesPerPlaintext), + unevenDimensions: false, keyCompression: .noCompression) + let processed = try ServerType.process( + database: databaseRows, + config: config, + with: context) + let shard = try ServerType(context: context, processed: processed) + return PirUsecase(context: context, keywordParams: config.parameter, shards: [shard]) + } +} diff --git a/Tests/PIRServiceTests/PIRServiceControllerTests.swift b/Tests/PIRServiceTests/PIRServiceControllerTests.swift index 9efe244..02cedea 100644 --- a/Tests/PIRServiceTests/PIRServiceControllerTests.swift +++ b/Tests/PIRServiceTests/PIRServiceControllerTests.swift @@ -22,29 +22,6 @@ import PrivateInformationRetrievalProtobuf import XCTest class PIRServiceControllerTests: XCTestCase { - private static var exampleUsecase: Usecase { - // swiftlint:disable:next force_try - try! buildExampleUsecase() - } - - private static func buildExampleUsecase() throws -> Usecase { - typealias ServerType = KeywordPirServer>> - let databaseRows = (0..<100) - .map { KeywordValuePair(keyword: [UInt8](String($0).utf8), value: [UInt8](String($0).utf8)) } - let context: Context = - try .init(encryptionParameters: .init(from: .n_4096_logq_27_28_28_logt_4)) - let config = try KeywordPirConfig( - dimensionCount: 2, - cuckooTableConfig: .defaultKeywordPir(maxSerializedBucketSize: context.bytesPerPlaintext), - unevenDimensions: false, keyCompression: .noCompression) - let processed = try ServerType.process( - database: databaseRows, - config: config, - with: context) - let shard = try ServerType(context: context, processed: processed) - return PirUsecase(context: context, keywordParams: config.parameter, shards: [shard]) - } - func testNoUserIdentifier() async throws { // Error message returned by Hummingbird struct ErrorMessage: Codable { @@ -98,7 +75,7 @@ class PIRServiceControllerTests: XCTestCase { func testConfigFetch() async throws { let usecaseStore = UsecaseStore() - let exampleUsecase = Self.exampleUsecase + let exampleUsecase = ExampleUsecase.hundred await usecaseStore.set(name: "test", usecase: exampleUsecase) let app = try await buildApplication(usecaseStore: usecaseStore) let user = UserIdentifier() @@ -187,107 +164,4 @@ class PIRServiceControllerTests: XCTestCase { } } } - - func testRequest() async throws { - typealias Scheme = Bfv - let usecaseStore = UsecaseStore() - let exampleUsecase = Self.exampleUsecase - await usecaseStore.set(name: "test", usecase: exampleUsecase) - let app = try await buildApplication(usecaseStore: usecaseStore) - let user = UserIdentifier() - // swiftlint:disable:next closure_body_length - try await app.test(.live) { client in - let context: Context = try .init(encryptionParameters: .init(from: .n_4096_logq_27_28_28_logt_4)) - - // MARK: get configuration - - var config = Apple_SwiftHomomorphicEncryption_Api_Pir_V1_PIRConfig() - var evaluationKeyConfig = Apple_SwiftHomomorphicEncryption_V1_EvaluationKeyConfig() - - try await client.execute( - uri: "/config", - userIdentifier: user, - message: Apple_SwiftHomomorphicEncryption_Api_Pir_V1_ConfigRequest()) - { response in - XCTAssertEqual(response.status, .ok) - let configResponse = try response - .message(as: Apple_SwiftHomomorphicEncryption_Api_Pir_V1_ConfigResponse.self) - - config = try XCTUnwrap(configResponse.configs["test"]).pirConfig - evaluationKeyConfig = configResponse.keyInfo[0].keyConfig - } - - let shardConfig = config.shardConfigs[0] - let keywordConfig: KeywordPirConfig = try .init( - dimensionCount: shardConfig.dimensions.count, - cuckooTableConfig: CuckooTableConfig - .defaultKeywordPir(maxSerializedBucketSize: context.bytesPerPlaintext), - unevenDimensions: false, keyCompression: .noCompression) - let pirParameter = shardConfig.native( - batchSize: Int(config.batchSize), - evaluationKeyConfig: evaluationKeyConfig.native()) - let keywordPirClient: KeywordPirClient> = .init( - keywordParameter: keywordConfig.parameter, - pirParameter: pirParameter, - context: context) - - // MARK: upload evaluation key - - let secretKey = try Scheme.generateSecretKey(context: context) - let evaluationKey = try keywordPirClient.generateEvaluationKey(using: secretKey) - - let serializedEvalKey = evaluationKey.serialize().proto() - let evalKeyMetadata = try Apple_SwiftHomomorphicEncryption_Api_Shared_V1_EvaluationKeyMetadata - .with { metadata in - metadata.timestamp = UInt64(Date.now.timeIntervalSince1970) - metadata.identifier = try evaluationKeyConfig.sha256() - } - let evalKey = Apple_SwiftHomomorphicEncryption_Api_Shared_V1_EvaluationKey.with { evalKey in - evalKey.metadata = evalKeyMetadata - evalKey.evaluationKey = serializedEvalKey - } - let evaluationKeys = Apple_SwiftHomomorphicEncryption_Api_Shared_V1_EvaluationKeys.with { evalKeys in - evalKeys.keys = [evalKey] - } - - try await client.execute(uri: "/key", userIdentifier: user, message: evaluationKeys) { response in - XCTAssertEqual(response.status, .ok) - } - - // MARK: query - - let queryKeyword = [UInt8]("23".utf8) - - let query = try keywordPirClient.generateQuery(at: queryKeyword, using: secretKey) - - let pirRequest = try Apple_SwiftHomomorphicEncryption_Api_Pir_V1_PIRRequest.with { pirRequest in - pirRequest.shardIndex = 0 - pirRequest.query = try query.proto() - pirRequest.evaluationKeyMetadata = evalKeyMetadata - // TODO: fill other fields? - } - let request = Apple_SwiftHomomorphicEncryption_Api_Pir_V1_Request.with { request in - request.usecase = "test" - request.pirRequest = pirRequest - } - let requests = Apple_SwiftHomomorphicEncryption_Api_Pir_V1_Requests.with { requests in - requests.requests = [request] - } - - var pirResponse = Apple_SwiftHomomorphicEncryption_Api_Pir_V1_PIRResponse() - try await client - .execute(uri: "/queries", userIdentifier: user, message: requests) { response in - XCTAssertEqual(response.status, .ok) - let responses = try response.message(as: Apple_SwiftHomomorphicEncryption_Api_Pir_V1_Responses.self) - pirResponse = responses.responses[0].pirResponse - } - - // MARK: decrypt response - - let response = try pirResponse.native(context: context) - let result = try keywordPirClient.decrypt(response: response, at: queryKeyword, using: secretKey) - - XCTAssertEqual(result, queryKeyword) - } - } } diff --git a/Tests/PIRServiceTests/PIRServiceTests.swift b/Tests/PIRServiceTests/PIRServiceTests.swift new file mode 100644 index 0000000..662e66c --- /dev/null +++ b/Tests/PIRServiceTests/PIRServiceTests.swift @@ -0,0 +1,56 @@ +// Copyright 2024 Apple Inc. and the Swift Homomorphic Encryption project authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +import HomomorphicEncryption +@testable import PIRService +import PIRServiceTesting +import PrivateInformationRetrieval +import XCTest + +class PIRServiceTests: XCTestCase { + func testRequest() async throws { + let usecaseStore = UsecaseStore() + await usecaseStore.set(name: "test", usecase: ExampleUsecase.hundred) + let app = try await buildApplication(usecaseStore: usecaseStore) + try await app.test(.live) { client in + var pirClient = PIRClient>>(connection: client) + let result = try await pirClient.request(keyword: "23") + XCTAssertEqual(result, "23") + } + } + + func testRequestWithPrivacyPass() async throws { + let usecaseStore = UsecaseStore() + await usecaseStore.set(name: "test", usecase: ExampleUsecase.hundred) + let userAuthenticator = UserAuthenticator() + await userAuthenticator.add(token: "ABCD", tier: .tier1) + let privacyPassState = try PrivacyPassState(userAuthenticator: userAuthenticator) + let app = try await buildApplication(usecaseStore: usecaseStore, privacyPassState: privacyPassState) + try await app.test(.live) { client in + var pirClient = PIRClient>>(connection: client, userToken: "ABCD") + let result = try await pirClient.request(keyword: "23") + XCTAssertEqual(result, "23") + } + } +} + +extension PIRClient { + mutating func request(keyword: String) async throws -> String? { + let response = try await request(keywords: [.init(keyword.utf8)], usecase: "test") + XCTAssertEqual(response.count, 1) + return response[0].map { value in + String(data: Data(value), encoding: .utf8) ?? "<\(value.count) bytes of binary response>" + } + } +} diff --git a/Tests/PIRServiceTests/TestClientProtocol+Protobuf.swift b/Tests/PIRServiceTests/TestClientProtocol+Protobuf.swift index f9e32ba..8e78df2 100644 --- a/Tests/PIRServiceTests/TestClientProtocol+Protobuf.swift +++ b/Tests/PIRServiceTests/TestClientProtocol+Protobuf.swift @@ -15,7 +15,9 @@ import Hummingbird import HummingbirdTesting @testable import PIRService +import PIRServiceTesting import SwiftProtobuf +import XCTest public extension TestClientProtocol { @discardableResult