From e1e8e9591361a36e650f385a0da06e892d69d73c Mon Sep 17 00:00:00 2001
From: Ruoyu Ying <ruoyu.ying@intel.com>
Date: Mon, 26 Feb 2024 13:33:13 +0800
Subject: [PATCH] vmsdk: add vsock support on TDX quote fetching

Signed-off-by: Ruoyu Ying <ruoyu.ying@intel.com>
---
 src/python/cctrusted_vm/cvm.py | 68 ++++++++++++++++++++++++++++++----
 1 file changed, 61 insertions(+), 7 deletions(-)

diff --git a/src/python/cctrusted_vm/cvm.py b/src/python/cctrusted_vm/cvm.py
index 9008e8b..181baa3 100644
--- a/src/python/cctrusted_vm/cvm.py
+++ b/src/python/cctrusted_vm/cvm.py
@@ -12,6 +12,7 @@
 import logging
 import struct
 import fcntl
+import socket
 from abc import abstractmethod
 from cctrusted_base.api import CCTrustedApi
 from cctrusted_base.imr import TcgIMR
@@ -19,7 +20,7 @@
 from cctrusted_base.tcg import TcgAlgorithmRegistry
 from cctrusted_base.tdx.common import TDX_VERSION_1_0, TDX_VERSION_1_5
 from cctrusted_base.tdx.rtmr import TdxRTMR
-from cctrusted_base.tdx.quote import TdxQuoteReq10, TdxQuoteReq15
+from cctrusted_base.tdx.quote import TdxQuoteReq10, TdxQuoteReq15, TdxQuote
 from cctrusted_base.tdx.report import TdxReportReq10, TdxReportReq15
 
 LOG = logging.getLogger(__name__)
@@ -216,6 +217,7 @@ class TdxVM(ConfidentialVM):
     ACPI_TABLE_FILE = "/sys/firmware/acpi/tables/CCEL"
     ACPI_TABLE_DATA_FILE = "/sys/firmware/acpi/tables/data/CCEL"
     IMA_DATA_FILE = "/sys/kernel/security/integrity/ima/ascii_runtime_measurements"
+    CFG_FILE_PATH = "/etc/tdx-attest.conf"
 
     def __init__(self):
         ConfidentialVM.__init__(self, CCTrustedApi.TYPE_CC_TDX)
@@ -391,6 +393,64 @@ def get_cc_report(self, nonce: bytearray, data: bytearray, extraArgs) -> CcRepor
         self.process_cc_report(input_data)
         report_bytes = self.tdreport.data
 
+        if self.version is TDX_VERSION_1_0:
+            quote_req = TdxQuoteReq10()
+        elif self.version is TDX_VERSION_1_5:
+            quote_req = TdxQuoteReq15()
+
+        tdvmcall_flag = False
+        # Use vsock to get TD Quote by default
+        with open(TdxVM.CFG_FILE_PATH, 'rb') as port_file:
+            port_info = port_file.read().rstrip().decode("utf-8")
+            if "port=" not in port_info:
+                LOG.info("No vsock port number specified.")
+                tdvmcall_flag = True
+            if not tdvmcall_flag:
+                port = int(port_info.partition("port=")[2])
+                if port <= 0 or port > 65535:
+                    LOG.error("Invalid vsock port number specified.")
+                    tdvmcall_flag = True
+
+        if not tdvmcall_flag:
+            with socket.socket(socket.AF_VSOCK, socket.SOCK_STREAM, 0) as sock:
+                sock.connect((socket.VMADDR_CID_HOST, port))
+
+                header_size = 4
+                # Generate p_blob_payload buffer
+                qgs_msg = quote_req.qgs_msg_quote_req(report_bytes)
+                msg_size = len(qgs_msg)
+
+                msg_size_bytes = [
+                    (msg_size >> 24) & 0xFF,
+                    (msg_size >> 16) & 0xFF,
+                    (msg_size >> 8) & 0xFF,
+                    msg_size & 0xFF
+                ]
+
+                p_blob_payload = bytearray(header_size)
+                for i, byte_value in enumerate(msg_size_bytes):
+                    p_blob_payload[i] = byte_value
+                p_blob_payload[header_size:header_size + msg_size] = qgs_msg[:msg_size]
+
+                nsent = sock.send(p_blob_payload)
+                LOG.debug("Sent %d bytes for Quote request.", nsent)
+
+                # Receive quote
+                header = sock.recv(header_size)
+                in_msg_size = 0
+                for i in range(header_size):
+                    in_msg_size = (in_msg_size << 8) + (header[i] & 0xFF)
+                qgs_resp = sock.recv(in_msg_size)
+                LOG.debug("Received %d bytes as Quote response", in_msg_size)
+
+                sock.close()
+            tdquote = quote_req.qgs_msg_quote_resp(qgs_resp)
+            return TdxQuote(tdquote)
+
+        # Fetch quote through tdvmcall
+        # pylint: disable=E1111
+        req_buf = quote_req.prepare_reqbuf(report_bytes)
+
         # Open TDX guest device node
         dev_path = self.DEVICE_NODE_PATH[self.version]
         try:
@@ -401,12 +461,6 @@ def get_cc_report(self, nonce: bytearray, data: bytearray, extraArgs) -> CcRepor
         LOG.debug("Successful open device node %s", dev_path)
 
         # Run ioctl command to get TD Quote
-        if self.version is TDX_VERSION_1_0:
-            quote_req = TdxQuoteReq10()
-        elif self.version is TDX_VERSION_1_5:
-            quote_req = TdxQuoteReq15()
-        # pylint: disable=E1111
-        req_buf = quote_req.prepare_reqbuf(report_bytes)
         try:
             fcntl.ioctl(tdx_dev, self.IOCTL_GET_QUOTE[self.version], req_buf)
         except OSError as e: