From aca2baf4580ef13d8c28354a5f2020f947c6f713 Mon Sep 17 00:00:00 2001
From: Stefano Garzarella <sgarzare@redhat.com>
Date: Tue, 16 Jun 2020 12:15:23 +0200
Subject: [PATCH] vsock: fixed flow control regression

This patch has been cherry-picked from the Firecracker tree. The
reference commit is 2da612a9cdce85c91fb54ab22d950ec6ccc93b27.

Fixed a bug introduced by a271d08f0b1ba0ee82761cd49244b6a8017bcede,
whereby the flow control accouting would be off by a few bytes, for
host-initiated connections.

The connection ack message ("OK <port_num><CR>") was accounted for as
data sent by the guest, so its length was substracted from the total
amount of data the guest was allowed to send.

This commit changes the way this ack message is sent, so that it
bypasses flow control accouting.

Signed-off-by: Dan Horobeanu <dhr@amazon.com>
Signed-off-by: Gabriel Ionescu <gbi@amazon.com>
Signed-off-by: Stefano Garzarella <sgarzare@redhat.com>
---
 vm-virtio/src/vsock/csm/connection.rs | 24 ++++++++++++++++++------
 vm-virtio/src/vsock/unix/muxer.rs     | 15 ++++++++++++---
 2 files changed, 30 insertions(+), 9 deletions(-)

diff --git a/vm-virtio/src/vsock/csm/connection.rs b/vm-virtio/src/vsock/csm/connection.rs
index 4735b4b27..bbcc06e4d 100644
--- a/vm-virtio/src/vsock/csm/connection.rs
+++ b/vm-virtio/src/vsock/csm/connection.rs
@@ -571,12 +571,29 @@ where
         self.pending_rx.insert(PendingRx::Rst);
     }
 
+    /// Return the connections state.
+    ///
+    pub fn state(&self) -> ConnState {
+        self.state
+    }
+
+    /// Send some raw, untracked, data straight to the underlying connected stream.
+    /// Returns: number of bytes written, or the error describing the write failure.
+    ///
+    /// Warning: this will bypass the connection state machine and write directly to the
+    /// underlying stream. No account of this write is kept, which includes bypassing
+    /// vsock flow control.
+    ///
+    pub fn send_bytes_raw(&mut self, buf: &[u8]) -> Result<usize> {
+        self.stream.write(buf).map_err(Error::StreamWrite)
+    }
+
     /// Send some raw data (a byte-slice) to the host stream.
     ///
     /// Raw data can either be sent straight to the host stream, or to our TX buffer, if the
     /// former fails.
     ///
-    pub fn send_bytes(&mut self, buf: &[u8]) -> Result<()> {
+    fn send_bytes(&mut self, buf: &[u8]) -> Result<()> {
         // If there is data in the TX buffer, that means we're already registered for EPOLLOUT
         // events on the underlying stream. Therefore, there's no point in attempting a write
         // at this point. `self.notify()` will get called when EPOLLOUT arrives, and it will
@@ -611,11 +628,6 @@ where
         Ok(())
     }
 
-    /// Return the connections state.
-    pub fn state(&self) -> ConnState {
-        self.state
-    }
-
     /// Check if the credit information the peer has last received from us is outdated.
     ///
     fn peer_needs_credit_update(&self) -> bool {
diff --git a/vm-virtio/src/vsock/unix/muxer.rs b/vm-virtio/src/vsock/unix/muxer.rs
index b2f57c285..e1aec2d8c 100644
--- a/vm-virtio/src/vsock/unix/muxer.rs
+++ b/vm-virtio/src/vsock/unix/muxer.rs
@@ -682,11 +682,20 @@ impl VsockMuxer {
             // If this is a host-initiated connection that has just become established, we'll have
             // to send an ack message to the host end.
             if prev_state == ConnState::LocalInit && conn.state() == ConnState::Established {
-                conn.send_bytes(format!("OK {}\n", key.local_port).as_bytes())
-                    .unwrap_or_else(|err| {
+                let msg = format!("OK {}\n", key.local_port);
+                match conn.send_bytes_raw(msg.as_bytes()) {
+                    Ok(written) if written == msg.len() => (),
+                    Ok(_) => {
+                        // If we can't write a dozen bytes to a pristine connection something
+                        // must be really wrong. Killing it.
+                        conn.kill();
+                        warn!("vsock: unable to fully write connection ack msg.");
+                    }
+                    Err(err) => {
                         conn.kill();
                         warn!("vsock: unable to ack host connection: {:?}", err);
-                    });
+                    }
+                };
             }
 
             // If the connection wasn't previously scheduled for RX, add it to our RX queue.