use tokio::io::{AsyncRead, AsyncReadExt, BufReader};
use crate::error::LspzError;
pub(crate) async fn read_frame<R: AsyncRead + Unpin>(
reader: &mut BufReader<R>,
) -> Result<Vec<u8>, LspzError> {
let mut buffer = Vec::with_capacity(4096);
loop {
let byte = reader.read_u8().await.map_err(|e| {
if e.kind() == std::io::ErrorKind::UnexpectedEof {
LspzError::ServerExited
} else {
LspzError::Io(e)
}
})?;
buffer.push(byte);
if buffer.len() >= 4 && buffer[buffer.len() - 4..] == [b'\r', b'\n', b'\r', b'\n'] {
break;
}
}
let header = std::str::from_utf8(&buffer)
.map_err(|_| LspzError::Protocol("header is not valid UTF-8".into()))?;
let content_length = parse_content_length(header)?;
let mut body = vec![0u8; content_length as usize];
reader.read_exact(&mut body).await.map_err(|e| {
if e.kind() == std::io::ErrorKind::UnexpectedEof {
LspzError::ServerExited
} else {
LspzError::Io(e)
}
})?;
let mut frame = Vec::with_capacity(buffer.len() + body.len());
frame.extend_from_slice(&buffer);
frame.extend_from_slice(&body);
Ok(frame)
}
pub(crate) fn parse_content_length(header: &str) -> Result<u64, LspzError> {
for line in header.lines() {
let line = line.trim();
if let Some(value) = line.to_lowercase().strip_prefix("content-length:") {
let value = value.trim();
return value
.parse::<u64>()
.map_err(|_| LspzError::Protocol(format!("invalid Content-Length: {value}")));
}
}
Err(LspzError::Protocol("missing Content-Length header".into()))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_parse_content_length_basic() {
let header = "Content-Length: 47\r\n\r\n";
assert_eq!(parse_content_length(header).unwrap(), 47);
}
#[test]
fn test_parse_content_length_case_insensitive() {
let header = "content-length: 128\r\n\r\n";
assert_eq!(parse_content_length(header).unwrap(), 128);
}
#[test]
fn test_parse_content_length_whitespace() {
let header = "Content-Length: 99 \r\n\r\n";
assert_eq!(parse_content_length(header).unwrap(), 99);
}
#[test]
fn test_missing_content_length() {
let header = "\r\n\r\n";
match parse_content_length(header) {
Err(LspzError::Protocol(msg)) => assert!(msg.contains("missing")),
_ => panic!("expected Protocol error"),
}
}
#[test]
fn test_invalid_content_length() {
let header = "Content-Length: abc\r\n\r\n";
match parse_content_length(header) {
Err(LspzError::Protocol(msg)) => assert!(msg.contains("invalid")),
_ => panic!("expected Protocol error"),
}
}
#[tokio::test]
async fn test_read_frame_from_buf() {
use std::io::Cursor;
let frame = b"Content-Length: 6\r\n\r\nhello!";
let mut reader = BufReader::new(Cursor::new(frame));
let result = read_frame(&mut reader).await.unwrap();
assert_eq!(result, frame);
}
#[tokio::test]
async fn test_read_frame_truncated_body() {
use std::io::Cursor;
let frame = b"Content-Length: 10\r\n\r\nhi";
let mut reader = BufReader::new(Cursor::new(frame));
let result = read_frame(&mut reader).await;
assert!(matches!(result, Err(LspzError::ServerExited)));
}
}