Wrap message's content into node's message (#306)

backend/src/node/connector.rs

@@ -10,7 +10,7 @@ use actix_http::ws::Item;
 use crate::aggregator::{Aggregator, AddNode};
 use crate::chain::{Chain, UpdateNode, RemoveNode};
 use crate::node::NodeId;
-use crate::node::message::{NodeMessage, Details, SystemConnected};
+use crate::node::message::{NodeMessage, Payload};
 use crate::util::LocateRequest;
 use crate::types::ConnId;
 
@@ -96,7 +96,7 @@ impl NodeConnector {
     }
 
     fn handle_message(&mut self, msg: NodeMessage, data: Bytes, ctx: &mut <Self as Actor>::Context) {
-        let conn_id = msg.id.unwrap_or(0);
+        let conn_id = msg.id();
 
         match self.multiplex.entry(conn_id).or_default() {
             ConnMultiplex::Connected { nid, chain } => {
@@ -107,8 +107,8 @@ impl NodeConnector {
                 });
             }
             ConnMultiplex::Waiting { backlog } => {
-                if let Details::SystemConnected(connected) = msg.details {
-                    let SystemConnected { network_id: _, mut node } = connected;
+                if let Payload::SystemConnected(connected) = msg.payload() {
+                    let mut node = connected.node.clone();
                     let rec = ctx.address().recipient();
 
                     // FIXME: Use genesis hash instead of names to avoid this mess
@@ -231,7 +231,7 @@ impl StreamHandler<Result<ws::Message, ws::ProtocolError>> for NodeConnector {
             #[cfg(debug)]
             Err(err) => {
                 let data: &[u8] = data.get(..512).unwrap_or_else(|| &data);
-                warn!("Failed to parse node message: {} {}", err, std::str::from_utf8(data).unwrap_or_else(|_| "INVALID UTF8"))
+                log::warn!("Failed to parse node message: {} {}", err, std::str::from_utf8(data).unwrap_or_else(|_| "INVALID UTF8"))
             },
             #[cfg(not(debug))]
             Err(_) => (),

backend/src/node/message.rs

@@ -1,5 +1,4 @@
 use actix::prelude::*;
-use chrono::{DateTime, Utc};
 use serde::Deserialize;
 use serde::de::IgnoredAny;
 use crate::node::NodeDetails;
@@ -7,18 +6,38 @@ use crate::types::{Block, BlockNumber, BlockHash, ConnId};
 
 #[derive(Deserialize, Debug, Message)]
 #[rtype(result = "()")]
-pub struct NodeMessage {
-    pub ts: DateTime<Utc>,
-    pub id: Option<ConnId>,
-    #[serde(flatten)]
-    pub details: Details,
+#[serde(untagged)]
+pub enum NodeMessage {
+    V1 {
+        #[serde(flatten)]
+        payload: Payload,
+    },
+    V2 {
+        id: ConnId,
+        payload: Payload,
+    },
+}
+
+impl NodeMessage {
+    /// Returns a reference to the payload.
+    pub fn payload(&self) -> &Payload {
+        match self {
+            NodeMessage::V1 { payload, .. } | NodeMessage::V2 { payload, .. } => payload,
+        }
+    }
+
+    /// Returns the connection ID or 0 if there is no ID.
+    pub fn id(&self) -> ConnId {
+        match self {
+            NodeMessage::V1 { .. } => 0,
+            NodeMessage::V2 { id, .. } => *id,
+        }
+    }
 }
 
 #[derive(Deserialize, Debug)]
 #[serde(tag = "msg")]
-pub enum Details {
-    #[serde(rename = "node.start")]
-    NodeStart(Block),
+pub enum Payload {
     #[serde(rename = "system.connected")]
     SystemConnected(SystemConnected),
     #[serde(rename = "system.interval")]
@@ -124,24 +143,24 @@ impl Block {
     }
 }
 
-impl Details {
+impl Payload {
     pub fn best_block(&self) -> Option<&Block> {
         match self {
-            Details::BlockImport(block) => Some(block),
-            Details::SystemInterval(SystemInterval { block, .. }) => block.as_ref(),
+            Payload::BlockImport(block) => Some(block),
+            Payload::SystemInterval(SystemInterval { block, .. }) => block.as_ref(),
             _ => None,
         }
     }
 
     pub fn finalized_block(&self) -> Option<Block> {
         match self {
-            Details::SystemInterval(ref interval) => {
+            Payload::SystemInterval(ref interval) => {
                 Some(Block {
                     hash: interval.finalized_hash?,
                     height: interval.finalized_height?,
                 })
             },
-            Details::NotifyFinalized(ref finalized) => {
+            Payload::NotifyFinalized(ref finalized) => {
                 Some(Block {
                     hash: finalized.hash,
                     height: finalized.height.parse().ok()?
@@ -151,3 +170,32 @@ impl Details {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn message_v1() {
+        let json = r#"{"msg":"notify.finalized","level":"INFO","ts":"2021-01-13T12:38:25.410794650+01:00","best":"0x031c3521ca2f9c673812d692fc330b9a18e18a2781e3f9976992f861fd3ea0cb","height":"50"}"#;
+        assert!(
+            matches!(
+                serde_json::from_str::<NodeMessage>(json).unwrap(),
+                NodeMessage::V1 { .. },
+            ),
+            "message did not match variant V1",
+        );
+    }
+
+    #[test]
+    fn message_v2() {
+        let json = r#"{"id":1,"ts":"2021-01-13T12:22:20.053527101+01:00","payload":{"best":"0xcc41708573f2acaded9dd75e07dac2d4163d136ca35b3061c558d7a35a09dd8d","height":"209","msg":"notify.finalized"}}"#;
+        assert!(
+            matches!(
+                serde_json::from_str::<NodeMessage>(json).unwrap(),
+                NodeMessage::V2 { .. },
+            ),
+            "message did not match variant V2",
+        );
+    }
+}