1use bytes::BytesMut;
12use tokio_util::codec::{Decoder, Encoder};
13
14use crate::error::DapError;
15
16pub struct ContentLengthCodec {
21 content_length: Option<usize>,
23}
24
25impl ContentLengthCodec {
26 #[must_use]
28 pub const fn new() -> Self {
29 Self {
30 content_length: None,
31 }
32 }
33}
34
35impl Default for ContentLengthCodec {
36 fn default() -> Self {
37 Self::new()
38 }
39}
40
41impl Decoder for ContentLengthCodec {
42 type Item = serde_json::Value;
43 type Error = DapError;
44
45 fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
46 if self.content_length.is_none() {
48 if let Some(header_end) = find_header_end(src) {
49 let header = std::str::from_utf8(&src[..header_end])
50 .map_err(|e| DapError::Codec(format!("invalid header encoding: {e}")))?;
51
52 let content_length = parse_content_length(header)?;
53 self.content_length = Some(content_length);
54
55 let _ = src.split_to(header_end + 4);
57 } else {
58 return Ok(None); }
60 }
61
62 if let Some(len) = self.content_length
64 && src.len() >= len
65 {
66 let body = src.split_to(len);
67 self.content_length = None;
68 let value: serde_json::Value = serde_json::from_slice(&body)?;
69 return Ok(Some(value));
70 }
71
72 Ok(None) }
74}
75
76impl Encoder<serde_json::Value> for ContentLengthCodec {
77 type Error = DapError;
78
79 fn encode(&mut self, item: serde_json::Value, dst: &mut BytesMut) -> Result<(), Self::Error> {
80 let body = serde_json::to_vec(&item)?;
81 let header = format!("Content-Length: {}\r\n\r\n", body.len());
82 dst.extend_from_slice(header.as_bytes());
83 dst.extend_from_slice(&body);
84 Ok(())
85 }
86}
87
88fn find_header_end(buf: &[u8]) -> Option<usize> {
90 buf.windows(4).position(|w| w == b"\r\n\r\n")
91}
92
93fn parse_content_length(header: &str) -> Result<usize, DapError> {
95 for line in header.lines() {
96 if let Some(value) = line.strip_prefix("Content-Length:") {
97 return value
98 .trim()
99 .parse::<usize>()
100 .map_err(|e| DapError::Codec(format!("invalid Content-Length: {e}")));
101 }
102 }
103 Err(DapError::Codec("missing Content-Length header".into()))
104}
105
106#[cfg(test)]
107#[allow(clippy::unwrap_used)]
108mod tests {
109 use super::*;
110 use bytes::BytesMut;
111 use tokio_util::codec::{Decoder, Encoder};
112
113 #[test]
114 fn decode_single_message() {
115 let mut codec = ContentLengthCodec::new();
116 let body = r#"{"seq":1,"type":"response"}"#;
117 let frame = format!("Content-Length: {}\r\n\r\n{body}", body.len());
118 let mut buf = BytesMut::from(frame.as_str());
119
120 let result = codec.decode(&mut buf).unwrap();
121 assert!(result.is_some());
122 let value = result.unwrap();
123 assert_eq!(value["seq"], 1);
124 }
125
126 #[test]
127 fn decode_partial_header() {
128 let mut codec = ContentLengthCodec::new();
129 let mut buf = BytesMut::from("Content-Len");
130
131 let result = codec.decode(&mut buf).unwrap();
132 assert!(result.is_none());
133 }
134
135 #[test]
136 fn decode_partial_body() {
137 let mut codec = ContentLengthCodec::new();
138 let body = r#"{"seq":1,"type":"response"}"#;
139 let frame = format!("Content-Length: {}\r\n\r\n", body.len());
140 let mut buf = BytesMut::from(frame.as_str());
142
143 let result = codec.decode(&mut buf).unwrap();
144 assert!(result.is_none());
145
146 buf.extend_from_slice(body.as_bytes());
148 let result = codec.decode(&mut buf).unwrap();
149 assert!(result.is_some());
150 }
151
152 #[test]
153 fn decode_two_messages() {
154 let mut codec = ContentLengthCodec::new();
155 let body1 = r#"{"seq":1}"#;
156 let body2 = r#"{"seq":2}"#;
157 let frame = format!(
158 "Content-Length: {}\r\n\r\n{body1}Content-Length: {}\r\n\r\n{body2}",
159 body1.len(),
160 body2.len()
161 );
162 let mut buf = BytesMut::from(frame.as_str());
163
164 let r1 = codec.decode(&mut buf).unwrap().unwrap();
165 assert_eq!(r1["seq"], 1);
166 let r2 = codec.decode(&mut buf).unwrap().unwrap();
167 assert_eq!(r2["seq"], 2);
168 }
169
170 #[test]
171 fn encode_message() {
172 let mut codec = ContentLengthCodec::new();
173 let value = serde_json::json!({"seq": 1, "type": "request"});
174 let mut buf = BytesMut::new();
175 codec.encode(value, &mut buf).unwrap();
176
177 let s = std::str::from_utf8(&buf).unwrap();
178 assert!(s.starts_with("Content-Length: "));
179 assert!(s.contains("\r\n\r\n"));
180 assert!(s.contains("\"seq\""));
181 }
182
183 #[test]
184 fn missing_content_length_header() {
185 let mut codec = ContentLengthCodec::new();
186 let mut buf = BytesMut::from("X-Custom: value\r\n\r\n{}");
187 let result = codec.decode(&mut buf);
188 assert!(result.is_err());
189 }
190}