From 5c9cfe445e48a26c8792fe58849dc0530ee0313b Mon Sep 17 00:00:00 2001
From: Fabian Boemer <fabian_boemer@apple.com>
Date: Thu, 29 Aug 2024 15:26:14 -0700
Subject: [PATCH] Refactor service protos to allow 3 service types. (#11)

The service types:
    * PIR-only
    * PNNS-only
    * PIR and PNNS
---
 .../api/pir/v1/api.proto                      | 80 +++++++++++++++++++
 .../{v1/api_pir.proto => pir/v1/pir.proto}    |  8 +-
 .../api/pnns/v1/api.proto                     | 80 +++++++++++++++++++
 .../{v1/api_pnns.proto => pnns/v1/pnns.proto} |  8 +-
 .../v1/api_shared.proto}                      |  2 +-
 .../api/v1/api.proto                          | 36 +++++----
 6 files changed, 189 insertions(+), 25 deletions(-)
 create mode 100644 apple/swift_homomorphic_encryption/api/pir/v1/api.proto
 rename apple/swift_homomorphic_encryption/api/{v1/api_pir.proto => pir/v1/pir.proto} (89%)
 create mode 100644 apple/swift_homomorphic_encryption/api/pnns/v1/api.proto
 rename apple/swift_homomorphic_encryption/api/{v1/api_pnns.proto => pnns/v1/pnns.proto} (89%)
 rename apple/swift_homomorphic_encryption/api/{v1/api_evaluation_key.proto => shared/v1/api_shared.proto} (96%)

diff --git a/apple/swift_homomorphic_encryption/api/pir/v1/api.proto b/apple/swift_homomorphic_encryption/api/pir/v1/api.proto
new file mode 100644
index 0000000..53b97ba
--- /dev/null
+++ b/apple/swift_homomorphic_encryption/api/pir/v1/api.proto
@@ -0,0 +1,80 @@
+// 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.
+
+syntax = "proto3";
+package apple.swift_homomorphic_encryption.api.pir.v1;
+
+import "apple/swift_homomorphic_encryption/api/pir/v1/pir.proto";
+import "apple/swift_homomorphic_encryption/api/shared/v1/api_shared.proto";
+
+// Request for server side configurations.
+message ConfigRequest {
+  // List of usecases to fetch configs for.
+  // When set to empty array, all configs will be returned.
+  repeated string usecases = 1;
+}
+
+// Usecase configuration.
+message Config {
+  // Configuration.
+  oneof config {
+    // Configuration for a PIR usecase.
+    apple.swift_homomorphic_encryption.api.pir.v1.PIRConfig pir_config = 1;
+  }
+  reserved 2;
+  // Unique identifier for the configuration.
+  bytes config_id = 3;
+}
+
+// Server side configurations.
+message ConfigResponse {
+  // usecases with associated configurations.
+  map<string, Config> configs = 1;
+  // Configuration & status of evaluation keys.
+  repeated apple.swift_homomorphic_encryption.api.shared.v1.KeyStatus key_info = 2;
+}
+
+// Container for multiple requests.
+message Requests {
+  // Requests.
+  repeated Request requests = 1;
+}
+
+// Container for multiple responses.
+message Responses {
+  // Responses.
+  repeated Response responses = 1;
+}
+
+// Generic request.
+message Request {
+  // Usecase identifier.
+  string usecase = 1;
+  // Generic request.
+  oneof request {
+    // PIR request.
+    apple.swift_homomorphic_encryption.api.pir.v1.PIRRequest pir_request = 2;
+  }
+  reserved 3;
+}
+
+// Generic response.
+message Response {
+  // Generic response.
+  oneof response {
+    // Response to a `PIRRequest`.
+    apple.swift_homomorphic_encryption.api.pir.v1.PIRResponse pir_response = 1;
+  }
+  reserved 2;
+}
diff --git a/apple/swift_homomorphic_encryption/api/v1/api_pir.proto b/apple/swift_homomorphic_encryption/api/pir/v1/pir.proto
similarity index 89%
rename from apple/swift_homomorphic_encryption/api/v1/api_pir.proto
rename to apple/swift_homomorphic_encryption/api/pir/v1/pir.proto
index c398e70..6499d9e 100644
--- a/apple/swift_homomorphic_encryption/api/v1/api_pir.proto
+++ b/apple/swift_homomorphic_encryption/api/pir/v1/pir.proto
@@ -13,9 +13,9 @@
 // limitations under the License.
 
 syntax = "proto3";
-package apple.swift_homomorphic_encryption.api.v1;
+package apple.swift_homomorphic_encryption.api.pir.v1;
 
-import "apple/swift_homomorphic_encryption/api/v1/api_evaluation_key.proto";
+import "apple/swift_homomorphic_encryption/api/shared/v1/api_shared.proto";
 import "apple/swift_homomorphic_encryption/pir/v1/pir.proto";
 import "apple/swift_homomorphic_encryption/pir/v1/pir_algorithm.proto";
 import "apple/swift_homomorphic_encryption/v1/he.proto";
@@ -60,13 +60,13 @@ message PIRRequest {
   // Encrypted query.
   apple.swift_homomorphic_encryption.pir.v1.EncryptedIndices query = 2;
   // Evaluation key metadata.
-  EvaluationKeyMetadata evaluation_key_metadata = 3;
+  apple.swift_homomorphic_encryption.api.shared.v1.EvaluationKeyMetadata evaluation_key_metadata = 3;
   // Hash of the `PIRConfig` used to construct the query.
   bytes configuration_hash = 4;
   // If set, route request to a shard with this `shard_id` instead of `shard_index`.
   optional string shard_id = 5;
   // If set, evaluation key to query with. Will override evaluation key stored server-side.
-  optional EvaluationKey evaluation_key = 6;
+  optional apple.swift_homomorphic_encryption.api.shared.v1.EvaluationKey evaluation_key = 6;
 }
 
 // PIR Response.
diff --git a/apple/swift_homomorphic_encryption/api/pnns/v1/api.proto b/apple/swift_homomorphic_encryption/api/pnns/v1/api.proto
new file mode 100644
index 0000000..a9b042e
--- /dev/null
+++ b/apple/swift_homomorphic_encryption/api/pnns/v1/api.proto
@@ -0,0 +1,80 @@
+// 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.
+
+syntax = "proto3";
+package apple.swift_homomorphic_encryption.api.pnns.v1;
+
+import "apple/swift_homomorphic_encryption/api/pnns/v1/pnns.proto";
+import "apple/swift_homomorphic_encryption/api/shared/v1/api_shared.proto";
+
+// Request for server side configurations.
+message ConfigRequest {
+  // List of usecases to fetch configs for.
+  // When set to empty array, all configs will be returned.
+  repeated string usecases = 1;
+}
+
+// Usecase configuration.
+message Config {
+  reserved 1;
+  // Configuration.
+  oneof config {
+    // Configuration for a PNNS usecase.
+    apple.swift_homomorphic_encryption.api.pnns.v1.PNNSConfig pnns_config = 2;
+  }
+  // Unique identifier for the configuration.
+  bytes config_id = 3;
+}
+
+// Server side configurations.
+message ConfigResponse {
+  // usecases with associated configurations.
+  map<string, Config> configs = 1;
+  // Configuration & status of evaluation keys.
+  repeated apple.swift_homomorphic_encryption.api.shared.v1.KeyStatus key_info = 2;
+}
+
+// Container for multiple requests.
+message Requests {
+  // Requests.
+  repeated Request requests = 1;
+}
+
+// Container for multiple responses.
+message Responses {
+  // Responses.
+  repeated Response responses = 1;
+}
+
+// Generic request.
+message Request {
+  // Usecase identifier.
+  string usecase = 1;
+  reserved 2;
+  // Generic request.
+  oneof request {
+    // PNNS request.
+    apple.swift_homomorphic_encryption.api.pnns.v1.PNNSRequest pnns_request = 3;
+  }
+}
+
+// Generic response.
+message Response {
+  reserved 1;
+  // Generic response.
+  oneof response {
+    // Response to a `PNNSRequest`.
+    apple.swift_homomorphic_encryption.api.pnns.v1.PNNSResponse pnns_response = 2;
+  }
+}
diff --git a/apple/swift_homomorphic_encryption/api/v1/api_pnns.proto b/apple/swift_homomorphic_encryption/api/pnns/v1/pnns.proto
similarity index 89%
rename from apple/swift_homomorphic_encryption/api/v1/api_pnns.proto
rename to apple/swift_homomorphic_encryption/api/pnns/v1/pnns.proto
index 9493894..4b5e8a3 100644
--- a/apple/swift_homomorphic_encryption/api/v1/api_pnns.proto
+++ b/apple/swift_homomorphic_encryption/api/pnns/v1/pnns.proto
@@ -13,9 +13,9 @@
 // limitations under the License.
 
 syntax = "proto3";
-package apple.swift_homomorphic_encryption.api.v1;
+package apple.swift_homomorphic_encryption.api.pnns.v1;
 
-import "apple/swift_homomorphic_encryption/api/v1/api_evaluation_key.proto";
+import "apple/swift_homomorphic_encryption/api/shared/v1/api_shared.proto";
 import "apple/swift_homomorphic_encryption/pnns/v1/pnns.proto";
 import "apple/swift_homomorphic_encryption/pnns/v1/pnns_distance_metric.proto";
 import "apple/swift_homomorphic_encryption/pnns/v1/pnns_matrix_packing.proto";
@@ -47,11 +47,11 @@ message PNNSRequest {
   // Encrypted query, one per plaintext CRT component
   repeated apple.swift_homomorphic_encryption.pnns.v1.SerializedCiphertextMatrix query = 2;
   // Key metadata
-  EvaluationKeyMetadata evaluation_key_metadata = 3;
+  apple.swift_homomorphic_encryption.api.shared.v1.EvaluationKeyMetadata evaluation_key_metadata = 3;
   // Identifier for the PNNSConfig used to construct the query
   bytes config_id = 4;
   // If set, evaluation key to query with. Will override evaluation key stored server-side
-  optional EvaluationKey evaluation_key = 5;
+  optional apple.swift_homomorphic_encryption.api.shared.v1.EvaluationKey evaluation_key = 5;
 }
 
 // PNNS Shard Response
diff --git a/apple/swift_homomorphic_encryption/api/v1/api_evaluation_key.proto b/apple/swift_homomorphic_encryption/api/shared/v1/api_shared.proto
similarity index 96%
rename from apple/swift_homomorphic_encryption/api/v1/api_evaluation_key.proto
rename to apple/swift_homomorphic_encryption/api/shared/v1/api_shared.proto
index bdebfc3..df7d42a 100644
--- a/apple/swift_homomorphic_encryption/api/v1/api_evaluation_key.proto
+++ b/apple/swift_homomorphic_encryption/api/shared/v1/api_shared.proto
@@ -13,7 +13,7 @@
 // limitations under the License.
 
 syntax = "proto3";
-package apple.swift_homomorphic_encryption.api.v1;
+package apple.swift_homomorphic_encryption.api.shared.v1;
 
 import "apple/swift_homomorphic_encryption/v1/he.proto";
 
diff --git a/apple/swift_homomorphic_encryption/api/v1/api.proto b/apple/swift_homomorphic_encryption/api/v1/api.proto
index 3f4d1ea..3431d67 100644
--- a/apple/swift_homomorphic_encryption/api/v1/api.proto
+++ b/apple/swift_homomorphic_encryption/api/v1/api.proto
@@ -15,34 +15,36 @@
 syntax = "proto3";
 package apple.swift_homomorphic_encryption.api.v1;
 
-import "apple/swift_homomorphic_encryption/api/v1/api_evaluation_key.proto";
-import "apple/swift_homomorphic_encryption/api/v1/api_pir.proto";
-
-// Request for server side configurations.
-message ConfigRequest {
-  // List of usecases to fetch configs for.
-  // When set to empty array, all configs will be returned.
-  repeated string usecases = 1;
-}
+import "apple/swift_homomorphic_encryption/api/pir/v1/pir.proto";
+import "apple/swift_homomorphic_encryption/api/pnns/v1/pnns.proto";
+import "apple/swift_homomorphic_encryption/api/shared/v1/api_shared.proto";
 
 // Usecase configuration.
 message Config {
   // Configuration.
   oneof config {
     // Configuration for a PIR usecase.
-    PIRConfig pir_config = 1;
+    apple.swift_homomorphic_encryption.api.pir.v1.PIRConfig pir_config = 1;
+    // Configuration for a PNNS usecase.
+    apple.swift_homomorphic_encryption.api.pnns.v1.PNNSConfig pnns_config = 2;
   }
-  reserved 2;
   // Unique identifier for the configuration.
   bytes config_id = 3;
 }
 
+// Request for server side configurations.
+message ConfigRequest {
+  // List of usecases to fetch configs for.
+  // When set to empty array, all configs will be returned.
+  repeated string usecases = 1;
+}
+
 // Server side configurations.
 message ConfigResponse {
   // usecases with associated configurations.
   map<string, Config> configs = 1;
   // Configuration & status of evaluation keys.
-  repeated KeyStatus key_info = 2;
+  repeated apple.swift_homomorphic_encryption.api.shared.v1.KeyStatus key_info = 2;
 }
 
 // Container for multiple requests.
@@ -64,9 +66,10 @@ message Request {
   // Generic request.
   oneof request {
     // PIR request.
-    PIRRequest pir_request = 2;
+    apple.swift_homomorphic_encryption.api.pir.v1.PIRRequest pir_request = 2;
+    // PNNS request.
+    apple.swift_homomorphic_encryption.api.pnns.v1.PNNSRequest pnns_request = 3;
   }
-  reserved 3;
 }
 
 // Generic response.
@@ -74,7 +77,8 @@ message Response {
   // Generic response.
   oneof response {
     // Response to a `PIRRequest`.
-    PIRResponse pir_response = 1;
+    apple.swift_homomorphic_encryption.api.pir.v1.PIRResponse pir_response = 1;
+    // Response to a `PNNSRequest`.
+    apple.swift_homomorphic_encryption.api.pnns.v1.PNNSResponse pnns_response = 2;
   }
-  reserved 2;
 }