Add length check for frames.

jjeffcaii · jjeffcaii · commit 0ebbc486aa2a · 2020-05-21T23:03:24.000+08:00
diff --git a/rsocket-test/tests/test_composite_metadata.rs b/rsocket-test/tests/test_composite_metadata.rs
@@ -1,9 +1,9 @@
-use bytes::BytesMut;
+use bytes::{BufMut, BytesMut};
 use rsocket_rust::extension::{self, CompositeMetadata, CompositeMetadataEntry, MimeType};
 use rsocket_rust::utils::Writeable;
 
 #[test]
-fn encode_and_decode_composite_metadata() {
+fn test_encode_and_decode() {
     let bingo = |metadatas: Vec<&CompositeMetadataEntry>| {
         assert_eq!(2, metadatas.len());
         assert_eq!(
@@ -29,3 +29,13 @@ fn encode_and_decode_composite_metadata() {
     let cm2 = CompositeMetadata::decode(&mut bf).unwrap();
     bingo(cm2.iter().collect());
 }
+
+#[test]
+fn test_bad() {
+    let mut bf = BytesMut::new();
+    bf.put_slice(b"must bad");
+    assert!(
+        CompositeMetadata::decode(&mut bf).is_err(),
+        "should be error"
+    )
+}
diff --git a/rsocket-test/tests/test_mimes.rs b/rsocket-test/tests/test_mimes.rs
@@ -2,6 +2,15 @@ extern crate rsocket_rust;
 
 use rsocket_rust::extension::{self, MimeType};
 
+#[test]
+fn test_from_str() {
+    let result = MimeType::from("foobar");
+    match result {
+        MimeType::WellKnown(_) => panic!("failed"),
+        MimeType::Normal(s) => assert_eq!(&s, "foobar"),
+    }
+}
+
 #[test]
 fn test_wellknown() {
     let well = MimeType::from("application/json");
@@ -10,3 +19,20 @@ fn test_wellknown() {
     let custom = MimeType::from("application/custom");
     assert_eq!(MimeType::Normal("application/custom".to_owned()), custom);
 }
+
+#[test]
+fn test_parse() {
+    assert!(MimeType::parse(0xFF).is_none(), "should be none");
+    for x in 0..0x29 {
+        assert!(MimeType::parse(x).is_some(), "should not be none");
+    }
+    for x in 0x29..0x7A {
+        assert!(MimeType::parse(x).is_none(), "should be none");
+    }
+    for x in 0x7A..0x80 {
+        assert!(MimeType::parse(x).is_some(), "should not be none");
+    }
+    for x in 0x80..0xFF {
+        assert!(MimeType::parse(x).is_none(), "should be none");
+    }
+}
diff --git a/rsocket/src/error.rs b/rsocket/src/error.rs
@@ -19,6 +19,7 @@ pub enum ErrorKind {
     WithDescription(String),
     IO(io::Error),
     Cancelled(),
+    LengthTooShort(usize),
 }
 
 #[derive(Debug)]
@@ -35,6 +36,7 @@ impl fmt::Display for RSocketError {
             ErrorKind::WithDescription(s) => write!(f, "{}", s),
             ErrorKind::IO(e) => write!(f, "{}", e),
             ErrorKind::Cancelled() => write!(f, "ERROR(CANCELLED)"),
+            ErrorKind::LengthTooShort(n) => write!(f, "ERROR(MINIMAL LENGTH {})", n),
         }
     }
 }
diff --git a/rsocket/src/extension/mime.rs b/rsocket/src/extension/mime.rs
@@ -16,9 +16,6 @@ lazy_static! {
         }
         m
     };
-}
-
-lazy_static! {
     static ref STR_TO_U8: HashMap<&'static str, u8> = {
         let mut m = HashMap::new();
         for it in list_all().iter() {
@@ -58,7 +55,7 @@ impl AsRef<str> for MimeType {
 
 impl From<&str> for MimeType {
     fn from(value: &str) -> MimeType {
-        match STR_TO_U8.get(&value) {
+        match STR_TO_U8.get(value) {
             Some(v) => Self::WellKnown(*v),
             None => Self::Normal(value.to_owned()),
         }
diff --git a/rsocket/src/frame/error.rs b/rsocket/src/frame/error.rs
@@ -1,3 +1,4 @@
+use super::utils::too_short;
 use super::{Body, Frame};
 use crate::utils::{RSocketResult, Writeable};
 use bytes::{Buf, BufMut, Bytes, BytesMut};
@@ -44,13 +45,17 @@ impl ErrorBuilder {
 
 impl Error {
     pub(crate) fn decode(flag: u16, bf: &mut BytesMut) -> RSocketResult<Error> {
-        let code = bf.get_u32();
-        let d: Option<Bytes> = if !bf.is_empty() {
-            Some(bf.to_bytes())
+        if bf.len() < 4 {
+            too_short(4)
         } else {
-            None
-        };
-        Ok(Error { code, data: d })
+            let code = bf.get_u32();
+            let data: Option<Bytes> = if !bf.is_empty() {
+                Some(bf.to_bytes())
+            } else {
+                None
+            };
+            Ok(Error { code, data })
+        }
     }
 
     pub fn builder(stream_id: u32, flag: u16) -> ErrorBuilder {
diff --git a/rsocket/src/frame/keepalive.rs b/rsocket/src/frame/keepalive.rs
@@ -1,3 +1,4 @@
+use super::utils::too_short;
 use super::{Body, Frame};
 use crate::utils::{RSocketResult, Writeable};
 use bytes::{Buf, BufMut, Bytes, BytesMut};
@@ -43,14 +44,18 @@ impl KeepaliveBuilder {
 
 impl Keepalive {
     pub(crate) fn decode(flag: u16, bf: &mut BytesMut) -> RSocketResult<Keepalive> {
-        let position = bf.get_u64();
-        let mut d: Option<Bytes> = None;
-        if !bf.is_empty() {
-            d = Some(bf.to_bytes());
+        if bf.len() < 8 {
+            return too_short(8);
         }
+        let position = bf.get_u64();
+        let data = if bf.is_empty() {
+            None
+        } else {
+            Some(bf.to_bytes())
+        };
         Ok(Keepalive {
             last_received_position: position,
-            data: d,
+            data,
         })
     }
 
diff --git a/rsocket/src/frame/lease.rs b/rsocket/src/frame/lease.rs
@@ -1,3 +1,4 @@
+use super::utils::too_short;
 use super::{Body, Frame};
 use crate::utils::{RSocketResult, Writeable};
 use bytes::{Buf, BufMut, Bytes, BytesMut};
@@ -51,6 +52,9 @@ impl LeaseBuilder {
 
 impl Lease {
     pub(crate) fn decode(flag: u16, bf: &mut BytesMut) -> RSocketResult<Lease> {
+        if bf.len() < 8 {
+            return too_short(8);
+        }
         let ttl = bf.get_u32();
         let n = bf.get_u32();
         let m = if flag & Frame::FLAG_METADATA != 0 {
diff --git a/rsocket/src/frame/mod.rs b/rsocket/src/frame/mod.rs
@@ -1,4 +1,4 @@
-use crate::error::RSocketError;
+use crate::error::{ErrorKind, RSocketError};
 use crate::utils::{RSocketResult, Writeable};
 use bytes::{Buf, BufMut, Bytes, BytesMut};
 
@@ -145,7 +145,9 @@ impl Frame {
     }
 
     pub fn decode(b: &mut BytesMut) -> RSocketResult<Frame> {
-        // TODO: check size
+        if b.len() < LEN_HEADER {
+            return Err(ErrorKind::LengthTooShort(LEN_HEADER).into());
+        }
         let sid = b.get_u32();
         let n = b.get_u16();
         let (flag, kind) = (n & 0x03FF, (n & 0xFC00) >> 10);
diff --git a/rsocket/src/frame/payload.rs b/rsocket/src/frame/payload.rs
@@ -1,3 +1,4 @@
+use super::utils::{read_payload, too_short};
 use super::{Body, Frame, PayloadSupport};
 use crate::utils::{RSocketResult, Writeable};
 use bytes::{BufMut, Bytes, BytesMut};
@@ -59,11 +60,7 @@ impl PayloadBuilder {
 
 impl Payload {
     pub(crate) fn decode(flag: u16, bf: &mut BytesMut) -> RSocketResult<Payload> {
-        let (m, d) = PayloadSupport::read(flag, bf);
-        Ok(Payload {
-            metadata: m,
-            data: d,
-        })
+        read_payload(flag, bf).map(|(metadata, data)| Payload { metadata, data })
     }
 
     pub fn builder(stream_id: u32, flag: u16) -> PayloadBuilder {
diff --git a/rsocket/src/frame/request_channel.rs b/rsocket/src/frame/request_channel.rs
@@ -1,3 +1,4 @@
+use super::utils::{read_payload, too_short};
 use super::{Body, Frame, PayloadSupport, REQUEST_MAX};
 use crate::utils::{RSocketResult, Writeable};
 use bytes::{Buf, BufMut, Bytes, BytesMut};
@@ -66,13 +67,16 @@ impl RequestChannelBuilder {
 
 impl RequestChannel {
     pub(crate) fn decode(flag: u16, bf: &mut BytesMut) -> RSocketResult<RequestChannel> {
-        let n = bf.get_u32();
-        let (m, d) = PayloadSupport::read(flag, bf);
-        Ok(RequestChannel {
-            initial_request_n: n,
-            metadata: m,
-            data: d,
-        })
+        if bf.len() < 4 {
+            too_short(4)
+        } else {
+            let initial_request_n = bf.get_u32();
+            read_payload(flag, bf).map(move |(metadata, data)| RequestChannel {
+                initial_request_n,
+                metadata,
+                data,
+            })
+        }
     }
 
     pub fn builder(stream_id: u32, flag: u16) -> RequestChannelBuilder {
diff --git a/rsocket/src/frame/request_fnf.rs b/rsocket/src/frame/request_fnf.rs
@@ -1,3 +1,4 @@
+use super::utils::{read_payload, too_short};
 use super::{Body, Frame, PayloadSupport};
 use crate::utils::{RSocketResult, Writeable};
 use bytes::{BufMut, Bytes, BytesMut};
@@ -59,8 +60,7 @@ impl RequestFNFBuilder {
 
 impl RequestFNF {
     pub(crate) fn decode(flag: u16, bf: &mut BytesMut) -> RSocketResult<RequestFNF> {
-        let (m, d) = PayloadSupport::read(flag, bf);
-        Ok(RequestFNF {
+        read_payload(flag, bf).map(|(m, d)| RequestFNF {
             metadata: m,
             data: d,
         })
diff --git a/rsocket/src/frame/request_n.rs b/rsocket/src/frame/request_n.rs
@@ -1,3 +1,4 @@
+use super::utils::too_short;
 use super::{Body, Frame, REQUEST_MAX};
 use crate::utils::{RSocketResult, Writeable};
 use bytes::{Buf, BufMut, Bytes, BytesMut};
@@ -34,8 +35,12 @@ impl RequestNBuilder {
 
 impl RequestN {
     pub(crate) fn decode(flag: u16, bf: &mut BytesMut) -> RSocketResult<RequestN> {
-        let n = bf.get_u32();
-        Ok(RequestN { n })
+        if bf.len() < 4 {
+            too_short(4)
+        } else {
+            let n = bf.get_u32();
+            Ok(RequestN { n })
+        }
     }
 
     pub fn builder(stream_id: u32, flag: u16) -> RequestNBuilder {
diff --git a/rsocket/src/frame/request_response.rs b/rsocket/src/frame/request_response.rs
@@ -1,3 +1,4 @@
+use super::utils::{read_payload, too_short};
 use super::{Body, Frame, PayloadSupport};
 use crate::utils::{RSocketResult, Writeable};
 use bytes::{BufMut, Bytes, BytesMut};
@@ -59,8 +60,7 @@ impl RequestResponseBuilder {
 
 impl RequestResponse {
     pub(crate) fn decode(flag: u16, bf: &mut BytesMut) -> RSocketResult<RequestResponse> {
-        let (m, d) = PayloadSupport::read(flag, bf);
-        Ok(RequestResponse {
+        read_payload(flag, bf).map(|(m, d)| RequestResponse {
             metadata: m,
             data: d,
         })
diff --git a/rsocket/src/frame/request_stream.rs b/rsocket/src/frame/request_stream.rs
@@ -1,3 +1,4 @@
+use super::utils::{read_payload, too_short};
 use super::{Body, Frame, PayloadSupport, REQUEST_MAX};
 use crate::utils::{RSocketResult, Writeable};
 use bytes::{Buf, BufMut, Bytes, BytesMut};
@@ -54,13 +55,16 @@ impl RequestStreamBuilder {
 
 impl RequestStream {
     pub(crate) fn decode(flag: u16, bf: &mut BytesMut) -> RSocketResult<RequestStream> {
-        let n = bf.get_u32();
-        let (m, d) = PayloadSupport::read(flag, bf);
-        Ok(RequestStream {
-            initial_request_n: n,
-            metadata: m,
-            data: d,
-        })
+        if bf.len() < 4 {
+            too_short(4)
+        } else {
+            let initial_request_n = bf.get_u32();
+            read_payload(flag, bf).map(move |(metadata, data)| RequestStream {
+                initial_request_n,
+                metadata,
+                data,
+            })
+        }
     }
 
     pub fn builder(stream_id: u32, flag: u16) -> RequestStreamBuilder {
diff --git a/rsocket/src/frame/resume.rs b/rsocket/src/frame/resume.rs
@@ -1,3 +1,4 @@
+use super::utils::too_short;
 use super::{Body, Frame, Version};
 use crate::utils::{RSocketResult, Writeable};
 use bytes::{Buf, BufMut, Bytes, BytesMut};
@@ -27,21 +28,30 @@ impl Resume {
     }
 
     pub(crate) fn decode(flag: u16, b: &mut BytesMut) -> RSocketResult<Resume> {
+        if b.len() < 6 {
+            return too_short(6);
+        }
         let major = b.get_u16();
         let minor = b.get_u16();
-        let token_size = b.get_u16();
+        let token_size = b.get_u16() as usize;
         let token = if token_size > 0 {
-            Some(b.split_to(token_size as usize).to_bytes())
+            if b.len() < token_size {
+                return too_short(token_size);
+            }
+            Some(b.split_to(token_size).freeze())
         } else {
             None
         };
-        let p1 = b.get_u64();
-        let p2 = b.get_u64();
+        if b.len() < 16 {
+            return too_short(16);
+        }
+        let last_received_server_position = b.get_u64();
+        let first_available_client_position = b.get_u64();
         Ok(Resume {
             version: Version::new(major, minor),
             token,
-            last_received_server_position: p1,
-            first_available_client_position: p2,
+            last_received_server_position,
+            first_available_client_position,
         })
     }
 
diff --git a/rsocket/src/frame/resume_ok.rs b/rsocket/src/frame/resume_ok.rs
@@ -1,3 +1,4 @@
+use super::utils::too_short;
 use super::{Body, Frame};
 use crate::utils::{RSocketResult, Writeable};
 use bytes::{Buf, BufMut, Bytes, BytesMut};
@@ -33,8 +34,13 @@ impl ResumeOKBuilder {
 
 impl ResumeOK {
     pub(crate) fn decode(flag: u16, bf: &mut BytesMut) -> RSocketResult<ResumeOK> {
-        let position = bf.get_u64();
-        Ok(ResumeOK { position })
+        if bf.len() < 8 {
+            too_short(8)
+        } else {
+            Ok(ResumeOK {
+                position: bf.get_u64(),
+            })
+        }
     }
 
     pub fn builder(stream_id: u32, flag: u16) -> ResumeOKBuilder {
diff --git a/rsocket/src/frame/setup.rs b/rsocket/src/frame/setup.rs
diff --git a/rsocket/src/frame/utils.rs b/rsocket/src/frame/utils.rs

Original file line number	Diff line number	Diff line change
`@@ -19,6 +19,7 @@ pub enum ErrorKind {`
`19`	`19`	`WithDescription(String),`
`20`	`20`	`IO(io::Error),`
`21`	`21`	`Cancelled(),`
	`22`	`+ LengthTooShort(usize),`
`22`	`23`	`}`
`23`	`24`
`24`	`25`	`#[derive(Debug)]`
`@@ -35,6 +36,7 @@ impl fmt::Display for RSocketError {`
`35`	`36`	`ErrorKind::WithDescription(s) => write!(f, "{}", s),`
`36`	`37`	`ErrorKind::IO(e) => write!(f, "{}", e),`
`37`	`38`	`ErrorKind::Cancelled() => write!(f, "ERROR(CANCELLED)"),`
	`39`	`+ ErrorKind::LengthTooShort(n) => write!(f, "ERROR(MINIMAL LENGTH {})", n),`
`38`	`40`	`}`
`39`	`41`	`}`
`40`	`42`	`}`
Original file line number	Diff line number	Diff line change
`@@ -16,9 +16,6 @@ lazy_static! {`
`16`	`16`	`}`
`17`	`17`	`m`
`18`	`18`	`};`
`19`		`-}`
`20`		`-`
`21`		`-lazy_static! {`
`22`	`19`	`static ref STR_TO_U8: HashMap<&'static str, u8> = {`
`23`	`20`	`let mut m = HashMap::new();`
`24`	`21`	`for it in list_all().iter() {`
`@@ -58,7 +55,7 @@ impl AsRef<str> for MimeType {`
`58`	`55`
`59`	`56`	`impl From<&str> for MimeType {`
`60`	`57`	`fn from(value: &str) -> MimeType {`
`61`		`- match STR_TO_U8.get(&value) {`
	`58`	`+ match STR_TO_U8.get(value) {`
`62`	`59`	`Some(v) => Self::WellKnown(*v),`
`63`	`60`	`None => Self::Normal(value.to_owned()),`
`64`	`61`	`}`