From 58b902d036ea8fc1bda7061b84366c5b924b8d34 Mon Sep 17 00:00:00 2001
From: Praveen K Paladugu <prapal@linux.microsoft.com>
Date: Fri, 9 Sep 2022 03:31:56 +0000
Subject: [PATCH] tpm: Add socket module

Add SocketDev struct. Methods in SocketDev will be used to read & write
to Ctrl channel created by swtpm.

Signed-off-by: Praveen K Paladugu <prapal@linux.microsoft.com>
Co-authored-by: Sean Yoo <t-seanyoo@microsoft.com>
---
 Cargo.lock        |   2 +
 tpm/Cargo.toml    |   2 +
 tpm/src/lib.rs    |   5 ++
 tpm/src/socket.rs | 144 ++++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 153 insertions(+)
 create mode 100644 tpm/src/socket.rs

diff --git a/Cargo.lock b/Cargo.lock
index fa2f808c2..6da228b8b 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1166,8 +1166,10 @@ version = "0.1.0"
 dependencies = [
  "anyhow",
  "byteorder",
+ "libc",
  "log",
  "thiserror",
+ "vmm-sys-util",
 ]
 
 [[package]]
diff --git a/tpm/Cargo.toml b/tpm/Cargo.toml
index 6aa00ee24..721535348 100644
--- a/tpm/Cargo.toml
+++ b/tpm/Cargo.toml
@@ -8,5 +8,7 @@ version = "0.1.0"
 [dependencies]
 anyhow = "1.0.56"
 byteorder = "1.4.3"
+libc = "0.2.133"
 log = "0.4.14"
 thiserror = "1.0.30"
+vmm-sys-util = "0.10.0"
diff --git a/tpm/src/lib.rs b/tpm/src/lib.rs
index 3f550d05f..1a05c7d6b 100644
--- a/tpm/src/lib.rs
+++ b/tpm/src/lib.rs
@@ -3,6 +3,11 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
+#[macro_use]
+extern crate log;
+
+pub mod socket;
+
 use anyhow::anyhow;
 use byteorder::{BigEndian, ReadBytesExt};
 use std::convert::TryInto;
diff --git a/tpm/src/socket.rs b/tpm/src/socket.rs
new file mode 100644
index 000000000..60db930ae
--- /dev/null
+++ b/tpm/src/socket.rs
@@ -0,0 +1,144 @@
+// Copyright © 2022, Microsoft Corporation
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+use crate::TPM_CRB_BUFFER_MAX;
+use anyhow::anyhow;
+use std::io::Read;
+use std::os::unix::io::{AsRawFd, RawFd};
+use std::os::unix::net::UnixStream;
+use thiserror::Error;
+use vmm_sys_util::sock_ctrl_msg::ScmSocket;
+
+#[derive(Error, Debug)]
+pub enum Error {
+    #[error("Cannot connect to tpm Socket: {0}")]
+    ConnectToSocket(#[source] anyhow::Error),
+    #[error("Failed to read from socket: {0}")]
+    ReadFromSocket(#[source] anyhow::Error),
+    #[error("Failed to write to socket: {0}")]
+    WriteToSocket(#[source] anyhow::Error),
+}
+type Result<T> = anyhow::Result<T, Error>;
+
+#[derive(PartialEq)]
+enum SocketDevState {
+    Disconnected,
+    Connecting,
+    Connected,
+}
+
+pub struct SocketDev {
+    state: SocketDevState,
+    stream: Option<UnixStream>,
+    // Fd sent to swtpm process for Data Channel
+    write_msgfd: RawFd,
+    // Data Channel used by Cloud-Hypervisor
+    data_fd: RawFd,
+    // Control Channel used by Cloud-Hypervisor
+    control_fd: RawFd,
+}
+
+impl Default for SocketDev {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl SocketDev {
+    pub fn new() -> Self {
+        Self {
+            state: SocketDevState::Disconnected,
+            stream: None,
+            write_msgfd: -1,
+            control_fd: -1,
+            data_fd: -1,
+        }
+    }
+
+    pub fn init(&mut self, path: String) -> Result<()> {
+        self.connect(&path)?;
+        Ok(())
+    }
+
+    pub fn connect(&mut self, socket_path: &str) -> Result<()> {
+        self.state = SocketDevState::Connecting;
+
+        let s = UnixStream::connect(socket_path).map_err(|e| {
+            Error::ConnectToSocket(anyhow!("Failed to connect to tpm Socket. Error: {:?}", e))
+        })?;
+        self.control_fd = s.as_raw_fd();
+        self.stream = Some(s);
+        self.state = SocketDevState::Connected;
+        debug!("Connected to tpm socket path : {:?}", socket_path);
+        Ok(())
+    }
+
+    pub fn set_datafd(&mut self, fd: RawFd) {
+        self.data_fd = fd;
+    }
+
+    pub fn set_msgfd(&mut self, fd: RawFd) {
+        self.write_msgfd = fd;
+    }
+
+    pub fn send_full(&self, buf: &[u8]) -> Result<usize> {
+        let write_fd = self.write_msgfd;
+
+        let size = self
+            .stream
+            .as_ref()
+            .unwrap()
+            .send_with_fd(buf, write_fd)
+            .map_err(|e| {
+                Error::WriteToSocket(anyhow!("Failed to write to Socket. Error: {:?}", e))
+            })?;
+
+        Ok(size)
+    }
+
+    pub fn write(&mut self, buf: &[u8]) -> Result<usize> {
+        if self.stream.is_none() {
+            return Err(Error::WriteToSocket(anyhow!(
+                "Stream for tpm socket was not initialized"
+            )));
+        }
+
+        if matches!(self.state, SocketDevState::Connected) {
+            let ret = self.send_full(buf)?;
+            // swtpm will receive data Fd after a successful send
+            // Reset cached write_msgfd after a successful send
+            // Ideally, write_msgfd is reset after first Ctrl Command
+            if ret > 0 && self.write_msgfd != 0 {
+                self.write_msgfd = 0;
+            }
+            Ok(ret)
+        } else {
+            Err(Error::WriteToSocket(anyhow!(
+                "TPM Socket was not in Connected State"
+            )))
+        }
+    }
+
+    pub fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
+        let newbuf: &mut [u8] = &mut [0; TPM_CRB_BUFFER_MAX];
+
+        if self.stream.is_none() {
+            return Err(Error::ReadFromSocket(anyhow!(
+                "Stream for tpm socket was not initialized"
+            )));
+        }
+        let mut socket = self.stream.as_ref().unwrap();
+        let size: usize = socket.read(newbuf).map_err(|e| {
+            Error::ReadFromSocket(anyhow!("Failed to read from socket. Error Code {:?}", e))
+        })?;
+        if buf.len() < size {
+            return Err(Error::ReadFromSocket(anyhow!(
+                "Input buffer is of insufficient size"
+            )));
+        }
+        buf[0..size].clone_from_slice(&newbuf[0..size]);
+        Ok(size)
+    }
+}