katcp/messages/
common.rs

1//! Common message types and traits
2use std::{
3    fmt::Display,
4    net::{IpAddr, SocketAddr},
5};
6
7use chrono::{DateTime, TimeZone, Utc};
8use katcp_derive::KatcpDiscrete;
9
10use crate::{
11    protocol::{KatcpError, Message, MessageResult},
12    utils::{escape, unescape},
13};
14
15/// The trait that specific katcp messages should implement
16pub trait KatcpMessage: TryFrom<Message> {
17    fn to_message(&self, id: Option<u32>) -> MessageResult;
18}
19
20/// Serializes the implemented type into an argument string
21/// Implemented for all fundamental katcp types as well as any user-defined types
22pub trait ToKatcpArgument {
23    /// Create a katcp message argument (String) from a self
24    fn to_argument(&self) -> String;
25}
26
27/// Deserializes an argument string into the implemented type
28/// Implemented for all fundamental katcp types as well as any user-defined types
29pub trait FromKatcpArgument
30where
31    Self: Sized,
32{
33    type Err; // Not Error as to not clash with Self being an enum with an `Error` variant
34    /// Create a self from a katcp message argument (String), potentially erroring
35    fn from_argument(s: impl AsRef<str>) -> Result<Self, Self::Err>;
36}
37
38/// A trait for serializing more complex types that return the full argument vector
39pub trait ToKatcpArguments {
40    fn to_arguments(&self) -> Vec<String>;
41}
42
43/// A trait for deserializing more complex types that consume an iterator of arguments
44pub trait FromKatcpArguments
45where
46    Self: Sized,
47{
48    type Err;
49    fn from_arguments(strings: &mut impl Iterator<Item = String>) -> Result<Self, Self::Err>;
50}
51
52/// Marker trait for implementers of both [`ToKatcpArgument`] and [`FromKatcpArgument`]
53pub trait KatcpArgument: ToKatcpArgument + FromKatcpArgument {}
54/// Marker trait for implementers of both [`ToKatcpArguments`] and [`FromKatcpArguments`]
55pub trait KatcpArguments: ToKatcpArguments + FromKatcpArguments {}
56impl<T> KatcpArgument for T where T: ToKatcpArgument + FromKatcpArgument {}
57impl<T> KatcpArguments for T where T: ToKatcpArguments + FromKatcpArguments {}
58
59/// Type alias for DateTime<Utc> from chrono
60pub type KatcpTimestamp = DateTime<Utc>;
61
62// ---- Implementations for the "core" KatcpTypes
63
64// str
65impl ToKatcpArgument for str {
66    fn to_argument(&self) -> String {
67        escape(self)
68    }
69}
70
71impl ToKatcpArgument for String {
72    fn to_argument(&self) -> String {
73        escape(self)
74    }
75}
76
77impl FromKatcpArgument for String {
78    type Err = KatcpError;
79
80    fn from_argument(s: impl AsRef<str>) -> Result<Self, Self::Err> {
81        Ok(unescape(s.as_ref()))
82    }
83}
84
85// KatcpTimestamp
86impl ToKatcpArgument for KatcpTimestamp {
87    fn to_argument(&self) -> String {
88        let secs = self.timestamp() as f64;
89        let nano = self.timestamp_subsec_nanos();
90        let frac = nano as f64 / 1e9;
91        format!("{}", secs + frac)
92    }
93}
94
95impl FromKatcpArgument for KatcpTimestamp {
96    type Err = KatcpError;
97
98    fn from_argument(s: impl AsRef<str>) -> Result<Self, Self::Err> {
99        let fractional: f64 = s.as_ref().parse().map_err(|_| KatcpError::BadArgument)?;
100        let secs = fractional as i64;
101        let nanos = (fractional.fract() * 1e9) as u32;
102        Ok(Utc.timestamp(secs, nanos))
103    }
104}
105
106// Option
107impl<T> ToKatcpArgument for Option<T>
108where
109    T: ToKatcpArgument,
110{
111    fn to_argument(&self) -> String {
112        match self {
113            Some(v) => v.to_argument(),
114            None => r"\@".to_owned(),
115        }
116    }
117}
118
119impl<E, T> FromKatcpArgument for Option<T>
120where
121    T: FromKatcpArgument<Err = E>,
122{
123    type Err = E;
124
125    fn from_argument(s: impl AsRef<str>) -> Result<Self, Self::Err> {
126        match s.as_ref() {
127            r"\@" => Ok(None),
128            _ => Ok(Some(T::from_argument(s)?)),
129        }
130    }
131}
132
133// Return Code
134#[derive(KatcpDiscrete, Debug, PartialEq, Eq, Copy, Clone)]
135/// Return codes that form the first parameter of replys
136pub enum RetCode {
137    /// Request successfully processed. Further arguments are request-specific
138    Ok,
139    /// Request malformed. Second argument is a human-readable description of the error
140    Invalid,
141    /// Valid request that could not be processed. Second argument is a human-readable description of the error.
142    Fail,
143}
144
145impl ToKatcpArgument for u32 {
146    fn to_argument(&self) -> String {
147        self.to_string()
148    }
149}
150
151impl FromKatcpArgument for u32 {
152    type Err = KatcpError;
153
154    fn from_argument(s: impl AsRef<str>) -> Result<Self, Self::Err> {
155        s.as_ref().parse().map_err(|_| KatcpError::BadArgument)
156    }
157}
158
159impl ToKatcpArgument for i32 {
160    fn to_argument(&self) -> String {
161        self.to_string()
162    }
163}
164
165impl FromKatcpArgument for i32 {
166    type Err = KatcpError;
167
168    fn from_argument(s: impl AsRef<str>) -> Result<Self, Self::Err> {
169        s.as_ref().parse().map_err(|_| KatcpError::BadArgument)
170    }
171}
172
173impl ToKatcpArgument for bool {
174    fn to_argument(&self) -> String {
175        (if *self { "1" } else { "0" }).to_owned()
176    }
177}
178
179impl FromKatcpArgument for bool {
180    type Err = KatcpError;
181
182    fn from_argument(s: impl AsRef<str>) -> Result<Self, Self::Err> {
183        match s.as_ref() {
184            "1" => Ok(true),
185            "0" => Ok(false),
186            _ => Err(KatcpError::BadArgument),
187        }
188    }
189}
190
191impl ToKatcpArgument for f32 {
192    fn to_argument(&self) -> String {
193        format!("{}", self)
194    }
195}
196
197impl FromKatcpArgument for f32 {
198    type Err = KatcpError;
199
200    fn from_argument(s: impl AsRef<str>) -> Result<Self, Self::Err> {
201        s.as_ref().parse().map_err(|_| KatcpError::BadArgument)
202    }
203}
204
205/// Katcp addresses optionally have a port, so we need a sum type for the two native rust
206/// types [`IpAddr`] and [`SocketAddr`], depending on whether we have a port
207#[derive(Debug, PartialEq, Eq, Copy, Clone)]
208pub enum KatcpAddress {
209    Ip(IpAddr),
210    Socket(SocketAddr),
211}
212
213impl ToKatcpArgument for KatcpAddress {
214    fn to_argument(&self) -> String {
215        match self {
216            KatcpAddress::Ip(addr) => match addr {
217                IpAddr::V4(addr) => addr.to_string(),
218                IpAddr::V6(addr) => format!("[{}]", addr),
219            },
220            KatcpAddress::Socket(addr) => match addr {
221                SocketAddr::V4(addr) => addr.to_string(),
222                SocketAddr::V6(addr) => addr.to_string(),
223            },
224        }
225    }
226}
227
228impl FromKatcpArgument for KatcpAddress {
229    type Err = KatcpError;
230
231    fn from_argument(s: impl AsRef<str>) -> Result<Self, Self::Err> {
232        if s.as_ref().is_empty() {
233            Err(KatcpError::BadArgument)
234        } else if let Ok(addr) = s.as_ref().parse() {
235            Ok(Self::Socket(addr))
236        } else if let Ok(addr) = s.as_ref().parse() {
237            Ok(Self::Ip(addr))
238        } else if s.as_ref().starts_with('[') && s.as_ref().ends_with(']') {
239            let slice = &s.as_ref()[1..s.as_ref().len() - 1];
240            if let Ok(addr) = slice.parse() {
241                Ok(Self::Ip(addr))
242            } else {
243                Err(KatcpError::BadArgument)
244            }
245        } else {
246            Err(KatcpError::BadArgument)
247        }
248    }
249}
250
251/// Convienence method for round-trip testing
252pub fn roundtrip_test<T, E>(message: T)
253where
254    E: std::fmt::Debug,
255    T: KatcpMessage + PartialEq + std::fmt::Debug + TryFrom<Message, Error = E>,
256{
257    let raw = message.to_message(None).unwrap();
258    let s = raw.to_string();
259    // Print the middle, we're using this in tests, so we'll only see it on fails
260    println!("Katcp Payload:\n{}", s);
261    let raw_test: Message = (s.as_str()).try_into().unwrap();
262    let message_test = raw_test.try_into().unwrap();
263    assert_eq!(message, message_test)
264}
265
266#[derive(KatcpDiscrete, Debug, PartialEq, Eq, Copy, Clone)]
267/// The datatypes that KATCP supports
268pub enum ArgumentType {
269    /// Represented by i32
270    Integer,
271    /// Represented by f32
272    Float,
273    Boolean,
274    /// Represented by [`KatcpTimestamp`]
275    Timestamp,
276    /// Represented by an enum
277    Discrete,
278    /// Represented by KatcpAddress
279    Address,
280    /// Will always be escaped and unescaped during serde
281    String,
282}
283
284#[derive(Debug, PartialEq, Clone)]
285/// The sum type of a vector of one of the primitive [`ArgumentType`]s
286pub enum ArgumentVec {
287    Integer(Vec<i32>),
288    Float(Vec<f32>),
289    Boolean(Vec<bool>),
290    Timestamp(Vec<KatcpTimestamp>),
291    String(Vec<String>),
292    Discrete(Vec<String>),
293    Address(Vec<KatcpAddress>),
294}
295
296impl Display for ArgumentVec {
297    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
298        let s = match self {
299            ArgumentVec::Integer(_) => "integer",
300            ArgumentVec::Float(_) => "float",
301            ArgumentVec::Boolean(_) => "boolean",
302            ArgumentVec::Timestamp(_) => "timestamp",
303            ArgumentVec::String(_) => "string",
304            ArgumentVec::Discrete(_) => "discrete",
305            ArgumentVec::Address(_) => "address",
306        }
307        .to_owned();
308
309        write! {f,"{}",s}
310    }
311}
312
313impl ToKatcpArguments for ArgumentVec {
314    fn to_arguments(&self) -> Vec<String> {
315        match self {
316            Self::Integer(v) => v.iter().map(|e| e.to_argument()).collect(),
317            Self::Float(v) => v.iter().map(|e| e.to_argument()).collect(),
318            Self::Boolean(v) => v.iter().map(|e| e.to_argument()).collect(),
319            Self::Timestamp(v) => v.iter().map(|e| e.to_argument()).collect(),
320            Self::String(v) => v.iter().map(|e| e.to_argument()).collect(),
321            Self::Discrete(v) => v.iter().map(|e| e.to_argument()).collect(),
322            Self::Address(v) => v.iter().map(|e| e.to_argument()).collect(),
323        }
324    }
325}
326
327pub(crate) fn from_argument_vec(
328    ty: &ArgumentType,
329    strings: &mut impl Iterator<Item = String>,
330) -> Result<ArgumentVec, KatcpError> {
331    Ok(match ty {
332        ArgumentType::Boolean => ArgumentVec::Boolean(
333            strings
334                .map(bool::from_argument)
335                .collect::<Result<Vec<_>, _>>()?,
336        ),
337        ArgumentType::Integer => ArgumentVec::Integer(
338            strings
339                .map(i32::from_argument)
340                .collect::<Result<Vec<_>, _>>()?,
341        ),
342        ArgumentType::Float => ArgumentVec::Float(
343            strings
344                .map(f32::from_argument)
345                .collect::<Result<Vec<_>, _>>()?,
346        ),
347        ArgumentType::Timestamp => ArgumentVec::Timestamp(
348            strings
349                .map(KatcpTimestamp::from_argument)
350                .collect::<Result<Vec<_>, _>>()?,
351        ),
352        ArgumentType::Discrete => ArgumentVec::Discrete(
353            strings
354                .map(String::from_argument)
355                .collect::<Result<Vec<_>, _>>()?,
356        ),
357        ArgumentType::Address => ArgumentVec::Address(
358            strings
359                .map(KatcpAddress::from_argument)
360                .collect::<Result<Vec<_>, _>>()?,
361        ),
362        ArgumentType::String => ArgumentVec::String(
363            strings
364                .map(String::from_argument)
365                .collect::<Result<Vec<_>, _>>()?,
366        ),
367    })
368}
369
370#[cfg(test)]
371mod test_arguments {
372    use super::*;
373
374    #[test]
375    fn test_string() {
376        let s = "This is a message with spaces\n";
377        assert_eq!(s, String::from_argument(s.to_argument()).unwrap());
378    }
379
380    #[test]
381    fn test_timestamp() {
382        let ts = Utc.timestamp(42069, 42069000);
383        assert_eq!(ts, KatcpTimestamp::from_argument(ts.to_argument()).unwrap());
384    }
385
386    #[test]
387    fn test_option() {
388        let s = Some("\tFoo a bar\n".to_owned());
389        assert_eq!(s, Option::<String>::from_argument(s.to_argument()).unwrap())
390    }
391
392    #[test]
393    fn test_ret_code() {
394        let code = RetCode::Invalid;
395        assert_eq!(code, RetCode::from_argument(code.to_argument()).unwrap())
396    }
397
398    #[test]
399    fn test_int() {
400        let pos_int = 12345;
401        let neg_int = -12345;
402        assert_eq!(pos_int, u32::from_argument(pos_int.to_argument()).unwrap());
403        assert_eq!(neg_int, i32::from_argument(neg_int.to_argument()).unwrap());
404    }
405
406    #[test]
407    fn test_bool() {
408        let a = true;
409        let b = false;
410        assert_eq!(a, bool::from_argument(a.to_argument()).unwrap());
411        assert_eq!(b, bool::from_argument(b.to_argument()).unwrap());
412    }
413
414    #[test]
415    fn test_float() {
416        let a = -1.234e-05;
417        let b = 1.7;
418        let c = 100.0;
419        assert_eq!(a, f32::from_argument(a.to_argument()).unwrap());
420        assert_eq!(b, f32::from_argument(b.to_argument()).unwrap());
421        assert_eq!(c, f32::from_argument(c.to_argument()).unwrap());
422    }
423
424    #[test]
425    fn test_addr() {
426        let v4_socket = "192.168.1.1:4000";
427        let v4_ip = "127.0.0.1";
428        let v6_socket = "[2001:db8:85a3::8a2e:370:7334]:4000";
429        let v6_ip = "[::1]";
430        assert_eq!(
431            v4_socket,
432            KatcpAddress::from_argument(v4_socket)
433                .unwrap()
434                .to_argument()
435        );
436        assert_eq!(
437            v6_socket,
438            KatcpAddress::from_argument(v6_socket)
439                .unwrap()
440                .to_argument()
441        );
442        assert_eq!(
443            v4_ip,
444            KatcpAddress::from_argument(v4_ip).unwrap().to_argument()
445        );
446        assert_eq!(
447            v6_ip,
448            KatcpAddress::from_argument(v6_ip).unwrap().to_argument()
449        );
450    }
451}