Merge pull request #590 from Enet4/bug/ul/get_client_pdu-mitigation

Fix issues around `dicom_ul::get_client_pdu`
Enet4 · Nov 2, 2024 · 6c44c43 · 6c44c43
2 parents f2fefa7 + 4ef8046
commit 6c44c43
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 16 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/scpproxy/Cargo.toml b/scpproxy/Cargo.toml
@@ -17,3 +17,4 @@ dicom-dictionary-std = { path = "../dictionary-std/", version = "0.7.0" }
 snafu = "0.8"
 tracing = "0.1.34"
 tracing-subscriber = "0.3.11"
+bytes = "1.6"
diff --git a/scpproxy/src/main.rs b/scpproxy/src/main.rs
@@ -1,7 +1,8 @@
+use bytes::BytesMut;
 use clap::{crate_version, value_parser, Arg, ArgAction, Command};
+use dicom_ul::association::client::get_client_pdu;
 use dicom_ul::pdu::writer::write_pdu;
 use dicom_ul::pdu::Pdu;
-use dicom_ul::association::client::get_client_pdu;
 use snafu::{Backtrace, OptionExt, Report, ResultExt, Snafu, Whatever};
 use std::io::Write;
 use std::net::{Shutdown, TcpListener, TcpStream};
@@ -93,10 +94,11 @@ fn run(
 
             {
                 let mut reader = scu_stream.try_clone().context(CloneSocketSnafu)?;
+                let mut buf = BytesMut::with_capacity(max_pdu_length as usize);
                 let message_tx = message_tx.clone();
                 scu_reader_thread = thread::spawn(move || {
                     loop {
-                        match get_client_pdu(&mut reader, max_pdu_length, strict) {
+                        match get_client_pdu(&mut reader, &mut buf, max_pdu_length, strict) {
                             Ok(pdu) => {
                                 message_tx
                                     .send(ThreadMessage::SendPdu {
@@ -105,7 +107,7 @@ fn run(
                                     })
                                     .context(SendMessageSnafu)?;
                             }
-                            Err(dicom_ul::association::client::Error::ReceiveResponse{ .. }) => {
+                            Err(dicom_ul::association::client::Error::ConnectionClosed) => {
                                 message_tx
                                     .send(ThreadMessage::Shutdown {
                                         initiator: ProviderType::Scu,
@@ -131,9 +133,10 @@ fn run(
 
             {
                 let mut reader = scp_stream.try_clone().context(CloneSocketSnafu)?;
+                let mut buf = BytesMut::with_capacity(max_pdu_length as usize);
                 scp_reader_thread = thread::spawn(move || {
                     loop {
-                        match get_client_pdu(&mut reader, max_pdu_length, strict) {
+                        match get_client_pdu(&mut reader, &mut buf, max_pdu_length, strict) {
                             Ok(pdu) => {
                                 message_tx
                                     .send(ThreadMessage::SendPdu {
@@ -142,7 +145,7 @@ fn run(
                                     })
                                     .context(SendMessageSnafu)?;
                             }
-                            Err(dicom_ul::association::client::Error::ReceiveResponse{ .. }) => {
+                            Err(dicom_ul::association::client::Error::ConnectionClosed) => {
                                 message_tx
                                     .send(ThreadMessage::Shutdown {
                                         initiator: ProviderType::Scp,

diff --git a/ul/src/association/client.rs b/ul/src/association/client.rs
@@ -137,15 +137,19 @@ pub enum Error {
 
 pub type Result<T, E = Error> = std::result::Result<T, E>;
 
-/// Helper function to get a PDU from a reader
-pub fn get_client_pdu<R: Read>(reader: &mut R, max_pdu_length: u32, strict: bool) -> Result<Pdu> {
-    // Receive response
-
-    let mut read_buffer = BytesMut::with_capacity(MAXIMUM_PDU_SIZE as usize);
+/// Helper function to get a PDU from a reader.
+///
+/// Chunks of data are read into `read_buffer`,
+/// which should be passed in subsequent calls
+/// to receive more PDUs from the same stream.
+pub fn get_client_pdu<R>(reader: &mut R, read_buffer: &mut BytesMut, max_pdu_length: u32, strict: bool) -> Result<Pdu>
+where
+    R: Read,
+{
     let mut reader = BufReader::new(reader);
-
     let msg = loop {
         let mut buf = Cursor::new(&read_buffer[..]);
+        // try to read a PDU according to what's in the buffer
         match read_pdu(&mut buf, max_pdu_length, strict).context(ReceiveResponseSnafu)? {
             Some(pdu) => {
                 read_buffer.advance(buf.position() as usize);
@@ -160,11 +164,11 @@ pub fn get_client_pdu<R: Read>(reader: &mut R, max_pdu_length: u32, strict: bool
         let recv = reader
             .fill_buf()
             .context(ReadPduSnafu)
-            .context(ReceiveSnafu)?
-            .to_vec();
-        reader.consume(recv.len());
+            .context(ReceiveSnafu)?;
+        let bytes_read = recv.len();
         read_buffer.extend_from_slice(&recv);
-        ensure!(!recv.is_empty(), ConnectionClosedSnafu);
+        reader.consume(bytes_read);
+        ensure!(bytes_read != 0, ConnectionClosedSnafu);
     };
     Ok(msg)
 }
@@ -678,7 +682,14 @@ impl<'a> ClientAssociationOptions<'a> {
         socket.write_all(&buffer).context(WireSendSnafu)?;
         buffer.clear();
 
-        let msg = get_client_pdu(&mut socket, MAXIMUM_PDU_SIZE, self.strict)?;
+        // !!!(#589) Soundness issue: if the SCP sends more PDUs in quick succession,
+        // more data may live in `buf` which may be lost,
+        // corrupting the PDU reader stream.
+        let mut buf = BytesMut::with_capacity(MAXIMUM_PDU_SIZE as usize);
+        let msg = get_client_pdu(&mut socket, &mut buf, MAXIMUM_PDU_SIZE, self.strict)?;
+        if !buf.is_empty() {
+            tracing::warn!("Received more data than expected in the first PDU, further issues may arise");
+        }
 
         match msg {
             Pdu::AssociationAC(AssociationAC {