Add Fixed PIR Shard Config (#48)

Package.swift

@@ -47,7 +47,7 @@ let package = Package(
         .executableTarget(
             name: "PIRService",
             dependencies: [
-                "PrivacyPass",
+                "PrivacyPass", "Util",
                 .product(name: "ArgumentParser", package: "swift-argument-parser"),
                 .product(name: "Crypto", package: "swift-crypto"),
                 .product(name: "HomomorphicEncryptionProtobuf", package: "swift-homomorphic-encryption"),
@@ -78,11 +78,13 @@ let package = Package(
         .target(
             name: "PrivacyPass",
             dependencies: [
+                "Util",
                 .product(name: "Crypto", package: "swift-crypto"),
                 .product(name: "SwiftASN1", package: "swift-asn1"),
                 .product(name: "_CryptoExtras", package: "swift-crypto"),
             ],
             swiftSettings: swiftSettings),
+        .target(name: "Util", swiftSettings: swiftSettings),
         .testTarget(
             name: "PrivacyPassTests",
             dependencies: [
@@ -98,12 +100,16 @@ let package = Package(
         .target(
             name: "PIRServiceTesting",
             dependencies: [
-                "PrivacyPass",
+                "PrivacyPass", "Util",
                 .product(name: "HomomorphicEncryptionProtobuf", package: "swift-homomorphic-encryption"),
                 .product(name: "HummingbirdTesting", package: "hummingbird"),
                 .product(name: "PrivateInformationRetrievalProtobuf", package: "swift-homomorphic-encryption"),
             ],
             swiftSettings: swiftSettings),
+        .testTarget(
+            name: "UtilTests",
+            dependencies: ["Util"],
+            swiftSettings: swiftSettings),
     ])
 
 #if canImport(Darwin)

Sources/PIRService/Application+build.swift

@@ -19,11 +19,13 @@ import Hummingbird
 import Logging
 import NIO
 import PrivateInformationRetrieval
+import Util
 
-struct AppContext: IdentifiedRequestContext, AuthenticatedRequestContext, RequestContext {
+struct AppContext: IdentifiedRequestContext, AuthenticatedRequestContext, PlatformRequestContext, RequestContext {
     var coreContext: CoreRequestContextStorage
     var userIdentifier: UserIdentifier
     var userTier: UserTier
+    var platform: Platform?
 
     // override upload size to 10MiB, the default 2MiB limit is too small for some evaluation keys.
     var maxUploadSize: Int {
@@ -32,6 +34,7 @@ struct AppContext: IdentifiedRequestContext, AuthenticatedRequestContext, Reques
 
     init(source: ApplicationRequestContextSource) {
         self.coreContext = .init(source: source)
+        self.platform = nil
         self.userIdentifier = UserIdentifier(identifier: "")
         self.userTier = .tier1
     }

Sources/PIRService/Controllers/PIRServiceController.swift

@@ -17,6 +17,7 @@ import HomomorphicEncryptionProtobuf
 import Hummingbird
 import HummingbirdCompression
 import PrivateInformationRetrievalProtobuf
+import Util
 
 struct PIRServiceController {
     let usecases: UsecaseStore
@@ -28,8 +29,10 @@ struct PIRServiceController {
 
     func addRoutes(to group: RouterGroup<AppContext>) {
         group.add(middleware: ExtractUserIdentifierMiddleware())
+            .add(middleware: ExtractPlatformMiddleware())
             .post("/key", use: key)
             .post("/queries", use: queries)
+            // only `config` uses response compression, since the key and queries are not compressible.
             .add(middleware: ResponseCompressionMiddleware())
             .post("/config", use: config)
     }
@@ -63,7 +66,11 @@ struct PIRServiceController {
         }
 
         let configs = try requestedUsecases.mapValues { usecase in
-            try usecase.config()
+            var config = try usecase.config()
+            if let platform = context.platform {
+                config.makeCompatible(with: platform)
+            }
+            return config
         }
 
         let keyConfigs = try requestedUsecases.values.map { try $0.evaluationKeyConfig() }

Sources/PIRService/Controllers/PirUsecase.swift

@@ -17,6 +17,7 @@ import HomomorphicEncryptionProtobuf
 import Hummingbird
 import PrivateInformationRetrieval
 import PrivateInformationRetrievalProtobuf
+import Util
 
 enum LoadingError: Error {
     case invalidParameters(shard: String, got: String, expected: String)
@@ -97,13 +98,32 @@ struct PirUsecase<PirScheme: IndexPirServer>: Usecase {
     func config() throws -> Apple_SwiftHomomorphicEncryption_Api_Pir_V1_Config {
         var pirConfig = Apple_SwiftHomomorphicEncryption_Api_Pir_V1_PIRConfig()
         pirConfig.encryptionParameters = try context.encryptionParameters.proto()
-        pirConfig.shardConfigs = shards.map { shard in
-            shard.indexPirParameter.proto()
+        guard let firstShard = shards.first else {
+            throw HTTPError(.internalServerError, message: "Empty shards")
         }
+
         pirConfig.keywordPirParams = keywordParams.proto()
         pirConfig.algorithm = .mulPir
-        pirConfig.batchSize = UInt64(shards.first?.indexPirParameter.batchSize ?? 1)
+        pirConfig.batchSize = UInt64(firstShard.indexPirParameter.batchSize)
         pirConfig.evaluationKeyConfigHash = try evaluationKeyConfig().sha256()
+        pirConfig.shardConfigs = shards.map { shard in
+            shard.indexPirParameter.proto()
+        }
+        let allShardsSame = shards.count > 1 && shards.dropFirst().allSatisfy { shard in
+            shard.indexPirParameter == firstShard.indexPirParameter
+        }
+        if allShardsSame {
+            pirConfig.pirShardConfigs = Apple_SwiftHomomorphicEncryption_Api_Pir_V1_PIRShardConfigs
+                .with { shardConfigs in
+                    shardConfigs.repeatedShardConfig = Apple_SwiftHomomorphicEncryption_Api_Pir_V1_PIRFixedShardConfig
+                        .with { fixedShardConfig in
+                            fixedShardConfig.shardCount = UInt32(shards.count)
+                            fixedShardConfig.shardConfig = firstShard.indexPirParameter.proto()
+                        }
+                }
+            pirConfig.shardConfigs = []
+        }
+
         return try Apple_SwiftHomomorphicEncryption_Api_Pir_V1_Config.with { config in
             config.pirConfig = pirConfig
             config.configID = try pirConfig.sha256()

Sources/PIRService/Controllers/Usecase.swift

@@ -14,8 +14,14 @@
 
 import HomomorphicEncryptionProtobuf
 import PrivateInformationRetrievalProtobuf
+import Util
 
 protocol Usecase: Sendable {
+    /// Returns the configuration.
+    ///
+    /// Note: may use features that are not compatible with older platforms.
+    /// ``Apple_SwiftHomomorphicEncryption_Api_Pir_V1_Config/makeCompatible(with:)`` can be used to make the
+    /// configuration compatible with older platforms.
     func config() throws -> Apple_SwiftHomomorphicEncryption_Api_Pir_V1_Config
     func evaluationKeyConfig() throws -> Apple_SwiftHomomorphicEncryption_V1_EvaluationKeyConfig
     func process(

Sources/PIRService/Extensions/PirConfig+compatible.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 PrivateInformationRetrievalProtobuf
+import Util
+
+extension Platform {
+    var supportsPirFixedShardConfig: Bool {
+        switch osType {
+        case .iOS:
+            osVersion >= .init(major: 18, minor: 2)
+        case .macOS:
+            osVersion >= .init(major: 15, minor: 2)
+        default:
+            fatalError("Unsupported platform \(self)")
+        }
+    }
+}
+
+public extension Apple_SwiftHomomorphicEncryption_Api_Pir_V1_Config {
+    /// Makes the configuration compatible with the given platform.
+    /// - Parameter platform: Device platform.
+    mutating func makeCompatible(with platform: Platform) {
+        if !platform.supportsPirFixedShardConfig {
+            // Check for PIRFixedShardConfig, introduced in iOS 18.2
+            switch pirConfig.pirShardConfigs.shardConfigs {
+            case let .repeatedShardConfig(repeatedConfig):
+                pirConfig.shardConfigs = Array(
+                    repeating: repeatedConfig.shardConfig,
+                    count: Int(repeatedConfig.shardCount))
+                pirConfig.clearPirShardConfigs()
+            case .none:
+                break
+            }
+        }
+    }
+}

Sources/PIRService/Middlewares/ExtractPlatform.swift

@@ -0,0 +1,38 @@
+// 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 Hummingbird
+import Util
+
+protocol PlatformRequestContext: RequestContext {
+    var platform: Platform? { get set }
+}
+
+/// Middleware that extracts the platform from the 'User-Agent' header
+struct ExtractPlatformMiddleware<Context: PlatformRequestContext>: RouterMiddleware {
+    func handle(
+        _ input: Request,
+        context: Context,
+        next: (Request, Context) async throws -> Response) async throws -> Response
+    {
+        var context = context
+        guard let userAgent = input.headers[.userAgent] else {
+            throw HTTPError(.badRequest, message: "Missing 'User-Agent' header")
+        }
+        context.platform = Platform(userAgent: userAgent)
+        return try await next(input, context)
+    }
+}

Sources/PIRService/PIRService.docc/HTTPEndpoints.md

@@ -26,6 +26,7 @@ Request        | Value              | Description
 Method         | POST               | HTTP method.
 Path           | `/config`          | HTTP path.
 Header         | `Authorization`    | The value will contain a private access token.
+Header         | `User-Agent`       | Identifier for the user's OS type and version.
 Header         | `User-Identifier`  | Pseudorandom identifier tied to a user.
 Request Body   | `ConfigRequest`    | Serialized Protobuf message that list the use-cases that the system is interested in.
 Response       | `ConfigResponse`   | Serialized Protobuf message. The `ConfigResponse` contains the `configs` and `key_info` response fields.
@@ -43,6 +44,7 @@ Request        | Value              | Description
 Method         | POST               | HTTP method.
 Path           | `/key`             | HTTP path.
 Header         | `Authorization`    | The value will contain a private access token.
+Header         | `User-Agent`       | Identifier for the user's OS type and version.
 Header         | `User-Identifier`  | Pseudorandom identifier tied to a user.
 Body           | `EvaluationKeys`   | Serialized Protobuf message that contains evaluation key(s).
 
@@ -57,6 +59,7 @@ Request        | Value              | Description
 Method         | POST               | HTTP method.
 Path           | `/queries`         | HTTP path.
 Header         | `Authorization`    | The value will contain a private access token.
+Header         | `User-Agent`       | Identifier for the user's OS type and version.
 Header         | `User-Identifier`  | Pseudorandom identifier tied to a user.
 Request Body   | `Requests`         | Serialized Protobuf message.
 Response       | `Responses`        | Serialized Protobuf message.

Sources/PIRService/PIRService.docc/TestingInstructions.md

@@ -261,6 +261,11 @@ By default `PIRService` will start listening on the loopback interface, but you
 make it listen on all network interfaces. The default port is `8080`, but it can be changed by using the `--port`
 option.
 
+### Features
+After introduction in iOS 18.0, `Live Caller ID Lookup` introduced further features.
+
+* `Fixed PIR Shard Config` (iOS 18.2). When all shard configurations are identical, `PIR Fixed Shard Config` allows for a more compact PIR config, saving bandwidth and client-side memory usage. To enable, set the `pirShardConfigs` field in the PIR config. iOS clients prior to iOS 18.2 will still require the `shardConfigs` field to be set. See [Reusing PIR Parameters]( https://swiftpackageindex.com/apple/swift-homomorphic-encryption/main/documentation/privateinformationretrieval/reusingpirparameters) for how to process the database such that all shard configurations are identical.
+
 ## Writing the application extension
 
 > Important: Please see [Getting up-to-date calling and blocking information for your

Sources/PIRServiceTesting/PIRClient.swift

@@ -20,6 +20,7 @@ import PrivacyPass
 import PrivateInformationRetrieval
 import PrivateInformationRetrievalProtobuf
 import SwiftProtobuf
+import Util
 
 /// PIRClient useful for testing `PIRService`.
 public struct PIRClient<PIRClient: IndexPirClient> {
@@ -63,6 +64,8 @@ public struct PIRClient<PIRClient: IndexPirClient> {
 
     /// User identifier.
     public var userID = UUID()
+    /// Platform the device is using.
+    public var platform: Platform
     /// Configuration cache.
     public var configCache: [String: Configuration]
     /// Stored secret keys.
@@ -76,20 +79,23 @@ public struct PIRClient<PIRClient: IndexPirClient> {
     /// - Parameters:
     ///   - connection: Connection to the service under test.
     ///   - userID: User identifier.
+    ///   - platform: Platform the device is using.
     ///   - 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(),
+        platform: Platform = .iOS18,
         configCache: [String: Configuration] = [:],
         secretKeys: [EvaluationKeyConfigHash: StoredSecretKey] = [:],
         tokens: [Token] = [],
         userToken: String? = nil)
     {
         self.connection = connection
         self.userID = userID
+        self.platform = platform
         self.configCache = configCache
         self.secretKeys = secretKeys
         self.tokens = tokens
@@ -136,7 +142,7 @@ public struct PIRClient<PIRClient: IndexPirClient> {
             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.shardIndex = try UInt32(config.shardIndex(for: keyword))
                 pirRequest.query = try query.proto()
                 pirRequest.evaluationKeyMetadata = .with { evaluationKeyMetadata in
                     evaluationKeyMetadata.timestamp = storedSecretKey.timestamp
@@ -173,7 +179,7 @@ public struct PIRClient<PIRClient: IndexPirClient> {
         config: Apple_SwiftHomomorphicEncryption_Api_Pir_V1_PIRConfig,
         context: Context<PIRClient.Scheme>) throws -> KeywordPirClient<PIRClient>
     {
-        let shardIndex = try config.shardindex(for: keyword)
+        let shardIndex = try config.shardIndex(for: keyword)
         let shardConfig = config.shardConfig(shardIndex: shardIndex)
         let evaluationKeyConfig = EvaluationKeyConfig()
         return KeywordPirClient<PIRClient>(
@@ -185,8 +191,10 @@ public struct PIRClient<PIRClient: IndexPirClient> {
     }
 
     mutating func post<Response: Message>(path: String, body: some Message) async throws -> Response {
-        var headers = HTTPFields()
-        headers[.userIdentifier] = userID.uuidString
+        var headers: HTTPFields = [
+            .userIdentifier: userID.uuidString,
+            .userAgent: platform.exampleUserAgent,
+        ]
         if userToken != nil {
             if tokens.isEmpty {
                 // Note: actual device behaviour is more complex than just fetching 4 tokens.

Sources/PIRServiceTesting/PIRConfig+ShardConfig.swift

@@ -14,8 +14,10 @@
 
 import PrivateInformationRetrieval
 import PrivateInformationRetrievalProtobuf
+import Util
 
-extension Apple_SwiftHomomorphicEncryption_Api_Pir_V1_PIRConfig {
+public extension Apple_SwiftHomomorphicEncryption_Api_Pir_V1_PIRConfig {
+    /// The number of shards in the configuration.
     var shardCount: Int {
         if let pirShardConfigs = pirShardConfigs.shardConfigs {
             switch pirShardConfigs {
@@ -26,6 +28,9 @@ extension Apple_SwiftHomomorphicEncryption_Api_Pir_V1_PIRConfig {
         return shardConfigs.count
     }
 
+    /// Returns the shard configuration at an index.
+    /// - Parameter shardIndex: Shard index.
+    /// - Returns: The shard configuration.
     func shardConfig(shardIndex: Int) -> Apple_SwiftHomomorphicEncryption_Api_Pir_V1_PIRShardConfig {
         if let pirShardConfigs = pirShardConfigs.shardConfigs {
             switch pirShardConfigs {
@@ -36,7 +41,10 @@ extension Apple_SwiftHomomorphicEncryption_Api_Pir_V1_PIRConfig {
         return shardConfigs[shardIndex]
     }
 
-    func shardindex(for keyword: KeywordValuePair.Keyword) throws -> Int {
+    /// Computes the shard index for a keyword.
+    /// - Parameter keyword: Keyword.
+    /// - Returns: The shard index for the keyword.
+    func shardIndex(for keyword: KeywordValuePair.Keyword) throws -> Int {
         if keywordPirParams.hasShardingFunction {
             switch keywordPirParams.shardingFunction.function {
             case .sha256: