feophantlib/codec/
network_frame.rs

1use bytes::Bytes;
2use bytes::{BufMut, BytesMut};
3use std::convert::TryFrom;
4use std::num::TryFromIntError;
5use thiserror::Error;
6
7use crate::constants::{PgErrorCodes, PgErrorLevels};
8use crate::engine::objects::SqlTuple;
9
10#[derive(Clone, Debug)]
11pub struct NetworkFrame {
12    pub message_type: u8,
13    pub payload: Bytes,
14}
15
16impl NetworkFrame {
17    pub fn new(message_type: u8, payload: Bytes) -> NetworkFrame {
18        NetworkFrame {
19            message_type,
20            payload,
21        }
22    }
23    pub fn authentication_ok() -> NetworkFrame {
24        NetworkFrame::new(b'R', Bytes::from_static(b"\0\0\0\0"))
25    }
26
27    pub fn command_complete(command_tag: String) -> NetworkFrame {
28        let mut buffer = BytesMut::new();
29
30        buffer.put(command_tag.as_bytes());
31        buffer.put_u8(b'\0');
32
33        NetworkFrame::new(b'C', buffer.freeze())
34    }
35
36    pub fn data_rows(rows: Vec<SqlTuple>) -> Result<Vec<NetworkFrame>, NetworkFrameError> {
37        let mut frames = vec![];
38
39        for row in rows {
40            let mut buffer = BytesMut::new();
41
42            let column_count = u16::try_from(row.0.len())?;
43            buffer.put_u16(column_count);
44
45            for field in row.0.into_iter() {
46                match field {
47                    Some(f) => {
48                        let f_str = format!("{}", f);
49                        let f_bytes = f_str.as_bytes();
50                        let f_len = i32::try_from(f_bytes.len())?;
51                        buffer.put_i32(f_len);
52                        buffer.put(f_bytes);
53                    }
54                    None => {
55                        buffer.put_i32(-1);
56                    }
57                }
58            }
59
60            frames.push(NetworkFrame::new(b'D', buffer.freeze()));
61        }
62
63        Ok(frames)
64    }
65
66    //Note this claims that the server is ALWAYS ready, even if its not
67    pub fn ready_for_query() -> NetworkFrame {
68        NetworkFrame::new(b'Z', Bytes::from_static(b"I"))
69    }
70
71    pub fn row_description(column_names: Vec<String>) -> Result<NetworkFrame, NetworkFrameError> {
72        let mut buffer = BytesMut::new();
73
74        let field_count = u16::try_from(column_names.len())?;
75        buffer.put_u16(field_count);
76
77        for column in column_names {
78            buffer.put(column.as_bytes());
79            buffer.put_u8(b'\0');
80
81            //The next following fields are going to be dummied out unless testing shows I need them.
82            //https://www.postgresql.org/docs/current/protocol-message-formats.html
83            buffer.put_u32(0); //Table OID
84            buffer.put_u16(0); //Table Column
85            buffer.put_u32(0); //Type OID
86            buffer.put_i16(0); //Type length
87            buffer.put_i32(0); //Type specifier
88            buffer.put_i16(0); //Format code, we're doing text for everything
89        }
90
91        Ok(NetworkFrame::new(b'T', buffer.freeze()))
92    }
93
94    //Valid severities can be found here: https://www.postgresql.org/docs/current/protocol-error-fields.html
95    //Valid error codes can be found here: https://www.postgresql.org/docs/current/errcodes-appendix.html
96    //TODO #26 need to figure out a better error method to show errors
97    pub fn error_response(
98        severity: PgErrorLevels,
99        code: PgErrorCodes,
100        message: String,
101    ) -> NetworkFrame {
102        let mut buffer = BytesMut::new();
103        buffer.put_u8(b'S'); //Severity
104        buffer.put(severity.value());
105        buffer.put_u8(b'\0');
106        buffer.put_u8(b'M'); //Code
107        buffer.put(message.as_bytes());
108        buffer.put_u8(b'\0');
109        buffer.put_u8(b'C'); //Code
110        buffer.put(code.value());
111        buffer.put_u8(b'\0');
112        buffer.put_u8(b'\0');
113
114        NetworkFrame::new(b'E', buffer.freeze())
115    }
116}
117
118#[derive(Error, Debug)]
119pub enum NetworkFrameError {
120    #[error(transparent)]
121    TooManyFields(#[from] TryFromIntError),
122}