From 401fd144cd666b595cd343cfc7cb45940c1eae08 Mon Sep 17 00:00:00 2001
From: Alex Janka <alex@alexjanka.com>
Date: Thu, 7 Mar 2024 09:24:20 +1100
Subject: [PATCH] timeout and proper chunk handling

---
 homekit-controller/src/homekit_http.rs | 87 +++++++++++++++++++++-----
 1 file changed, 70 insertions(+), 17 deletions(-)

diff --git a/homekit-controller/src/homekit_http.rs b/homekit-controller/src/homekit_http.rs
index 8b109f7..80d8f7b 100644
--- a/homekit-controller/src/homekit_http.rs
+++ b/homekit-controller/src/homekit_http.rs
@@ -1,4 +1,8 @@
-use std::{collections::HashMap, time::Duration};
+use std::{
+    collections::{HashMap, VecDeque},
+    mem,
+    time::Duration,
+};
 
 use chacha20poly1305::{
     aead::generic_array::GenericArray, AeadInPlace, ChaCha20Poly1305, KeyInit, Nonce,
@@ -228,30 +232,30 @@ impl AccessorySocket {
         {
             if transfer_encoding.value == b"chunked" {
                 loop {
-                    packet.append(&mut self.get_next().await?);
+                    let chunked = ChunkedTransfer::clone_from(&packet).collect::<Vec<Vec<u8>>>();
 
-                    let utf8_decoded = String::from_utf8(packet.clone())?;
+                    if let Some(last) = chunked.chunks(2).last() {
+                        if last.len() == 2 && last[0] == b"0" {
+                            break;
+                        }
+                    }
 
-                    let split = utf8_decoded
-                        .split_terminator("\r\n")
-                        .map(String::from)
-                        .collect::<Vec<_>>();
-                    if let Some(last) = split.chunks(2).last() {
-                        if last.len() == 2 && last[0] == "0" {
+                    match tokio::time::timeout(Duration::from_secs(2), self.get_next()).await {
+                        Ok(next) => {
+                            packet.append(&mut next?);
+                        }
+                        Err(_) => {
+                            log::error!("timed out");
                             break;
                         }
                     }
                 }
 
-                let utf8_decoded = String::from_utf8(std::mem::take(&mut packet))?;
+                let mut chunked =
+                    ChunkedTransfer::from(std::mem::take(&mut packet)).collect::<Vec<Vec<u8>>>();
 
-                let split = utf8_decoded
-                    .split_terminator("\r\n")
-                    .map(String::from)
-                    .collect::<Vec<_>>();
-
-                for chunk in split.chunks_exact(2) {
-                    packet.extend_from_slice(chunk[1].as_bytes())
+                for chunk in chunked.chunks_exact_mut(2) {
+                    packet.append(&mut chunk[1]);
                 }
             }
         }
@@ -329,3 +333,52 @@ pub enum DiscoveryError {
     #[error("not found")]
     NotFound,
 }
+
+struct ChunkedTransfer {
+    data: Vec<u8>,
+}
+
+impl From<Vec<u8>> for ChunkedTransfer {
+    fn from(value: Vec<u8>) -> Self {
+        Self { data: value }
+    }
+}
+
+impl ChunkedTransfer {
+    fn clone_from(value: &[u8]) -> Self {
+        Self {
+            data: value.to_owned(),
+        }
+    }
+}
+
+impl Iterator for ChunkedTransfer {
+    type Item = Vec<u8>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.data.is_empty() {
+            None
+        } else {
+            let mut crlf_index = None;
+            for (i, val) in self.data.iter().enumerate() {
+                if let Some(next) = self.data.get(i + 1) {
+                    if *val == b'\r' && *next == b'\n' {
+                        crlf_index = Some(i);
+                        break;
+                    }
+                }
+            }
+
+            if let Some(i) = crlf_index {
+                let remainder = self.data.split_off(i);
+                let mut remainder = VecDeque::from(remainder);
+                let _ = remainder.pop_front().map(|v| v as char);
+                let _ = remainder.pop_front().map(|v| v as char);
+
+                Some(mem::replace(&mut self.data, Vec::from(remainder)))
+            } else {
+                Some(mem::take(&mut self.data))
+            }
+        }
+    }
+}