diff --git a/quinn-proto/src/packet.rs b/quinn-proto/src/packet.rs index b048c4091..c256be262 100644 --- a/quinn-proto/src/packet.rs +++ b/quinn-proto/src/packet.rs @@ -5,9 +5,7 @@ use err_derive::Error; use crate::{ coding::{self, BufExt, BufMutExt}, - crypto, - shared::ConnectionId, - MAX_CID_SIZE, VERSION, + crypto, ConnectionId, VERSION, }; // Due to packet number encryption, it is impossible to fully decode a header @@ -261,7 +259,8 @@ impl Header { } => { w.write(u8::from(LongHeaderType::Initial) | number.tag()); w.write(VERSION); - Self::encode_cids(w, dst_cid, src_cid); + dst_cid.encode_long(w); + src_cid.encode_long(w); w.write_var(token.len() as u64); w.put_slice(token); w.write::(0); // Placeholder for payload length; see `set_payload_length` @@ -280,7 +279,8 @@ impl Header { } => { w.write(u8::from(LongHeaderType::Standard(ty)) | number.tag()); w.write(VERSION); - Self::encode_cids(w, dst_cid, src_cid); + dst_cid.encode_long(w); + src_cid.encode_long(w); w.write::(0); // Placeholder for payload length; see `set_payload_length` number.encode(w); PartialEncode { @@ -295,7 +295,8 @@ impl Header { } => { w.write(u8::from(LongHeaderType::Retry)); w.write(VERSION); - Self::encode_cids(w, dst_cid, src_cid); + dst_cid.encode_long(w); + src_cid.encode_long(w); PartialEncode { start, header_len: w.len() - start, @@ -329,7 +330,8 @@ impl Header { } => { w.write(0x80u8 | random); w.write::(0); - Self::encode_cids(w, dst_cid, src_cid); + dst_cid.encode_long(w); + src_cid.encode_long(w); PartialEncode { start, header_len: w.len() - start, @@ -339,13 +341,6 @@ impl Header { } } - fn encode_cids(w: &mut W, dst_cid: &ConnectionId, src_cid: &ConnectionId) { - w.put_u8(dst_cid.len() as u8); - w.put_slice(dst_cid); - w.put_u8(src_cid.len() as u8); - w.put_slice(src_cid); - } - /// Whether the packet is encrypted on the wire pub(crate) fn is_protected(&self) -> bool { match *self { @@ -519,7 +514,8 @@ impl PlainHeader { let first = buf.get::()?; if first & LONG_HEADER_FORM == 0 { let spin = first & SPIN_BIT != 0; - let dst_cid = Self::get_cid(buf, local_cid_len)?; + let dst_cid = ConnectionId::new(&buf.bytes()[..local_cid_len]); + buf.advance(local_cid_len); Ok(PlainHeader::Short { first, @@ -529,11 +525,10 @@ impl PlainHeader { } else { let version = buf.get::()?; - let dcil = buf.get::()? as usize; - let dst_cid = Self::get_cid(buf, dcil)?; - - let scil = buf.get::()? as usize; - let src_cid = Self::get_cid(buf, scil)?; + let dst_cid = ConnectionId::decode_long(buf) + .ok_or(PacketDecodeError::InvalidHeader("malformed cid"))?; + let src_cid = ConnectionId::decode_long(buf) + .ok_or(PacketDecodeError::InvalidHeader("malformed cid"))?; // TODO: Support long CIDs for compatibility with future QUIC versions if version == 0 { @@ -576,22 +571,6 @@ impl PlainHeader { } } } - - fn get_cid(buf: &mut R, len: usize) -> Result { - if len > MAX_CID_SIZE { - return Err(PacketDecodeError::InvalidHeader( - "illegal connection ID length", - )); - } - if buf.remaining() < len { - return Err(PacketDecodeError::InvalidHeader( - "connection ID longer than packet", - )); - } - let cid = ConnectionId::new(&buf.bytes()[..len]); - buf.advance(len); - Ok(cid) - } } // An encoded packet number diff --git a/quinn-proto/src/retry_token.rs b/quinn-proto/src/retry_token.rs index 4671b7acc..59796c12a 100644 --- a/quinn-proto/src/retry_token.rs +++ b/quinn-proto/src/retry_token.rs @@ -4,13 +4,12 @@ use std::{ time::{Duration, SystemTime, UNIX_EPOCH}, }; -use bytes::{Buf, BufMut}; +use bytes::BufMut; use crate::{ coding::{BufExt, BufMutExt}, crypto::HmacKey, shared::ConnectionId, - MAX_CID_SIZE, }; // TODO: Use AEAD to hide token details from clients for better stability guarantees: @@ -29,18 +28,15 @@ impl RetryToken { pub fn encode(&self, key: &impl HmacKey, address: &SocketAddr) -> Vec { let mut buf = Vec::new(); - buf.write(self.src_cid.len() as u8); - buf.put_slice(&self.src_cid); - - buf.write(self.dst_cid.len() as u8); - buf.put_slice(&self.dst_cid); - + self.src_cid.encode_long(&mut buf); + self.dst_cid.encode_long(&mut buf); buf.write::( self.issued .duration_since(UNIX_EPOCH) .map(|x| x.as_secs()) .unwrap_or(0), ); + let signature_pos = buf.len(); match address.ip() { IpAddr::V4(x) => buf.put_slice(&x.octets()), @@ -56,23 +52,12 @@ impl RetryToken { pub fn from_bytes(key: &impl HmacKey, address: &SocketAddr, data: &[u8]) -> Result { let mut reader = io::Cursor::new(data); - let src_cid_len = reader.get::().map_err(|_| ())? as usize; - if src_cid_len > reader.remaining() || src_cid_len > MAX_CID_SIZE { - return Err(()); - } - let src_cid = ConnectionId::new(&reader.bytes()[..src_cid_len]); - reader.advance(src_cid_len); - - let dst_cid_len = reader.get::().map_err(|_| ())? as usize; - if dst_cid_len > reader.remaining() || dst_cid_len > MAX_CID_SIZE { - return Err(()); - } - let dst_cid = ConnectionId::new(&reader.bytes()[..dst_cid_len]); - reader.advance(dst_cid_len); + let src_cid = ConnectionId::decode_long(&mut reader).ok_or(())?; + let dst_cid = ConnectionId::decode_long(&mut reader).ok_or(())?; let issued = UNIX_EPOCH + Duration::new(reader.get::().map_err(|_| ())?, 0); - let signature_start = reader.position() as usize; + let signature_start = reader.position() as usize; let mut buf = Vec::new(); buf.put_slice(&data[0..signature_start]); match address.ip() { @@ -96,6 +81,7 @@ mod test { #[test] fn token_sanity() { use super::*; + use crate::MAX_CID_SIZE; use rand::RngCore; use ring::hmac; use std::{ diff --git a/quinn-proto/src/shared.rs b/quinn-proto/src/shared.rs index 3153b625b..dc43f1f69 100644 --- a/quinn-proto/src/shared.rs +++ b/quinn-proto/src/shared.rs @@ -1,9 +1,9 @@ use std::{cmp, fmt, net::SocketAddr, time::Instant}; -use bytes::BytesMut; +use bytes::{Buf, BufMut, BytesMut}; use rand::Rng; -use crate::{packet::PartialDecode, MAX_CID_SIZE, RESET_TOKEN_SIZE}; +use crate::{coding::BufExt, packet::PartialDecode, MAX_CID_SIZE, RESET_TOKEN_SIZE}; /// Events sent from an Endpoint to a Connection #[derive(Debug)] @@ -87,6 +87,23 @@ impl ConnectionId { res.bytes[..len].clone_from_slice(&rng_bytes[..len]); res } + + /// Decode from long header format + pub(crate) fn decode_long(buf: &mut impl Buf) -> Option { + let len = buf.get::().ok()? as usize; + if len > MAX_CID_SIZE || buf.remaining() < len { + return None; + } + let cid = ConnectionId::new(&buf.bytes()[..len]); + buf.advance(len); + Some(cid) + } + + /// Encode in long header format + pub(crate) fn encode_long(&self, buf: &mut impl BufMut) { + buf.put_u8(self.len() as u8); + buf.put_slice(self); + } } impl ::std::ops::Deref for ConnectionId {