From 5f5431ecdae75c7e8aba0f7aebcfc2e0102b70dc Mon Sep 17 00:00:00 2001
From: junderw <jonathan.underwood4649@gmail.com>
Date: Wed, 4 Sep 2024 22:38:35 +0900
Subject: [PATCH] Fix EOF panics

---
 src/version1.rs | 2 ++
 src/version2.rs | 4 ++--
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/version1.rs b/src/version1.rs
index 19fc38b..c85b1ee 100644
--- a/src/version1.rs
+++ b/src/version1.rs
@@ -175,6 +175,8 @@ pub(crate) fn parse(buf: &mut impl Buf) -> Result<super::ProxyHeader, ParseError
     buf.advance(count);
 
     ensure!(buf.get_u8() == CR, IllegalHeaderEnding);
+    // We only checked up to the CR
+    ensure!(buf.remaining() >= 1, UnexpectedEof);
     ensure!(buf.get_u8() == LF, IllegalHeaderEnding);
 
     let addresses = match (source, destination) {
diff --git a/src/version2.rs b/src/version2.rs
index 6148e62..b1ca63f 100644
--- a/src/version2.rs
+++ b/src/version2.rs
@@ -142,7 +142,7 @@ pub(crate) fn parse(buf: &mut impl Buf) -> Result<super::ProxyHeader, ParseError
                 needs: 108usize * 2,
             },
         );
-        ensure!(buf.remaining() >= 108 * 2, UnexpectedEof);
+        ensure!(buf.remaining() >= length, UnexpectedEof);
         let mut source = [0u8; 108];
         let mut destination = [0u8; 108];
         buf.copy_to_slice(&mut source[..]);
@@ -177,7 +177,7 @@ pub(crate) fn parse(buf: &mut impl Buf) -> Result<super::ProxyHeader, ParseError
         },
     );
     ensure!(
-        buf.remaining() >= port_length + address_length,
+        buf.remaining() >= length,
         UnexpectedEof,
     );