瀏覽代碼

Merge pull request #834 from hunhoffe/bug/peek_slice

Match the behavior of peek_slice to recv_slice
Dario Nieuwenhuis 1 年之前
父節點
當前提交
28a5dd17ab
共有 1 個文件被更改,包括 57 次插入5 次删除
  1. 57 5
      src/socket/tcp.rs

+ 57 - 5
src/socket/tcp.rs

@@ -1173,10 +1173,7 @@ impl<'a> Socket<'a> {
     ///
     /// This function otherwise behaves identically to [recv_slice](#method.recv_slice).
     pub fn peek_slice(&mut self, data: &mut [u8]) -> Result<usize, RecvError> {
-        let buffer = self.peek(data.len())?;
-        let data = &mut data[..buffer.len()];
-        data.copy_from_slice(buffer);
-        Ok(buffer.len())
+        Ok(self.rx_buffer.read_allocated(0, data))
     }
 
     /// Return the amount of octets queued in the transmit buffer.
@@ -1868,7 +1865,10 @@ impl<'a> Socket<'a> {
         let assembler_was_empty = self.assembler.is_empty();
 
         // Try adding payload octets to the assembler.
-        let Ok(contig_len) = self.assembler.add_then_remove_front(payload_offset, payload_len) else {
+        let Ok(contig_len) = self
+            .assembler
+            .add_then_remove_front(payload_offset, payload_len)
+        else {
             net_debug!(
                 "assembler: too many holes to add {} octets at offset {}",
                 payload_len,
@@ -3667,6 +3667,58 @@ mod test {
         assert_eq!(s.rx_buffer.dequeue_many(6), &b"abcdef"[..]);
     }
 
+    #[test]
+    fn test_peek_slice() {
+        const BUF_SIZE: usize = 10;
+
+        let send_buf = b"0123456";
+
+        let mut s = socket_established_with_buffer_sizes(BUF_SIZE, BUF_SIZE);
+
+        // Populate the recv buffer
+        send!(
+            s,
+            TcpRepr {
+                seq_number: REMOTE_SEQ + 1,
+                ack_number: Some(LOCAL_SEQ + 1),
+                payload: &send_buf[..],
+                ..SEND_TEMPL
+            }
+        );
+
+        // Peek into the recv buffer
+        let mut peeked_buf = [0u8; BUF_SIZE];
+        let actually_peeked = s.peek_slice(&mut peeked_buf[..]).unwrap();
+        let mut recv_buf = [0u8; BUF_SIZE];
+        let actually_recvd = s.recv_slice(&mut recv_buf[..]).unwrap();
+        assert_eq!(
+            &mut peeked_buf[..actually_peeked],
+            &mut recv_buf[..actually_recvd]
+        );
+    }
+
+    #[test]
+    fn test_peek_slice_buffer_wrap() {
+        const BUF_SIZE: usize = 10;
+
+        let send_buf = b"0123456789";
+
+        let mut s = socket_established_with_buffer_sizes(BUF_SIZE, BUF_SIZE);
+
+        let _ = s.rx_buffer.enqueue_slice(&send_buf[..8]);
+        let _ = s.rx_buffer.dequeue_many(6);
+        let _ = s.rx_buffer.enqueue_slice(&send_buf[..5]);
+
+        let mut peeked_buf = [0u8; BUF_SIZE];
+        let actually_peeked = s.peek_slice(&mut peeked_buf[..]).unwrap();
+        let mut recv_buf = [0u8; BUF_SIZE];
+        let actually_recvd = s.recv_slice(&mut recv_buf[..]).unwrap();
+        assert_eq!(
+            &mut peeked_buf[..actually_peeked],
+            &mut recv_buf[..actually_recvd]
+        );
+    }
+
     fn setup_rfc2018_cases() -> (TestSocket, Vec<u8>) {
         // This is a utility function used by the tests for RFC 2018 cases. It configures a socket
         // in a particular way suitable for those cases.