Add more checks for PirAlgorithm

Sources/PIRProcessDatabase/ProcessDatabase.swift

@@ -307,10 +307,10 @@ struct ProcessDatabase: ParsableCommand {
             keywordPirConfig: keywordConfig)
 
         let encryptionParameters = try EncryptionParameters<Scheme>(from: config.rlweParameters)
-        let processArgs = ProcessKeywordDatabase.Arguments<Scheme>(databaseConfig: databaseConfig,
-                                                                   encryptionParameters: encryptionParameters,
-                                                                   algorithm: config.algorithm,
-                                                                   trialsPerShard: config.trialsPerShard)
+        let processArgs = try ProcessKeywordDatabase.Arguments<Scheme>(databaseConfig: databaseConfig,
+                                                                       encryptionParameters: encryptionParameters,
+                                                                       algorithm: config.algorithm,
+                                                                       trialsPerShard: config.trialsPerShard)
 
         var evaluationKeyConfig = EvaluationKeyConfiguration()
         let context = try Context(encryptionParameters: processArgs.encryptionParameters)

Sources/PrivateInformationRetrieval/Error.swift

@@ -33,6 +33,7 @@ public enum PirError: Error, Hashable, Codable {
     case invalidHashBucketEntryValueSize(maxSize: Int)
     case invalidHashBucketSlotCount(maxCount: Int)
     case invalidIndex(index: Int, numberOfEntries: Int)
+    case invalidPirAlgorithm(_ pirAlgorithm: PirAlgorithm)
     case invalidReply(ciphertextCount: Int, expected: Int)
     case invalidResponse(replyCount: Int, expected: Int)
     case invalidSharding(_ description: String)
@@ -94,6 +95,8 @@ extension PirError: LocalizedError {
             "Invalid hash bucket slot count; maximum is \(maxCount)"
         case let .invalidIndex(index, numberOfEntries):
             "Index \(index) should between 0 and \(numberOfEntries)"
+        case let .invalidPirAlgorithm(pirAlgorithm):
+            "Invalid PIR algorithm: \(pirAlgorithm)"
         case let .invalidReply(ciphertextCount, expected):
             "Reply has \(ciphertextCount) ciphertexts, expected \(expected)"
         case let .invalidResponse(replyCount, expected):

Sources/PrivateInformationRetrieval/KeywordDatabase.swift

@@ -279,12 +279,19 @@ public enum ProcessKeywordDatabase {
         ///   - encryptionParameters: Encryption parameters.
         ///   - algorithm: PIR algorithm to process with.
         ///   - trialsPerShard: number of test queries per shard.
+        ///  - Throws: Error upon invalid arguments
         public init(
             databaseConfig: KeywordDatabaseConfig,
             encryptionParameters: EncryptionParameters<Scheme>,
             algorithm: PirAlgorithm,
-            trialsPerShard: Int)
+            trialsPerShard: Int) throws
         {
+            guard trialsPerShard >= 0 else {
+                throw PirError.validationError("trialsPerShard \(trialsPerShard) must be > 0")
+            }
+            guard algorithm == .mulPir else {
+                throw PirError.invalidPirAlgorithm(algorithm)
+            }
             self.databaseConfig = databaseConfig
             self.encryptionParameters = encryptionParameters
             self.algorithm = algorithm
@@ -360,6 +367,9 @@ public enum ProcessKeywordDatabase {
     {
         let keywordConfig = arguments.databaseConfig.keywordPirConfig
         let context = try Context(encryptionParameters: arguments.encryptionParameters)
+        guard arguments.algorithm == .mulPir else {
+            throw PirError.invalidPirAlgorithm(arguments.algorithm)
+        }
         return try KeywordPirServer<MulPirServer<Scheme>>.process(database: shard,
                                                                   config: keywordConfig,
                                                                   with: context)
@@ -457,6 +467,9 @@ public enum ProcessKeywordDatabase {
 
         var processedShards = [String: ProcessedDatabaseWithParameters<Scheme>]()
         for (shardID, shardedDatabase) in keywordDatabase.shards where !shardedDatabase.isEmpty {
+            guard arguments.algorithm == .mulPir else {
+                throw PirError.invalidPirAlgorithm(arguments.algorithm)
+            }
             let processed = try KeywordPirServer<MulPirServer<Scheme>>.process(database: shardedDatabase,
                                                                                config: keywordConfig,
                                                                                with: context)

Tests/PrivateInformationRetrievalTests/KeywordPirTests.swift

@@ -359,6 +359,29 @@ class KeywordPirTests: XCTestCase {
             MulPirServer<Bfv<UInt64>>.self, client: MulPirClient<Bfv<UInt64>>.self)
     }
 
+    func testInvalidArguments() throws {
+        let cuckooConfig = try CuckooTableConfig(
+            hashFunctionCount: 2,
+            maxEvictionCount: 100,
+            maxSerializedBucketSize: 10 * 5,
+            bucketCount: .allowExpansion(expansionFactor: 1.1, targetLoadFactor: 0.5))
+        let keywordConfig = try KeywordPirConfig(
+            dimensionCount: 2,
+            cuckooTableConfig: cuckooConfig,
+            unevenDimensions: true)
+        let databaseConfig = KeywordDatabaseConfig(
+            sharding: Sharding.shardCount(1),
+            keywordPirConfig: keywordConfig)
+        let encryptionParameters = try EncryptionParameters<Bfv<UInt32>>(
+            from: .n_4096_logq_27_28_28_logt_5)
+        XCTAssertThrowsError(try ProcessKeywordDatabase.Arguments(
+            databaseConfig: databaseConfig,
+            encryptionParameters: encryptionParameters,
+            algorithm: PirAlgorithm.aclsPir,
+            trialsPerShard: 1),
+        error: PirError.invalidPirAlgorithm(PirAlgorithm.aclsPir))
+    }
+
     func testSharding() throws {
         func runTest<PirServer: IndexPirServer, PirClient: IndexPirClient>(
             rlweParameters: PredefinedRlweParameters,
@@ -386,7 +409,7 @@ class KeywordPirTests: XCTestCase {
                 keywordPirConfig: keywordConfig)
             let testDatabase = PirTestUtils.getTestTable(rowCount: rowCount, valueSize: valueSize)
 
-            let args = ProcessKeywordDatabase.Arguments(
+            let args = try ProcessKeywordDatabase.Arguments(
                 databaseConfig: databaseConfig,
                 encryptionParameters: encryptionParameters,
                 algorithm: PirAlgorithm.mulPir,