From 8385dda201d97b3aef8b7af31fec8c6542d53db4 Mon Sep 17 00:00:00 2001
From: Jack Grigg <thestr4d@gmail.com>
Date: Tue, 10 Dec 2019 13:22:21 +0000
Subject: [PATCH] Check buffer length in set_length

---
 src/metadata.rs      |  2 +-
 src/serialization.rs | 30 +++++++++++++++++++-----------
 2 files changed, 20 insertions(+), 12 deletions(-)

diff --git a/src/metadata.rs b/src/metadata.rs
index e2cc3ef..5c513cb 100644
--- a/src/metadata.rs
+++ b/src/metadata.rs
@@ -136,7 +136,7 @@ pub(crate) fn set_item(
     // Re-encode item and insert
     if cb_item != 0 {
         offset -= cb_len;
-        offset += set_length(&mut data[offset..], cb_item);
+        offset += set_length(&mut data[offset..], cb_item)?;
         data[offset..offset + cb_item].copy_from_slice(p_item);
     }
 
diff --git a/src/serialization.rs b/src/serialization.rs
index 55fc699..a12cd75 100644
--- a/src/serialization.rs
+++ b/src/serialization.rs
@@ -85,8 +85,7 @@ impl<'a> Tlv<'a> {
         }
         buffer[0] = tag;
 
-        // TODO: Raise error
-        let offset = 1 + set_length(&mut buffer[1..], value.len());
+        let offset = 1 + set_length(&mut buffer[1..], value.len())?;
 
         if buffer.len() < offset + value.len() {
             return Err(Error::SizeError);
@@ -113,8 +112,7 @@ impl<'a> Tlv<'a> {
         }
         buffer[0] = tag;
 
-        // TODO: Raise error
-        let offset = 1 + set_length(&mut buffer[1..], length);
+        let offset = 1 + set_length(&mut buffer[1..], length)?;
 
         if buffer.len() < offset + length {
             return Err(Error::SizeError);
@@ -127,19 +125,29 @@ impl<'a> Tlv<'a> {
 
 /// Set length
 #[cfg(feature = "untested")]
-pub(crate) fn set_length(buffer: &mut [u8], length: usize) -> usize {
+pub(crate) fn set_length(buffer: &mut [u8], length: usize) -> Result<usize, Error> {
     if length < 0x80 {
-        buffer[0] = length as u8;
-        1
+        if buffer.is_empty() {
+            Err(Error::SizeError)
+        } else {
+            buffer[0] = length as u8;
+            Ok(1)
+        }
     } else if length < 0x100 {
-        buffer[0] = 0x81;
-        buffer[1] = length as u8;
-        2
+        if buffer.len() < 2 {
+            Err(Error::SizeError)
+        } else {
+            buffer[0] = 0x81;
+            buffer[1] = length as u8;
+            Ok(2)
+        }
+    } else if buffer.len() < 3 {
+        Err(Error::SizeError)
     } else {
         buffer[0] = 0x82;
         buffer[1] = ((length >> 8) & 0xff) as u8;
         buffer[2] = (length & 0xff) as u8;
-        3
+        Ok(3)
     }
 }