ostrich_core/
lib.rs

1use std::io;
2// use std::fmt;
3use std::ops::Range;
4use std::convert::TryFrom;
5
6use num_derive::FromPrimitive;    
7use num_traits::FromPrimitive;
8
9/// Ostrich packet size, 1024 Bytes (1K)
10pub const PCK_SIZE: usize = 1024;
11
12/*  Ostrich packet format:
13 *  Some fields are empty for some messages.
14 *  For example, when sending an Error command
15 *  fields sender and receiver are empty.
16 *
17 * 1   B : Command code (0)
18 * 16  B : Sender name or empty (1-15) NOTE: First byte for text length
19 * 16  B : Receiver or empty (16-31) NOTE: First byte for text length
20 * 2   B : Text length in bytes (32-33)
21 * 991 B : Text or empty (34-1023)
22 */
23pub const CMD_BYTES: Range<usize> = (0..0);
24pub const SENDER_LEN: usize = 1;
25pub const SENDER_BYTES: Range<usize> = (2..15);
26pub const RECV_LEN: usize = 16;
27pub const RECV_BYTES: Range<usize> = (17..31);
28pub const TXT_LEN: Range<usize> = (32..33);
29pub const TXT_BYTES: Range<usize> = (34..1023);
30
31#[allow(dead_code)]
32#[derive(Debug, Clone, Copy, PartialEq, Eq, FromPrimitive)]
33#[repr(u8)]
34pub enum CommandCode {
35    Ok = 0,
36    Err = 1,
37    Get = 2,
38    Msg = 3,
39}
40
41#[derive(Debug, Clone, PartialEq)]
42pub enum Command {
43    Ok,
44    Err(String),
45    Get,
46    Msg(String, String, String),
47}
48
49pub struct RawMessage;
50
51impl RawMessage {
52
53    pub fn from_raw(raw: &[u8]) -> Result<Command, io::Error> {
54        // Check the first byte for command code
55        match FromPrimitive::from_u8(raw[0]) {
56            Some(CommandCode::Ok) => Ok(Command::Ok),
57            Some(CommandCode::Err) => {
58                // Get error message length
59                let mut range = [0u8;2];
60                range[0] = raw[TXT_LEN.start];
61                range[1] = raw[TXT_LEN.end];
62                let n: usize = u16::from_ne_bytes(range) as usize;
63
64                // Code the error into an utf-8 string
65                let error = String::from_utf8_lossy(&raw[TXT_BYTES][..n]);
66                Ok(Command::Err(error.to_string()))
67            },
68            Some(CommandCode::Get) => Ok(Command::Get),
69            Some(CommandCode::Msg) => {
70                let n = raw[SENDER_LEN] as usize;
71                let sender = String::from_utf8_lossy(&raw[SENDER_BYTES][..n]);
72                let n = raw[RECV_LEN] as usize;
73                let recv = String::from_utf8_lossy(&raw[RECV_BYTES][..n]);
74                
75                // Get the length of the message text
76                let mut range = [0u8;2];
77                range[0] = raw[TXT_LEN.start];
78                range[1] = raw[TXT_LEN.end];
79                let n: usize = u16::from_ne_bytes(range) as usize;
80
81                // Convert txt to string
82                let text = String::from_utf8_lossy(&raw[TXT_BYTES][..n]);
83
84                Ok(Command::Msg(sender.to_string(), recv.to_string(), text.to_string()))
85            },
86            None => Err(io::Error::new(io::ErrorKind::InvalidData, 
87                                       format!("Incorrect command byte: {}", raw[0]))),
88        }
89    }
90
91    fn put(buffer: &mut [u8], 
92           content: &[u8], 
93           range: Range<usize>) -> Result<(), io::Error> {
94        
95        // Check the range is inside the buffer's bounds
96        if range.end > buffer.len() {
97            let err = format!("Rage out of bounds: range end {}, buffer len {}", 
98                              range.end, buffer.len());
99            return Err(io::Error::new(io::ErrorKind::InvalidInput, err));
100        }
101        
102        content.iter()
103            .enumerate()
104            .skip_while(|(i, _)| *i > range.end) // Check if the content size is larger than the range
105            .for_each(|(i, x)| buffer[range.start+i] = *x);
106
107        Ok(())
108    }
109
110    pub fn to_raw(command: &Command) -> Result<[u8; 1024], io::Error> {
111        // Init buffer
112        // let mut buffer = BytesMut::with_capacity(PCK_SIZE);
113        let mut buffer = [0u8; PCK_SIZE];
114
115        // Set command code
116        match command {
117            Command::Ok => buffer[0] = CommandCode::Ok as u8,
118            Command::Err(err) => {
119                // Set command code
120                buffer[0] = CommandCode::Err as u8;
121
122                // Set the error message length bytes
123                let n = match u16::try_from(err.len()) {
124                    Ok(n) => n.to_ne_bytes(),
125                    Err(_) => return Err(io::Error::new(io::ErrorKind::InvalidInput, 
126                                                  "Error message length exceded"))
127                };
128
129                RawMessage::put(&mut buffer, &n, TXT_LEN)?;
130                // Append the error's bytes to the buffer's text section
131                let err = err.as_bytes();
132                RawMessage::put(&mut buffer, err, TXT_BYTES)?;
133            },
134            Command::Get => buffer[0] = CommandCode::Get as u8,
135            Command::Msg(s,r,t) => {
136                // Append MSG code
137                buffer[0] = CommandCode::Msg as u8;
138                // Add sender name
139                let s = s.as_bytes();
140                buffer[SENDER_LEN] = s.len() as u8;
141                RawMessage::put(&mut buffer, s, SENDER_BYTES)?;
142                // Add receiver name
143                let r = r.as_bytes();
144                buffer[RECV_LEN] = r.len() as u8;
145                RawMessage::put(&mut buffer, r, RECV_BYTES)?;
146
147                // Set the txt message length bytes
148                let t = t.as_bytes();
149                let n = match u16::try_from(t.len()) {
150                    Ok(n) => n.to_ne_bytes(),
151                    Err(_) => return Err(io::Error::new(io::ErrorKind::InvalidInput, 
152                                                  "Error message length exceded"))
153                };
154                RawMessage::put(&mut buffer, &n, TXT_LEN)?;
155                // Add the messgase's body
156                RawMessage::put(&mut buffer, t, TXT_BYTES)?;
157            },
158        }
159
160        Ok(buffer)
161    }
162}
163
164#[test]
165fn test_ok() {
166    // Test ok command
167    let command = Command::Ok; 
168    let mesg = RawMessage::to_raw(&command).unwrap();
169    let recovered = RawMessage::from_raw(&mesg).unwrap();
170    assert_eq!(mesg[0], 0);
171    assert_eq!(command, recovered);
172
173}
174
175#[test]
176fn test_get() {
177    // Test ok command
178    let command = Command::Get; 
179    let mesg = RawMessage::to_raw(&command).unwrap();
180    let recovered = RawMessage::from_raw(&mesg).unwrap();
181    assert_eq!(mesg[0], 2);
182    assert_eq!(command, recovered);
183
184}
185
186#[test]
187fn test_err() {
188    // Test error command
189    let command = Command::Err("Some fatal error".to_string());
190    let mesg = RawMessage::to_raw(&command).unwrap();
191    let recovered = RawMessage::from_raw(&mesg).unwrap();
192    assert_eq!(mesg[0], 1);
193    assert_eq!(command, recovered);
194}
195
196#[test]
197fn test_msg() {
198    // Test error command
199    let command = Command::Msg("sender".to_string(),
200                               "receiver".to_string(),
201                               "The super secret message".to_string());
202
203    let mesg = RawMessage::to_raw(&command).unwrap();
204    let recovered = RawMessage::from_raw(&mesg).unwrap();
205    assert_eq!(mesg[0], 3);
206    assert_eq!(command, recovered);
207}