Protocol message versioning (#1140)

* Add version number to subscribe message and check it

* Doccomments

* Comments

* Changelog

CHANGELOG.md

@@ -19,6 +19,9 @@ At the moment this project **does not** adhere to
   was added to the staking extension pallet's `Config` trait.
 - In [#1134](https://github.com/entropyxyz/entropy-core/pull/1134/) the ```no-sync``` option was removed
 
+### Added
+- Protocol message versioning ([#1140](https://github.com/entropyxyz/entropy-core/pull/1140))
+
 ### Changed
 - Use correct key rotation endpoint in OCW ([#1104](https://github.com/entropyxyz/entropy-core/pull/1104))
 - Change attestation flow to be pull based ([#1109](https://github.com/entropyxyz/entropy-core/pull/1109/))

crates/protocol/src/lib.rs

@@ -46,6 +46,13 @@ use synedrion::{
     AuxInfo, ThresholdKeyShare,
 };
 
+/// The current version number of the protocol message format or protocols themselves
+pub const PROTOCOL_MESSAGE_VERSION: u32 = 1;
+
+/// Currently supported protocol message versions for backward compatibility
+/// This must contain the current version
+pub const SUPPORTED_PROTOCOL_MESSAGE_VERSIONS: [u32; 1] = [PROTOCOL_MESSAGE_VERSION];
+
 /// Identifies a party participating in a protocol session
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)]
 pub struct PartyId(AccountId32);

crates/protocol/src/protocol_transport/errors.rs

@@ -65,3 +65,12 @@ pub enum EncryptedConnectionErr {
     #[error("Could not get remote public key")]
     RemotePublicKey,
 }
+
+/// Error when checking supported protocol versions
+#[derive(Debug, Error, PartialEq)]
+pub enum ProtocolVersionMismatchError {
+    #[error("This version of the protocol is newer than ours - we are on version {0}")]
+    VersionTooNew(u32),
+    #[error("This version of the protocol is no longer supported - the oldest we support is {0}")]
+    VersionTooOld(u32),
+}

crates/protocol/src/protocol_transport/subscribe_message.rs

@@ -13,11 +13,13 @@
 // You should have received a copy of the GNU Affero General Public License
 // along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
-use crate::SessionId;
+use crate::{SessionId, PROTOCOL_MESSAGE_VERSION, SUPPORTED_PROTOCOL_MESSAGE_VERSIONS};
 use serde::{Deserialize, Serialize};
 use sp_core::{sr25519, Pair};
 use subxt::utils::AccountId32;
 
+use super::errors::ProtocolVersionMismatchError;
+
 /// A message sent by a party when initiating a websocket connection to participate
 /// in the signing or DKG protcol
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -29,14 +31,20 @@ pub struct SubscribeMessage {
     pub public_key: sr25519::Public,
     /// Signature to authenticate connecting party
     pub signature: sr25519::Signature,
+    /// Specifies the version of the protocol messages which will be used for this session
+    pub version: u32,
 }
 
 impl SubscribeMessage {
     pub fn new(session_id: SessionId, pair: &sr25519::Pair) -> Result<Self, bincode::Error> {
         let session_id_serialized = bincode::serialize(&session_id)?;
-
         let signature = pair.sign(&session_id_serialized);
-        Ok(Self { session_id, public_key: pair.public(), signature })
+        Ok(Self {
+            session_id,
+            public_key: pair.public(),
+            signature,
+            version: PROTOCOL_MESSAGE_VERSION,
+        })
     }
 
     pub fn account_id(&self) -> AccountId32 {
@@ -47,4 +55,49 @@ impl SubscribeMessage {
         let session_id_serialized = bincode::serialize(&self.session_id)?;
         Ok(sr25519::Pair::verify(&self.signature, session_id_serialized, &self.public_key))
     }
+
+    pub fn check_supported(&self) -> Result<(), ProtocolVersionMismatchError> {
+        if self.version > PROTOCOL_MESSAGE_VERSION {
+            Err(ProtocolVersionMismatchError::VersionTooNew(PROTOCOL_MESSAGE_VERSION))
+        } else if !SUPPORTED_PROTOCOL_MESSAGE_VERSIONS.contains(&self.version) {
+            Err(ProtocolVersionMismatchError::VersionTooOld(
+                *SUPPORTED_PROTOCOL_MESSAGE_VERSIONS
+                    .iter()
+                    .min()
+                    .expect("At least one protocol message version must be supported"),
+            ))
+        } else {
+            Ok(())
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_protocol_version_check() {
+        let session_id = SessionId::Dkg { block_number: 0 };
+        let pair = sr25519::Pair::from_seed(&[0; 32]);
+        let subscribe_message = SubscribeMessage::new(session_id.clone(), &pair).unwrap();
+        assert!(subscribe_message.check_supported().is_ok());
+
+        let session_id_serialized = bincode::serialize(&session_id).unwrap();
+        let signature = pair.sign(&session_id_serialized);
+        let mut subscribe_message =
+            SubscribeMessage { session_id, public_key: pair.public(), signature, version: 0 };
+        assert_eq!(
+            subscribe_message.check_supported(),
+            Err(ProtocolVersionMismatchError::VersionTooOld(
+                *SUPPORTED_PROTOCOL_MESSAGE_VERSIONS.iter().min().unwrap()
+            ))
+        );
+
+        subscribe_message.version = 2;
+        assert_eq!(
+            subscribe_message.check_supported(),
+            Err(ProtocolVersionMismatchError::VersionTooNew(PROTOCOL_MESSAGE_VERSION))
+        );
+    }
 }

crates/threshold-signature-server/src/signing_client/errors.rs

@@ -139,6 +139,10 @@ pub enum SubscribeErr {
     UserError(String),
     #[error("Listener: {0}")]
     Listener(#[from] entropy_protocol::errors::ListenerErr),
+    #[error("Protocol version mismatch: {0}")]
+    VersionMismatch(
+        #[from] entropy_protocol::protocol_transport::errors::ProtocolVersionMismatchError,
+    ),
 }
 
 impl IntoResponse for SubscribeErr {

crates/threshold-signature-server/src/signing_client/protocol_transport.rs

@@ -84,6 +84,8 @@ pub async fn open_protocol_connections(
                 .map_err(|e| ProtocolErr::EncryptedConnection(e.to_string()))?;
             let subscribe_response: Result<(), String> = bincode::deserialize(&response_message)?;
             if let Err(error_message) = subscribe_response {
+                // In future versions, we can check here if the error is VersionTooNew(version)
+                // and if possible the downgrade protocol messages used to be backward compatible
                 return Err(ProtocolErr::BadSubscribeMessage(error_message));
             }
 
@@ -158,6 +160,10 @@ async fn handle_initial_incoming_ws_message(
     let msg: SubscribeMessage = bincode::deserialize(&serialized_subscribe_message)?;
     tracing::info!("Got ws connection, with message: {msg:?}");
 
+    // In future versions we may have backwards compatibility with the
+    // old version and be able to recover here
+    msg.check_supported()?;
+
     if !msg.verify()? {
         return Err(SubscribeErr::InvalidSignature("Invalid signature."));
     }