socket_flow/
message.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
use crate::error::Error;
use crate::frame::{Frame, OpCode};

#[derive(Debug, Clone, PartialEq)]
pub enum Message {
    Text(String),
    Binary(Vec<u8>),
}

impl Message {
    // Converts a Frame into a Message variant
    pub fn from_frame(frame: Frame) -> Result<Self, Error> {
        match frame.opcode {
            OpCode::Text => Ok(Message::Text(String::from_utf8(frame.payload)?)),
            OpCode::Binary => Ok(Message::Binary(frame.payload)),
            _ => Err(Error::InvalidOpcode),
        }
    }

    // Function to get the payload as binary (Vec<u8>)
    pub fn as_binary(&self) -> Vec<u8> {
        match self {
            Message::Text(text) => text.as_bytes().to_vec(),
            Message::Binary(data) => data.clone(),
        }
    }

    // Function to get the payload as a String
    pub fn as_text(&self) -> Result<String, Error> {
        match self {
            Message::Text(text) => Ok(text.clone()),
            Message::Binary(data) => Ok(String::from_utf8(data.clone())?),
        }
    }

    // Function to convert Message to a Frame
    pub fn to_frame(self, final_fragment: bool) -> Frame {
        let opcode = match self {
            Message::Text(_) => OpCode::Text,
            Message::Binary(_) => OpCode::Binary,
        };

        let payload = match self {
            Message::Text(text) => text.into_bytes(),
            Message::Binary(data) => data,
        };

        Frame {
            final_fragment,
            opcode,
            payload,
        }
    }
}