File size: 5,307 Bytes
ea39c0e | 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 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 | use std::io;
use std::mem::size_of;
use serde::Serialize;
use serde::de::DeserializeOwned;
use tokio::io::AsyncRead;
use tokio::io::AsyncReadExt;
use tokio::io::AsyncWrite;
use tokio::io::AsyncWriteExt;
/// Maximum JSON payload size accepted for one code-mode host frame.
pub const MAX_FRAME_BYTES: usize = 64 * 1024 * 1024;
/// A serialized IPC frame that has already passed the payload size limit.
#[derive(Clone, Debug)]
pub struct EncodedFrame {
payload: Vec<u8>,
}
impl EncodedFrame {
pub fn encode<T>(message: &T) -> io::Result<Self>
where
T: Serialize,
{
let payload = serde_json::to_vec(message).map_err(|err| {
io::Error::new(
io::ErrorKind::InvalidData,
format!("failed to encode code-mode IPC frame: {err}"),
)
})?;
if payload.len() > MAX_FRAME_BYTES {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!(
"code-mode IPC frame length {} exceeds {MAX_FRAME_BYTES} bytes",
payload.len()
),
));
}
Ok(Self { payload })
}
/// Returns the complete length-prefixed representation of this frame.
pub fn into_framed_bytes(self) -> Vec<u8> {
let mut bytes = Vec::with_capacity(size_of::<u32>() + self.payload.len());
bytes.extend_from_slice(&(self.payload.len() as u32).to_le_bytes());
bytes.extend_from_slice(&self.payload);
bytes
}
/// Decodes exactly one complete length-prefixed frame.
pub fn decode_framed<T>(bytes: &[u8]) -> io::Result<T>
where
T: DeserializeOwned,
{
let length_bytes: [u8; size_of::<u32>()] = bytes
.get(..size_of::<u32>())
.and_then(|length_bytes| length_bytes.try_into().ok())
.ok_or_else(|| {
io::Error::new(
io::ErrorKind::InvalidData,
"code-mode IPC frame is missing its length prefix",
)
})?;
let length = u32::from_le_bytes(length_bytes) as usize;
if length > MAX_FRAME_BYTES {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!("code-mode IPC frame length {length} exceeds {MAX_FRAME_BYTES} bytes"),
));
}
let payload = &bytes[size_of::<u32>()..];
if payload.len() != length {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!(
"code-mode IPC frame declares {length} payload bytes but contains {}",
payload.len()
),
));
}
serde_json::from_slice(payload).map_err(|err| {
io::Error::new(
io::ErrorKind::InvalidData,
format!("failed to decode code-mode IPC frame: {err}"),
)
})
}
}
/// Decodes JSON messages prefixed by a four-byte little-endian payload length.
pub struct FramedReader<R> {
reader: R,
}
impl<R> FramedReader<R>
where
R: AsyncRead + Unpin,
{
pub fn new(reader: R) -> Self {
Self { reader }
}
/// Reads the next frame, returning `None` only for EOF at a frame boundary.
pub async fn read<T>(&mut self) -> io::Result<Option<T>>
where
T: DeserializeOwned,
{
let mut length_bytes = [0_u8; size_of::<u32>()];
if self.reader.read(&mut length_bytes[..1]).await? == 0 {
return Ok(None);
}
self.reader.read_exact(&mut length_bytes[1..]).await?;
let length = u32::from_le_bytes(length_bytes) as usize;
if length > MAX_FRAME_BYTES {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!("code-mode IPC frame length {length} exceeds {MAX_FRAME_BYTES} bytes"),
));
}
let mut payload = vec![0; length];
self.reader.read_exact(&mut payload).await?;
serde_json::from_slice(&payload).map(Some).map_err(|err| {
io::Error::new(
io::ErrorKind::InvalidData,
format!("failed to decode code-mode IPC frame: {err}"),
)
})
}
}
/// Encodes JSON messages with a four-byte little-endian payload length.
pub struct FramedWriter<W> {
writer: W,
}
impl<W> FramedWriter<W>
where
W: AsyncWrite + Unpin,
{
pub fn new(writer: W) -> Self {
Self { writer }
}
/// Writes and flushes one complete frame.
pub async fn write<T>(&mut self, message: &T) -> io::Result<()>
where
T: Serialize,
{
self.write_frame(&EncodedFrame::encode(message)?).await
}
/// Writes and flushes a frame encoded before it entered an I/O queue.
pub async fn write_frame(&mut self, frame: &EncodedFrame) -> io::Result<()> {
let length = u32::try_from(frame.payload.len()).map_err(|_| {
io::Error::new(
io::ErrorKind::InvalidData,
"code-mode IPC frame length exceeds u32",
)
})?;
self.writer.write_all(&length.to_le_bytes()).await?;
self.writer.write_all(&frame.payload).await?;
self.writer.flush().await
}
}
|