Skip to main content

postgres_protocol/message/
shared.rs

1#![allow(missing_docs)]
2
3use std::io::Read;
4use std::{cmp, io};
5
6use byteorder::{BigEndian, ByteOrder};
7use bytes::Bytes;
8use memchr::memchr;
9
10#[derive(Debug, Copy, Clone)]
11pub struct Header {
12    tag: u8,
13    len: i32,
14}
15
16#[allow(clippy::len_without_is_empty)]
17impl Header {
18    #[inline]
19    pub fn parse(buf: &[u8]) -> io::Result<Option<Header>> {
20        if buf.len() < 5 {
21            return Ok(None);
22        }
23
24        let tag = buf[0];
25        let len = BigEndian::read_i32(&buf[1..]);
26
27        if len < 4 {
28            return Err(io::Error::new(
29                io::ErrorKind::InvalidData,
30                "invalid message length: header length < 4",
31            ));
32        }
33
34        Ok(Some(Header { tag, len }))
35    }
36
37    #[inline]
38    pub fn tag(self) -> u8 {
39        self.tag
40    }
41
42    #[inline]
43    pub fn len(self) -> i32 {
44        self.len
45    }
46}
47
48pub struct Buffer {
49    pub bytes: Bytes,
50    pub idx: usize,
51}
52
53impl Buffer {
54    #[inline]
55    pub fn slice(&self) -> &[u8] {
56        &self.bytes[self.idx..]
57    }
58
59    #[inline]
60    pub fn is_empty(&self) -> bool {
61        self.slice().is_empty()
62    }
63
64    #[inline]
65    pub fn read_cstr(&mut self) -> io::Result<Bytes> {
66        match memchr(0, self.slice()) {
67            Some(pos) => {
68                let start = self.idx;
69                let end = start + pos;
70                let cstr = self.bytes.slice(start..end);
71                self.idx = end + 1;
72                Ok(cstr)
73            }
74            None => Err(io::Error::new(
75                io::ErrorKind::UnexpectedEof,
76                "unexpected EOF",
77            )),
78        }
79    }
80
81    #[inline]
82    pub fn read_all(&mut self) -> Bytes {
83        let buf = self.bytes.slice(self.idx..);
84        self.idx = self.bytes.len();
85        buf
86    }
87}
88
89impl Read for Buffer {
90    #[inline]
91    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
92        let len = {
93            let slice = self.slice();
94            let len = cmp::min(slice.len(), buf.len());
95            buf[..len].copy_from_slice(&slice[..len]);
96            len
97        };
98        self.idx += len;
99        Ok(len)
100    }
101}