diff --git a/src/lib.rs b/src/lib.rs index 2f429f9..261f0c4 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -653,8 +653,10 @@ impl ZK { return Ok(res.payload.into_owned()); } - // Parse size: prioritize offset 1 (standard) if length >= 5 - let size = if res.payload.len() >= 5 { + // Parse size: prioritize offset 1 (standard) if length >= 5. + // Fallback for offset misalignment: if size parsed at offset 1 exceeds MAX_RESPONSE_SIZE, + // and payload length >= 4, check if offset 0 size is valid (<= MAX_RESPONSE_SIZE). + let mut size = if res.payload.len() >= 5 { byteorder::LittleEndian::read_u32(&res.payload[1..5]) as usize } else if res.payload.len() >= 4 { // Fallback for short ACK_OK payloads @@ -666,6 +668,18 @@ impl ZK { ))); }; + if res.payload.len() >= 5 && size > MAX_RESPONSE_SIZE { + let alt_size = byteorder::LittleEndian::read_u32(&res.payload[0..4]) as usize; + if alt_size <= MAX_RESPONSE_SIZE { + log::warn!( + "Oversized size parsed at offset 1 ({}), falling back to offset 0 ({})", + size, + alt_size + ); + size = alt_size; + } + } + if size > MAX_RESPONSE_SIZE { return Err(ZKError::InvalidData(format!( "Buffered response size {} exceeds maximum {}", diff --git a/tests/misalignment_tests.rs b/tests/misalignment_tests.rs index eac2f6d..7f192d4 100644 --- a/tests/misalignment_tests.rs +++ b/tests/misalignment_tests.rs @@ -245,3 +245,151 @@ fn test_receive_chunk_misalignment_tcp() { server_handle.join().unwrap(); } + +#[test] +fn test_prepare_buffer_offset_misalignment_tcp() { + let _ = env_logger::builder().is_test(true).try_init(); + let listener = TcpListener::bind("127.0.0.1:0").expect("Failed to bind mock server"); + let addr = listener.local_addr().unwrap(); + let port = addr.port(); + + let server_handle = thread::spawn(move || { + let (mut stream, _) = listener.accept().expect("Failed to accept connection"); + let session_id = 1234; + + // 1. Handle Connect + let mut header = [0u8; 8]; + stream.read_exact(&mut header).unwrap(); + let (length, _) = TCPWrapper::decode_header(&header).unwrap(); + let mut body = vec![0u8; length]; + stream.read_exact(&mut body).unwrap(); + let packet = ZKPacket::from_bytes_owned(body).unwrap(); + assert_eq!(packet.command, CMD_CONNECT); + + let res = ZKPacket::new(CMD_ACK_OK, session_id, packet.reply_id, vec![]); + stream + .write_all(&TCPWrapper::wrap(&res.to_bytes())) + .unwrap(); + + // 2. Handle CMD_GET_FREE_SIZES (called by read_sizes) + stream.read_exact(&mut header).unwrap(); + let (length, _) = TCPWrapper::decode_header(&header).unwrap(); + let mut body = vec![0u8; length]; + stream.read_exact(&mut body).unwrap(); + let packet = ZKPacket::from_bytes_owned(body).unwrap(); + assert_eq!(packet.command, CMD_GET_FREE_SIZES); + + let mut sizes_payload = vec![0u8; 92]; + LittleEndian::write_i32(&mut sizes_payload[16..20], 1); + let res = ZKPacket::new(CMD_ACK_OK, session_id, packet.reply_id, sizes_payload); + stream + .write_all(&TCPWrapper::wrap(&res.to_bytes())) + .unwrap(); + + // 2.5 Handle TZAdj Sync (from read_sizes) + stream.read_exact(&mut header).unwrap(); + let (length, _) = TCPWrapper::decode_header(&header).unwrap(); + let mut body = vec![0u8; length]; + stream.read_exact(&mut body).unwrap(); + let packet = ZKPacket::from_bytes_owned(body).unwrap(); + assert_eq!(packet.command, CMD_OPTIONS_RRQ); + let res = ZKPacket::new( + CMD_ACK_OK, + session_id, + packet.reply_id, + b"TZAdj=7\0".to_vec(), + ); + stream + .write_all(&TCPWrapper::wrap(&res.to_bytes())) + .unwrap(); + + // 3. Handle CMD_USERTEMP_RRQ (from get_users -> read_with_buffer) + stream.read_exact(&mut header).unwrap(); + let (length, _) = TCPWrapper::decode_header(&header).unwrap(); + let mut body = vec![0u8; length]; + stream.read_exact(&mut body).unwrap(); + let packet = ZKPacket::from_bytes_owned(body).unwrap(); + assert_eq!(packet.command, _CMD_PREPARE_BUFFER); + + // Send CMD_PREPARE_DATA with size 28 (one user) + trailing garbage (0xAB) + // This will simulate the payload offset misalignment issue. + let mut prepare_payload = Vec::new(); + prepare_payload.write_u32::(28).unwrap(); + prepare_payload.push(0xAB); // 5th byte: trailing garbage shifts offset 1 to a huge number + + let res = ZKPacket::new( + CMD_PREPARE_DATA, + session_id, + packet.reply_id, + prepare_payload, + ); + stream + .write_all(&TCPWrapper::wrap(&res.to_bytes())) + .unwrap(); + + // 4. Handle _CMD_READ_BUFFER (from read_chunk) + stream.read_exact(&mut header).unwrap(); + let (length, _) = TCPWrapper::decode_header(&header).unwrap(); + let mut body = vec![0u8; length]; + stream.read_exact(&mut body).unwrap(); + let packet = ZKPacket::from_bytes_owned(body).unwrap(); + assert_eq!(packet.command, _CMD_READ_BUFFER); + + // Send CMD_PREPARE_DATA for the chunk + let mut chunk_prepare_payload = Vec::new(); + chunk_prepare_payload.write_u32::(32).unwrap(); + let res = ZKPacket::new( + CMD_PREPARE_DATA, + session_id, + packet.reply_id, + chunk_prepare_payload, + ); + stream + .write_all(&TCPWrapper::wrap(&res.to_bytes())) + .unwrap(); + + // 5. Client enters receive_chunk loop + // User struct: 28 bytes + let mut correct_data = Vec::new(); + correct_data.write_u32::(28).unwrap(); // size + correct_data.write_u16::(1).unwrap(); // uid + correct_data.write_u8(0).unwrap(); // priv + correct_data.extend_from_slice(b"12345"); // pass + correct_data.extend_from_slice(b"TestUser"); // name + correct_data.write_u32::(0).unwrap(); // card + correct_data.write_u8(0).unwrap(); // pad + correct_data.write_u8(1).unwrap(); // group + correct_data.write_u16::(0).unwrap(); // tz + correct_data.write_u32::(100).unwrap(); // user_id + + let correct_res = ZKPacket::new(CMD_DATA, session_id, packet.reply_id, correct_data); + stream + .write_all(&TCPWrapper::wrap(&correct_res.to_bytes())) + .unwrap(); + stream.flush().unwrap(); + + // 6. Handle CMD_FREE_DATA + stream.read_exact(&mut header).unwrap(); + let (length, _) = TCPWrapper::decode_header(&header).unwrap(); + let mut body = vec![0u8; length]; + stream.read_exact(&mut body).unwrap(); + let packet = ZKPacket::from_bytes_owned(body).unwrap(); + assert_eq!(packet.command, CMD_FREE_DATA); + + let res = ZKPacket::new(CMD_ACK_OK, session_id, packet.reply_id, vec![]); + stream + .write_all(&TCPWrapper::wrap(&res.to_bytes())) + .unwrap(); + stream.flush().unwrap(); + }); + + let mut zk = ZK::new("127.0.0.1", port); + zk.connect(ZKProtocol::TCP).unwrap(); + + let users = zk.get_users().unwrap(); + assert_eq!(users.len(), 1); + assert_eq!(users[0].name, "TestUser"); + assert_eq!(users[0].user_id, "100"); + + server_handle.join().unwrap(); +}