Add serialization of protocol messages

willow/proto/willow/BUILD

@@ -14,6 +14,7 @@
 
 load("@protobuf//bazel:cc_proto_library.bzl", "cc_proto_library")
 load("@protobuf//bazel:proto_library.bzl", "proto_library")
+load("@protobuf//rust:defs.bzl", "rust_proto_library")
 
 package(
     default_visibility = ["//visibility:public"],
@@ -48,3 +49,22 @@ cc_proto_library(
     name = "input_spec_cc_proto",
     deps = [":input_spec_proto"],
 )
+
+proto_library(
+    name = "messages_proto",
+    srcs = ["messages.proto"],
+    deps = [
+        "//willow/proto/shell:shell_ciphertexts_proto",
+        "//willow/proto/zk:proofs_proto",
+    ],
+)
+
+cc_proto_library(
+    name = "messages_cc_proto",
+    deps = [":messages_proto"],
+)
+
+rust_proto_library(
+    name = "messages_rust_proto",
+    deps = [":messages_proto"],
+)

willow/proto/willow/messages.proto

@@ -0,0 +1,49 @@
+// Copyright 2025 Google LLC
+//
+// 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 secure_aggregation.willow;
+
+import "willow/proto/shell/ciphertexts.proto";
+import "willow/proto/zk/proofs.proto";
+
+option java_multiple_files = true;
+option java_outer_classname = "MessagesProto";
+
+message ClientMessage {
+  ShellKaheCiphertext kahe_ciphertext = 1;
+  ShellAheCiphertext ahe_ciphertext = 2;
+  RlweRelationProofListProto proof = 3;
+  bytes nonce = 4;
+}
+
+message PartialDecryptionRequest {
+  ShellAhePartialDecCiphertext partial_dec_ciphertext = 1;
+}
+
+message PartialDecryptionResponse {
+  ShellAhePartialDecryption partial_decryption = 1;
+}
+
+message CiphertextContribution {
+  ShellKaheCiphertext kahe_ciphertext = 1;
+  ShellAheRecoverCiphertext ahe_recover_ciphertext = 2;
+}
+
+message DecryptionRequestContribution {
+  ShellAhePartialDecCiphertext partial_dec_ciphertext = 1;
+  RlweRelationProofListProto proof = 2;
+  bytes nonce = 3;
+}

willow/src/traits/BUILD

@@ -85,7 +85,13 @@ rust_library(
     deps = [
         ":ahe_traits",
         ":kahe_traits",
+        ":proto_serialization_traits",
         ":vahe_traits",
+        "@protobuf//rust:protobuf",
+        "//shell_wrapper:status",
+        "//willow/proto/shell:shell_ciphertexts_rust_proto",
+        "//willow/proto/willow:messages_rust_proto",
+        "//willow/proto/zk:proofs_rust_proto",
     ],
 )
 

willow/src/traits/messages.rs

@@ -13,9 +13,23 @@
 // limitations under the License.
 
 use ahe_traits::AheBase;
-use kahe_traits::KaheBase;
+use kahe_traits::{HasKahe, KaheBase};
+use messages_rust_proto::{
+    CiphertextContribution as CiphertextContributionProto, ClientMessage as ClientMessageProto,
+    DecryptionRequestContribution as DecryptionRequestContributionProto,
+    PartialDecryptionRequest as PartialDecryptionRequestProto,
+    PartialDecryptionResponse as PartialDecryptionResponseProto,
+};
+use proofs_rust_proto::RlweRelationProofListProto;
+use proto_serialization_traits::{FromProto, ToProto};
+use protobuf::{proto, AsView};
+use shell_ciphertexts_rust_proto::{
+    ShellAheCiphertext, ShellAhePartialDecCiphertext, ShellAhePartialDecryption,
+    ShellAheRecoverCiphertext, ShellKaheCiphertext,
+};
+use status::StatusError;
 use std::fmt::Debug;
-use vahe_traits::VaheBase;
+use vahe_traits::{HasVahe, VaheBase};
 
 pub type DecryptorPublicKeyShare<Vahe: VaheBase> = <Vahe as AheBase>::PublicKeyShare;
 
@@ -30,6 +44,52 @@ pub struct ClientMessage<Kahe: KaheBase, Vahe: VaheBase> {
     pub nonce: Vec<u8>,
 }
 
+impl<'a, C, Kahe, Vahe> ToProto<&'a C> for ClientMessage<Kahe, Vahe>
+where
+    C: HasKahe<Kahe = Kahe> + HasVahe<Vahe = Vahe>,
+    Kahe: KaheBase + 'a,
+    Vahe: VaheBase + 'a,
+    Kahe::Ciphertext: ToProto<&'a Kahe, Proto = ShellKaheCiphertext>,
+    Vahe::Ciphertext: ToProto<&'a Vahe, Proto = ShellAheCiphertext>,
+    Vahe::EncryptionProof: ToProto<Proto = RlweRelationProofListProto>,
+{
+    type Proto = ClientMessageProto;
+
+    fn to_proto(&self, context: &'a C) -> Result<Self::Proto, StatusError> {
+        Ok(proto!(ClientMessageProto {
+            kahe_ciphertext: self.kahe_ciphertext.to_proto(context.kahe())?,
+            ahe_ciphertext: self.ahe_ciphertext.to_proto(context.vahe())?,
+            proof: self.proof.to_proto(())?,
+            nonce: self.nonce.clone(),
+        }))
+    }
+}
+
+impl<'a, C, Kahe, Vahe> FromProto<&'a C> for ClientMessage<Kahe, Vahe>
+where
+    C: HasKahe<Kahe = Kahe> + HasVahe<Vahe = Vahe>,
+    Kahe: KaheBase + 'a,
+    Vahe: VaheBase + 'a,
+    Kahe::Ciphertext: FromProto<&'a Kahe, Proto = ShellKaheCiphertext>,
+    Vahe::Ciphertext: FromProto<&'a Vahe, Proto = ShellAheCiphertext>,
+    Vahe::EncryptionProof: FromProto<Proto = RlweRelationProofListProto>,
+{
+    type Proto = ClientMessageProto;
+
+    fn from_proto(
+        proto: impl AsView<Proxied = Self::Proto>,
+        context: &'a C,
+    ) -> Result<Self, StatusError> {
+        let proto = proto.as_view();
+        Ok(ClientMessage {
+            kahe_ciphertext: Kahe::Ciphertext::from_proto(proto.kahe_ciphertext(), context.kahe())?,
+            ahe_ciphertext: Vahe::Ciphertext::from_proto(proto.ahe_ciphertext(), context.vahe())?,
+            proof: Vahe::EncryptionProof::from_proto(proto.proof(), ())?,
+            nonce: proto.nonce().to_vec(),
+        })
+    }
+}
+
 impl<Kahe: KaheBase, Vahe: VaheBase> Clone for ClientMessage<Kahe, Vahe> {
     fn clone(self: &ClientMessage<Kahe, Vahe>) -> ClientMessage<Kahe, Vahe> {
         ClientMessage {
@@ -46,6 +106,43 @@ pub struct PartialDecryptionRequest<Vahe: VaheBase> {
     pub partial_dec_ciphertext: Vahe::PartialDecCiphertext,
 }
 
+impl<'a, C, Vahe> ToProto<&'a C> for PartialDecryptionRequest<Vahe>
+where
+    C: HasVahe<Vahe = Vahe>,
+    Vahe: VaheBase + 'a,
+    Vahe::PartialDecCiphertext: ToProto<&'a Vahe, Proto = ShellAhePartialDecCiphertext>,
+{
+    type Proto = PartialDecryptionRequestProto;
+
+    fn to_proto(&self, context: &'a C) -> Result<Self::Proto, StatusError> {
+        Ok(proto!(PartialDecryptionRequestProto {
+            partial_dec_ciphertext: self.partial_dec_ciphertext.to_proto(context.vahe())?,
+        }))
+    }
+}
+
+impl<'a, C, Vahe> FromProto<&'a C> for PartialDecryptionRequest<Vahe>
+where
+    C: HasVahe<Vahe = Vahe>,
+    Vahe: VaheBase + 'a,
+    Vahe::PartialDecCiphertext: FromProto<&'a Vahe, Proto = ShellAhePartialDecCiphertext>,
+{
+    type Proto = PartialDecryptionRequestProto;
+
+    fn from_proto(
+        proto: impl AsView<Proxied = Self::Proto>,
+        context: &'a C,
+    ) -> Result<Self, StatusError> {
+        let proto = proto.as_view();
+        Ok(PartialDecryptionRequest {
+            partial_dec_ciphertext: Vahe::PartialDecCiphertext::from_proto(
+                proto.partial_dec_ciphertext(),
+                context.vahe(),
+            )?,
+        })
+    }
+}
+
 /// We manually implement clone for PartialDecryptionRequest because Vahe is not cloneable.
 impl<Vahe: VaheBase> Clone for PartialDecryptionRequest<Vahe> {
     fn clone(self: &PartialDecryptionRequest<Vahe>) -> PartialDecryptionRequest<Vahe> {
@@ -65,12 +162,92 @@ pub struct PartialDecryptionResponse<Vahe: VaheBase> {
     pub partial_decryption: Vahe::PartialDecryption,
 }
 
+impl<'a, C, Vahe> ToProto<&'a C> for PartialDecryptionResponse<Vahe>
+where
+    C: HasVahe<Vahe = Vahe>,
+    Vahe: VaheBase + 'a,
+    Vahe::PartialDecryption: ToProto<&'a Vahe, Proto = ShellAhePartialDecryption>,
+{
+    type Proto = PartialDecryptionResponseProto;
+
+    fn to_proto(&self, context: &'a C) -> Result<Self::Proto, StatusError> {
+        Ok(proto!(PartialDecryptionResponseProto {
+            partial_decryption: self.partial_decryption.to_proto(context.vahe())?,
+        }))
+    }
+}
+
+impl<'a, C, Vahe> FromProto<&'a C> for PartialDecryptionResponse<Vahe>
+where
+    C: HasVahe<Vahe = Vahe>,
+    Vahe: VaheBase + 'a,
+    Vahe::PartialDecryption: FromProto<&'a Vahe, Proto = ShellAhePartialDecryption>,
+{
+    type Proto = PartialDecryptionResponseProto;
+
+    fn from_proto(
+        proto: impl AsView<Proxied = Self::Proto>,
+        context: &'a C,
+    ) -> Result<Self, StatusError> {
+        let proto = proto.as_view();
+        Ok(PartialDecryptionResponse {
+            partial_decryption: Vahe::PartialDecryption::from_proto(
+                proto.partial_decryption(),
+                context.vahe(),
+            )?,
+        })
+    }
+}
+
 /// The part of the client message that the verifier needn't check
 pub struct CiphertextContribution<Kahe: KaheBase, Vahe: VaheBase> {
     pub kahe_ciphertext: Kahe::Ciphertext,
     pub ahe_recover_ciphertext: Vahe::RecoverCiphertext,
 }
 
+impl<'a, C, Kahe, Vahe> ToProto<&'a C> for CiphertextContribution<Kahe, Vahe>
+where
+    C: HasKahe<Kahe = Kahe> + HasVahe<Vahe = Vahe>,
+    Kahe: KaheBase + 'a,
+    Vahe: VaheBase + 'a,
+    Kahe::Ciphertext: ToProto<&'a Kahe, Proto = ShellKaheCiphertext>,
+    Vahe::RecoverCiphertext: ToProto<&'a Vahe, Proto = ShellAheRecoverCiphertext>,
+{
+    type Proto = CiphertextContributionProto;
+
+    fn to_proto(&self, context: &'a C) -> Result<Self::Proto, StatusError> {
+        Ok(proto!(CiphertextContributionProto {
+            kahe_ciphertext: self.kahe_ciphertext.to_proto(context.kahe())?,
+            ahe_recover_ciphertext: self.ahe_recover_ciphertext.to_proto(context.vahe())?,
+        }))
+    }
+}
+
+impl<'a, C, Kahe, Vahe> FromProto<&'a C> for CiphertextContribution<Kahe, Vahe>
+where
+    C: HasKahe<Kahe = Kahe> + HasVahe<Vahe = Vahe>,
+    Kahe: KaheBase + 'a,
+    Vahe: VaheBase + 'a,
+    Kahe::Ciphertext: FromProto<&'a Kahe, Proto = ShellKaheCiphertext>,
+    Vahe::RecoverCiphertext: FromProto<&'a Vahe, Proto = ShellAheRecoverCiphertext>,
+{
+    type Proto = CiphertextContributionProto;
+
+    fn from_proto(
+        proto: impl AsView<Proxied = Self::Proto>,
+        context: &'a C,
+    ) -> Result<Self, StatusError> {
+        let proto = proto.as_view();
+        Ok(CiphertextContribution {
+            kahe_ciphertext: Kahe::Ciphertext::from_proto(proto.kahe_ciphertext(), context.kahe())?,
+            ahe_recover_ciphertext: Vahe::RecoverCiphertext::from_proto(
+                proto.ahe_recover_ciphertext(),
+                context.vahe(),
+            )?,
+        })
+    }
+}
+
 impl<Kahe: KaheBase, Vahe: VaheBase> Clone for CiphertextContribution<Kahe, Vahe> {
     fn clone(&self) -> CiphertextContribution<Kahe, Vahe> {
         CiphertextContribution {
@@ -88,6 +265,49 @@ pub struct DecryptionRequestContribution<Vahe: VaheBase> {
     pub nonce: Vec<u8>,
 }
 
+impl<'a, C, Vahe> ToProto<&'a C> for DecryptionRequestContribution<Vahe>
+where
+    C: HasVahe<Vahe = Vahe>,
+    Vahe: VaheBase + 'a,
+    Vahe::PartialDecCiphertext: ToProto<&'a Vahe, Proto = ShellAhePartialDecCiphertext>,
+    Vahe::EncryptionProof: ToProto<Proto = RlweRelationProofListProto>,
+{
+    type Proto = DecryptionRequestContributionProto;
+
+    fn to_proto(&self, context: &'a C) -> Result<Self::Proto, StatusError> {
+        Ok(proto!(DecryptionRequestContributionProto {
+            partial_dec_ciphertext: self.partial_dec_ciphertext.to_proto(context.vahe())?,
+            proof: self.proof.to_proto(())?,
+            nonce: self.nonce.clone(),
+        }))
+    }
+}
+
+impl<'a, C, Vahe> FromProto<&'a C> for DecryptionRequestContribution<Vahe>
+where
+    C: HasVahe<Vahe = Vahe>,
+    Vahe: VaheBase + 'a,
+    Vahe::PartialDecCiphertext: FromProto<&'a Vahe, Proto = ShellAhePartialDecCiphertext>,
+    Vahe::EncryptionProof: FromProto<Proto = RlweRelationProofListProto>,
+{
+    type Proto = DecryptionRequestContributionProto;
+
+    fn from_proto(
+        proto: impl AsView<Proxied = Self::Proto>,
+        context: &'a C,
+    ) -> Result<Self, StatusError> {
+        let proto = proto.as_view();
+        Ok(DecryptionRequestContribution {
+            partial_dec_ciphertext: Vahe::PartialDecCiphertext::from_proto(
+                proto.partial_dec_ciphertext(),
+                context.vahe(),
+            )?,
+            proof: Vahe::EncryptionProof::from_proto(proto.proof(), ())?,
+            nonce: proto.nonce().to_vec(),
+        })
+    }
+}
+
 impl<Vahe: VaheBase> Clone for DecryptionRequestContribution<Vahe> {
     fn clone(&self) -> DecryptionRequestContribution<Vahe> {
         DecryptionRequestContribution {

willow/tests/BUILD

@@ -39,6 +39,7 @@ rust_test(
         "//willow/src/traits:kahe_traits",
         "//willow/src/traits:messages",
         "//willow/src/traits:prng_traits",
+        "//willow/src/traits:proto_serialization_traits",
         "//willow/src/traits:server_traits",
         "//willow/src/traits:vahe_traits",
         "//willow/src/traits:verifier_traits",