snops_common/state/
node_state.rs

1use std::net::SocketAddr;
2
3use indexmap::IndexMap;
4
5use super::{AgentId, HeightRequest, InternedId, NodeKey};
6use crate::format::{DataFormat, DataFormatReader, DataHeaderOf, PackedUint};
7
8#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
9pub struct NodeState {
10    pub node_key: NodeKey,
11    pub private_key: KeyState,
12    /// Increment the usize whenever the request is updated.
13    pub height: (usize, HeightRequest),
14
15    pub online: bool,
16    pub peers: Vec<AgentPeer>,
17    pub validators: Vec<AgentPeer>,
18    pub env: IndexMap<String, String>,
19    pub binary: Option<InternedId>,
20}
21
22#[derive(Debug, Clone)]
23pub struct NodeStateFormatHeader {
24    version: u8,
25    node_key: DataHeaderOf<NodeKey>,
26    key_state: DataHeaderOf<KeyState>,
27    height: DataHeaderOf<HeightRequest>,
28    peer: DataHeaderOf<AgentPeer>,
29}
30
31impl DataFormat for NodeStateFormatHeader {
32    type Header = u8;
33    const LATEST_HEADER: Self::Header = 1;
34
35    fn write_data<W: std::io::prelude::Write>(
36        &self,
37        writer: &mut W,
38    ) -> Result<usize, crate::format::DataWriteError> {
39        let mut written = 0;
40        written += self.version.write_data(writer)?;
41        written += self.node_key.write_data(writer)?;
42        written += self.key_state.write_data(writer)?;
43        written += self.height.write_data(writer)?;
44        written += self.peer.write_data(writer)?;
45        Ok(written)
46    }
47
48    fn read_data<R: std::io::prelude::Read>(
49        reader: &mut R,
50        header: &Self::Header,
51    ) -> Result<Self, crate::format::DataReadError> {
52        if *header != Self::LATEST_HEADER {
53            return Err(crate::format::DataReadError::unsupported(
54                "NodeStateFormatHeader",
55                Self::LATEST_HEADER,
56                *header,
57            ));
58        }
59
60        Ok(NodeStateFormatHeader {
61            version: reader.read_data(&())?,
62            node_key: reader.read_data(&((), ()))?,
63            key_state: reader.read_data(&())?,
64            height: reader.read_data(&((), ()))?,
65            peer: reader.read_data(&())?,
66        })
67    }
68}
69
70impl DataFormat for NodeState {
71    type Header = NodeStateFormatHeader;
72    const LATEST_HEADER: Self::Header = NodeStateFormatHeader {
73        version: 2,
74        node_key: NodeKey::LATEST_HEADER,
75        key_state: KeyState::LATEST_HEADER,
76        height: HeightRequest::LATEST_HEADER,
77        peer: AgentPeer::LATEST_HEADER,
78    };
79
80    fn write_data<W: std::io::prelude::Write>(
81        &self,
82        writer: &mut W,
83    ) -> Result<usize, crate::format::DataWriteError> {
84        let mut written = 0;
85        written += self.node_key.write_data(writer)?;
86        written += self.private_key.write_data(writer)?;
87        written += PackedUint::from(self.height.0).write_data(writer)?;
88        written += self.height.1.write_data(writer)?;
89        written += self.online.write_data(writer)?;
90        written += self.peers.write_data(writer)?;
91        written += self.validators.write_data(writer)?;
92        written += self.env.write_data(writer)?;
93        written += self.binary.write_data(writer)?;
94        Ok(written)
95    }
96
97    fn read_data<R: std::io::prelude::Read>(
98        reader: &mut R,
99        header: &Self::Header,
100    ) -> Result<Self, crate::format::DataReadError> {
101        if header.version == 0 || header.version > Self::LATEST_HEADER.version {
102            return Err(crate::format::DataReadError::unsupported(
103                "NodeState",
104                format!("1 or {}", Self::LATEST_HEADER.version),
105                header.version,
106            ));
107        }
108
109        let node_key = reader.read_data(&header.node_key)?;
110        let private_key = reader.read_data(&header.key_state)?;
111        let height_inc = PackedUint::read_data(reader, &())?;
112        let height_req = reader.read_data(&header.height)?;
113        let online = reader.read_data(&())?;
114        let peers = reader.read_data(&header.peer)?;
115        let validators = reader.read_data(&header.peer)?;
116        let env = reader.read_data(&((), ()))?;
117        let binary = if header.version > 1 {
118            reader.read_data(&())?
119        } else {
120            None
121        };
122
123        Ok(NodeState {
124            node_key,
125            private_key,
126            height: (height_inc.into(), height_req),
127            online,
128            peers,
129            validators,
130            env,
131            binary,
132        })
133    }
134}
135
136/// A representation of which key to use for the agent.
137#[derive(Default, Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
138pub enum KeyState {
139    /// No private key provided
140    #[default]
141    None,
142    /// A private key is provided by the agent
143    Local,
144    /// A literal private key
145    Literal(String),
146    // TODO: generated?/new
147}
148
149impl From<Option<String>> for KeyState {
150    fn from(s: Option<String>) -> Self {
151        match s {
152            Some(s) => Self::Literal(s),
153            None => Self::None,
154        }
155    }
156}
157
158impl KeyState {
159    pub fn try_string(&self) -> Option<String> {
160        match self {
161            Self::Literal(s) => Some(s.to_owned()),
162            _ => None,
163        }
164    }
165
166    pub fn is_none(&self) -> bool {
167        matches!(self, KeyState::None)
168    }
169}
170
171#[derive(
172    Debug, Clone, Copy, serde::Serialize, serde::Deserialize, PartialEq, Eq, Hash, PartialOrd, Ord,
173)]
174pub enum AgentPeer {
175    Internal(AgentId, u16),
176    External(SocketAddr),
177}
178
179impl AgentPeer {
180    /// Get the port from the peer
181    pub fn port(&self) -> u16 {
182        match self {
183            Self::Internal(_, port) => *port,
184            Self::External(addr) => addr.port(),
185        }
186    }
187
188    /// Return a new peer with the given port.
189    pub fn with_port(&self, port: u16) -> Self {
190        match self {
191            Self::Internal(ip, _) => Self::Internal(*ip, port),
192            Self::External(addr) => Self::External(SocketAddr::new(addr.ip(), port)),
193        }
194    }
195}
196
197impl DataFormat for KeyState {
198    type Header = u8;
199    const LATEST_HEADER: Self::Header = 1;
200
201    fn write_data<W: std::io::prelude::Write>(
202        &self,
203        writer: &mut W,
204    ) -> Result<usize, crate::format::DataWriteError> {
205        match self {
206            Self::None => 0u8.write_data(writer),
207            Self::Local => 1u8.write_data(writer),
208            Self::Literal(s) => Ok(2u8.write_data(writer)? + s.write_data(writer)?),
209        }
210    }
211
212    fn read_data<R: std::io::prelude::Read>(
213        reader: &mut R,
214        header: &Self::Header,
215    ) -> Result<Self, crate::format::DataReadError> {
216        if *header != Self::LATEST_HEADER {
217            return Err(crate::format::DataReadError::unsupported(
218                "KeyState",
219                Self::LATEST_HEADER,
220                *header,
221            ));
222        }
223
224        match reader.read_data(&())? {
225            0u8 => Ok(Self::None),
226            1u8 => Ok(Self::Local),
227            2u8 => Ok(Self::Literal(reader.read_data(&())?)),
228            n => Err(crate::format::DataReadError::Custom(format!(
229                "Invalid KeyState discriminant: {n}",
230            ))),
231        }
232    }
233}
234
235impl DataFormat for AgentPeer {
236    type Header = u8;
237    const LATEST_HEADER: Self::Header = 1;
238
239    fn write_data<W: std::io::prelude::Write>(
240        &self,
241        writer: &mut W,
242    ) -> Result<usize, crate::format::DataWriteError> {
243        match self {
244            Self::Internal(id, port) => {
245                Ok(0u8.write_data(writer)? + id.write_data(writer)? + port.write_data(writer)?)
246            }
247            Self::External(addr) => Ok(1u8.write_data(writer)? + addr.write_data(writer)?),
248        }
249    }
250
251    fn read_data<R: std::io::prelude::Read>(
252        reader: &mut R,
253        header: &Self::Header,
254    ) -> Result<Self, crate::format::DataReadError> {
255        if *header != Self::LATEST_HEADER {
256            return Err(crate::format::DataReadError::unsupported(
257                "AgentPeer",
258                Self::LATEST_HEADER,
259                *header,
260            ));
261        }
262
263        match reader.read_data(&())? {
264            0u8 => Ok(Self::Internal(
265                reader.read_data(&())?,
266                reader.read_data(&())?,
267            )),
268            1u8 => Ok(Self::External(reader.read_data(&())?)),
269            n => Err(crate::format::DataReadError::Custom(format!(
270                "Invalid AgentPeer discriminant: {n}",
271            ))),
272        }
273    }
274}
275
276#[cfg(test)]
277mod tests {
278
279    use crate::{
280        format::{read_dataformat, write_dataformat, DataFormat},
281        prelude::{HeightRequest, KeyState, NodeState, NodeStateFormatHeader},
282    };
283
284    macro_rules! case {
285        ($name:ident, $ty:ty, $a:expr, $b:expr) => {
286            #[test]
287            fn $name() -> Result<(), Box<dyn std::error::Error>> {
288                let mut data = Vec::new();
289                write_dataformat(&mut data, &$a)?;
290                assert_eq!(data, $b);
291
292                let mut reader = &data[..];
293                let read_value = read_dataformat::<_, $ty>(&mut reader)?;
294
295                // write the data again because not every type implements PartialEq
296                let mut data2 = Vec::new();
297                write_dataformat(&mut data2, &read_value)?;
298                assert_eq!(data, data2);
299                Ok(())
300            }
301        };
302    }
303
304    case!(
305        node_state,
306        NodeState,
307        NodeState {
308            node_key: "client/foo".parse()?,
309            private_key: KeyState::None,
310            height: (0, HeightRequest::Top),
311            online: true,
312            peers: vec![],
313            validators: vec![],
314            env: Default::default(),
315            binary: None,
316        },
317        [
318            NodeStateFormatHeader::LATEST_HEADER.to_byte_vec()?,
319            NodeState::LATEST_HEADER.to_byte_vec()?,
320            NodeState {
321                node_key: "client/foo".parse()?,
322                private_key: KeyState::None,
323                height: (0, HeightRequest::Top),
324                online: true,
325                peers: vec![],
326                validators: vec![],
327                env: Default::default(),
328                binary: None,
329            }
330            .to_byte_vec()?,
331        ]
332        .concat()
333    );
334}
snops_common/state/node_state.rs

snops_common/state/
node_state.rs
