Implement client side for testing

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"

Package.swift

@@ -28,14 +28,15 @@ 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"),
         .package(url: "https://github.com/apple/swift-asn1.git", from: "1.0.0"),
         .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)

Sources/PIRService/Controllers/PrivacyPassController.swift

@@ -29,7 +29,7 @@ struct PrivacyPassController<UserAuthenticator: UserTokenAuthenticator> {
         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
     }

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<PIRClient.Scheme>(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
+    }
+}

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..<count {
+                group.addTask {
+                    let preparedRequest = try publicKey.request(challenge: challenge.bytes())
+                    let response = try await connection.post(
+                        path: "/issue",
+                        body: preparedRequest.tokenRequest.bytes(),
+                        headers: [.authorization: "Bearer \(userToken)"])
+                    let body = Array(buffer: response.body)
+                    guard response.status == .ok else {
+                        throw PIRClientError.failedToFetchToken(
+                            status: response.status,
+                            message: String(data: Data(body), encoding: .utf8) ??
+                                "<\(body.count) bytes of binary response>")
+                    }
+                    let tokenResponse = try TokenResponse(from: body)
+                    return try preparedRequest.finalize(response: tokenResponse)
+                }
+            }
+
+            for try await token in group {
+                tokens.append(token)
+            }
+        }
+    }
+}