Merge "pvmfw/avb: Add Capability::UefiSupport" into main
diff --git a/build/debian/build.sh b/build/debian/build.sh
index 3d3820a..2177b17 100755
--- a/build/debian/build.sh
+++ b/build/debian/build.sh
@@ -37,6 +37,7 @@
 	apt update
 	DEBIAN_FRONTEND=noninteractive \
 	apt install --no-install-recommends --assume-yes \
+		binfmt-support \
 		ca-certificates \
 		debsums \
 		dosfstools \
@@ -49,10 +50,11 @@
 		python3-marshmallow \
 		python3-pytest \
 		python3-yaml \
+		qemu-system-arm \
+		qemu-user-static \
 		qemu-utils \
 		udev \
-		qemu-system-arm \
-		qemu-user-static
+
 
         sed -i s/losetup\ -f/losetup\ -P\ -f/g /usr/sbin/fai-diskimage
         sed -i 's/wget \$/wget -t 0 \$/g' /usr/share/debootstrap/functions
diff --git a/build/debian/fai_config/hooks/extrbase.BASE b/build/debian/fai_config/hooks/extrbase.BASE
deleted file mode 100755
index 05d1e96..0000000
--- a/build/debian/fai_config/hooks/extrbase.BASE
+++ /dev/null
@@ -1,6 +0,0 @@
-#!/bin/sh
-set -euE
-
-touch "${LOGDIR}/skip.extrbase"
-
-debootstrap --verbose --variant minbase --arch "$DEBOOTSTRAP_ARCH" "$SUITE" "$FAI_ROOT" "$DEBOOTSTRAP_MIRROR"
diff --git a/build/debian/fai_config/hooks/partition.ARM64 b/build/debian/fai_config/hooks/partition.ARM64
deleted file mode 100755
index b3b603b..0000000
--- a/build/debian/fai_config/hooks/partition.ARM64
+++ /dev/null
@@ -1,53 +0,0 @@
-#!/bin/sh
-set -eu
-touch $LOGDIR/skip.partition
-
-set -- $disklist
-device=/dev/$1
-
-wait_for_device() {
-  for s in $(seq 10); do
-    if [ -e "$1" ]; then
-      break
-    fi
-    sleep 1
-  done
-}
-
-sfdisk "$device" << EOF
-label: gpt
-unit: sectors
-
-# EFI system
-p15 : start=2048, size=260096, type="EFI System", uuid=${PARTUUID_ESP}
-# Linux
-p1 : start=262144, type="Linux root (ARM-64)", uuid=${PARTUUID_ROOT}
-EOF
-
-file=$(losetup -O BACK-FILE ${device} | tail -1)
-
-root_offset=$(parted -m ${device} unit B print | awk -F '[B:]' '/1:/{ print $2 }')
-root_size=$(  parted -m ${device} unit B print | awk -F '[B:]' '/1:/{ print $6 }')
-efi_offset=$( parted -m ${device} unit B print | awk -F '[B:]' '/15:/{ print $2 }')
-efi_size=$(   parted -m ${device} unit B print | awk -F '[B:]' '/15:/{ print $6 }')
-device_root=$(losetup -o ${root_offset} --sizelimit ${root_size} --show -f ${file})
-device_efi=$(losetup -o ${efi_offset} --sizelimit ${efi_size} --show -f ${file})
-rm -f ${device}p1
-rm -f ${device}p15
-ln -sf ${device_root} ${device}p1
-ln -sf ${device_efi} ${device}p15
-
-ls -al /dev/loop*
-losetup -a -l
-parted ${device} unit B print
-
-partprobe "$device"
-
-wait_for_device "$device_root"
-mkfs.ext4 -U "$FSUUID_ROOT" "$device_root"
-tune2fs -c 0 -i 0 "$device_root"
-
-wait_for_device "$device_efi"
-mkfs.vfat "$device_efi"
-
-parted ${device} unit B print
diff --git a/build/debian/forwarder_guest/Cargo.toml b/build/debian/forwarder_guest/Cargo.toml
new file mode 100644
index 0000000..e70dcd4
--- /dev/null
+++ b/build/debian/forwarder_guest/Cargo.toml
@@ -0,0 +1,11 @@
+[package]
+name = "forwarder_guest"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+clap = { version = "4.5.19", features = ["derive"] }
+forwarder = { path = "../../../libs/libforwarder" }
+poll_token_derive = "0.1.0"
+remain = "0.2.14"
+vmm-sys-util = "0.12.1"
diff --git a/build/debian/forwarder_guest/src/main.rs b/build/debian/forwarder_guest/src/main.rs
new file mode 100644
index 0000000..6ebd4ef
--- /dev/null
+++ b/build/debian/forwarder_guest/src/main.rs
@@ -0,0 +1,123 @@
+// Copyright 2024 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Copied from ChromiumOS with relicensing:
+// src/platform2/vm_tools/chunnel/src/bin/chunnel.rs
+
+//! Guest-side stream socket forwarder
+
+use std::fmt;
+use std::result;
+
+use clap::Parser;
+use forwarder::forwarder::{ForwarderError, ForwarderSession};
+use forwarder::stream::{StreamSocket, StreamSocketError};
+use poll_token_derive::PollToken;
+use vmm_sys_util::poll::{PollContext, PollToken};
+
+#[remain::sorted]
+#[derive(Debug)]
+enum Error {
+    ConnectSocket(StreamSocketError),
+    Forward(ForwarderError),
+    PollContextAdd(vmm_sys_util::errno::Error),
+    PollContextDelete(vmm_sys_util::errno::Error),
+    PollContextNew(vmm_sys_util::errno::Error),
+    PollWait(vmm_sys_util::errno::Error),
+}
+
+type Result<T> = result::Result<T, Error>;
+
+impl fmt::Display for Error {
+    #[remain::check]
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        use self::Error::*;
+
+        #[remain::sorted]
+        match self {
+            ConnectSocket(e) => write!(f, "failed to connect socket: {}", e),
+            Forward(e) => write!(f, "failed to forward traffic: {}", e),
+            PollContextAdd(e) => write!(f, "failed to add fd to poll context: {}", e),
+            PollContextDelete(e) => write!(f, "failed to delete fd from poll context: {}", e),
+            PollContextNew(e) => write!(f, "failed to create poll context: {}", e),
+            PollWait(e) => write!(f, "failed to wait for poll: {}", e),
+        }
+    }
+}
+
+fn run_forwarder(local_stream: StreamSocket, remote_stream: StreamSocket) -> Result<()> {
+    #[derive(PollToken)]
+    enum Token {
+        LocalStreamReadable,
+        RemoteStreamReadable,
+    }
+    let poll_ctx: PollContext<Token> = PollContext::new().map_err(Error::PollContextNew)?;
+    poll_ctx.add(&local_stream, Token::LocalStreamReadable).map_err(Error::PollContextAdd)?;
+    poll_ctx.add(&remote_stream, Token::RemoteStreamReadable).map_err(Error::PollContextAdd)?;
+
+    let mut forwarder = ForwarderSession::new(local_stream, remote_stream);
+
+    loop {
+        let events = poll_ctx.wait().map_err(Error::PollWait)?;
+
+        for event in events.iter_readable() {
+            match event.token() {
+                Token::LocalStreamReadable => {
+                    let shutdown = forwarder.forward_from_local().map_err(Error::Forward)?;
+                    if shutdown {
+                        poll_ctx
+                            .delete(forwarder.local_stream())
+                            .map_err(Error::PollContextDelete)?;
+                    }
+                }
+                Token::RemoteStreamReadable => {
+                    let shutdown = forwarder.forward_from_remote().map_err(Error::Forward)?;
+                    if shutdown {
+                        poll_ctx
+                            .delete(forwarder.remote_stream())
+                            .map_err(Error::PollContextDelete)?;
+                    }
+                }
+            }
+        }
+        if forwarder.is_shut_down() {
+            return Ok(());
+        }
+    }
+}
+
+#[derive(Parser)]
+/// Flags for running command
+pub struct Args {
+    /// Local socket address
+    #[arg(long)]
+    #[arg(alias = "local")]
+    local_sockaddr: String,
+
+    /// Remote socket address
+    #[arg(long)]
+    #[arg(alias = "remote")]
+    remote_sockaddr: String,
+}
+
+// TODO(b/370897694): Support forwarding for datagram socket
+fn main() -> Result<()> {
+    let args = Args::parse();
+
+    let local_stream = StreamSocket::connect(&args.local_sockaddr).map_err(Error::ConnectSocket)?;
+    let remote_stream =
+        StreamSocket::connect(&args.remote_sockaddr).map_err(Error::ConnectSocket)?;
+
+    run_forwarder(local_stream, remote_stream)
+}
diff --git a/build/debian/kokoro/gcp_ubuntu_docker/build.sh b/build/debian/kokoro/gcp_ubuntu_docker/build.sh
index c5745d0..4598d1c 100644
--- a/build/debian/kokoro/gcp_ubuntu_docker/build.sh
+++ b/build/debian/kokoro/gcp_ubuntu_docker/build.sh
@@ -3,11 +3,7 @@
 set -e
 
 cd "${KOKORO_ARTIFACTS_DIR}/git/avf/build/debian/"
-
-# FAI needs it
-pyenv install 3.10
-pyenv global 3.10
-python --version
-
 sudo losetup -D
-sudo -E ./build.sh
+grep vmx /proc/cpuinfo || true
+sudo ./build.sh
+cp image.raw ${KOKORO_ARTIFACTS_DIR}
diff --git a/build/debian/kokoro/gcp_ubuntu_docker/continuous.cfg b/build/debian/kokoro/gcp_ubuntu_docker/continuous.cfg
index d92031e..111096d 100644
--- a/build/debian/kokoro/gcp_ubuntu_docker/continuous.cfg
+++ b/build/debian/kokoro/gcp_ubuntu_docker/continuous.cfg
@@ -5,3 +5,9 @@
 # Location of the bash script. Should have value <git_on_borg_scm.name>/<path_from_repository_root>.
 # git_on_borg_scm.name is specified in the job configuration (next section).
 build_file: "avf/build/debian/kokoro/gcp_ubuntu_docker/build.sh"
+
+action {
+  define_artifacts {
+    regex: "image.raw"
+  }
+}
diff --git a/build/debian/kokoro/gcp_ubuntu_docker/presubmit.cfg b/build/debian/kokoro/gcp_ubuntu_docker/hourly.cfg
similarity index 85%
rename from build/debian/kokoro/gcp_ubuntu_docker/presubmit.cfg
rename to build/debian/kokoro/gcp_ubuntu_docker/hourly.cfg
index d92031e..111096d 100644
--- a/build/debian/kokoro/gcp_ubuntu_docker/presubmit.cfg
+++ b/build/debian/kokoro/gcp_ubuntu_docker/hourly.cfg
@@ -5,3 +5,9 @@
 # Location of the bash script. Should have value <git_on_borg_scm.name>/<path_from_repository_root>.
 # git_on_borg_scm.name is specified in the job configuration (next section).
 build_file: "avf/build/debian/kokoro/gcp_ubuntu_docker/build.sh"
+
+action {
+  define_artifacts {
+    regex: "image.raw"
+  }
+}
diff --git a/libs/libforwarder/Android.bp b/libs/libforwarder/Android.bp
new file mode 100644
index 0000000..48307e7
--- /dev/null
+++ b/libs/libforwarder/Android.bp
@@ -0,0 +1,15 @@
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+rust_library {
+    name: "libforwarder",
+    crate_name: "forwarder",
+    edition: "2021",
+    srcs: ["src/lib.rs"],
+    rustlibs: [
+        "liblibc",
+        "libvsock",
+    ],
+    proc_macros: ["libremain"],
+}
diff --git a/libs/libforwarder/Cargo.toml b/libs/libforwarder/Cargo.toml
new file mode 100644
index 0000000..9f3f341
--- /dev/null
+++ b/libs/libforwarder/Cargo.toml
@@ -0,0 +1,9 @@
+[package]
+name = "forwarder"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+libc = "0.2.159"
+remain = "0.2.14"
+vsock = "0.5.1"
diff --git a/libs/libforwarder/src/forwarder.rs b/libs/libforwarder/src/forwarder.rs
new file mode 100644
index 0000000..3600ab2
--- /dev/null
+++ b/libs/libforwarder/src/forwarder.rs
@@ -0,0 +1,170 @@
+// Copyright 2024 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Copied from ChromiumOS with relicensing:
+// src/platform2/vm_tools/chunnel/src/forwarder.rs
+
+//! This module contains forwarding mechanism between stream sockets.
+
+use std::fmt;
+use std::io::{self, Read, Write};
+use std::result;
+
+use crate::stream::StreamSocket;
+
+// This was picked arbitrarily. crosvm doesn't yet use VIRTIO_NET_F_MTU, so there's no reason to
+// opt for massive 65535 byte frames.
+const MAX_FRAME_SIZE: usize = 8192;
+
+/// Errors that can be encountered by a ForwarderSession.
+#[remain::sorted]
+#[derive(Debug)]
+pub enum ForwarderError {
+    /// An io::Error was encountered while reading from a stream.
+    ReadFromStream(io::Error),
+    /// An io::Error was encountered while shutting down writes on a stream.
+    ShutDownStream(io::Error),
+    /// An io::Error was encountered while writing to a stream.
+    WriteToStream(io::Error),
+}
+
+type Result<T> = result::Result<T, ForwarderError>;
+
+impl fmt::Display for ForwarderError {
+    #[remain::check]
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        use self::ForwarderError::*;
+
+        #[remain::sorted]
+        match self {
+            ReadFromStream(e) => write!(f, "failed to read from stream: {}", e),
+            ShutDownStream(e) => write!(f, "failed to shut down stream: {}", e),
+            WriteToStream(e) => write!(f, "failed to write to stream: {}", e),
+        }
+    }
+}
+
+/// A ForwarderSession owns stream sockets that it forwards traffic between.
+pub struct ForwarderSession {
+    local: StreamSocket,
+    remote: StreamSocket,
+}
+
+fn forward(from_stream: &mut StreamSocket, to_stream: &mut StreamSocket) -> Result<bool> {
+    let mut buf = [0u8; MAX_FRAME_SIZE];
+
+    let count = from_stream.read(&mut buf).map_err(ForwarderError::ReadFromStream)?;
+    if count == 0 {
+        to_stream.shut_down_write().map_err(ForwarderError::ShutDownStream)?;
+        return Ok(true);
+    }
+
+    to_stream.write_all(&buf[..count]).map_err(ForwarderError::WriteToStream)?;
+    Ok(false)
+}
+
+impl ForwarderSession {
+    /// Creates a forwarder session from a local and remote stream socket.
+    pub fn new(local: StreamSocket, remote: StreamSocket) -> Self {
+        ForwarderSession { local, remote }
+    }
+
+    /// Forwards traffic from the local socket to the remote socket.
+    /// Returns true if the local socket has reached EOF and the
+    /// remote socket has been shut down for further writes.
+    pub fn forward_from_local(&mut self) -> Result<bool> {
+        forward(&mut self.local, &mut self.remote)
+    }
+
+    /// Forwards traffic from the remote socket to the local socket.
+    /// Returns true if the remote socket has reached EOF and the
+    /// local socket has been shut down for further writes.
+    pub fn forward_from_remote(&mut self) -> Result<bool> {
+        forward(&mut self.remote, &mut self.local)
+    }
+
+    /// Returns a reference to the local stream socket.
+    pub fn local_stream(&self) -> &StreamSocket {
+        &self.local
+    }
+
+    /// Returns a reference to the remote stream socket.
+    pub fn remote_stream(&self) -> &StreamSocket {
+        &self.remote
+    }
+
+    /// Returns true if both sockets are completely shut down and the session can be dropped.
+    pub fn is_shut_down(&self) -> bool {
+        self.local.is_shut_down() && self.remote.is_shut_down()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::io::{Read, Write};
+    use std::net::Shutdown;
+    use std::os::unix::net::UnixStream;
+
+    #[test]
+    fn forward_unix() {
+        // Local streams.
+        let (mut london, folkestone) = UnixStream::pair().unwrap();
+        // Remote streams.
+        let (coquelles, mut paris) = UnixStream::pair().unwrap();
+
+        // Connect the local and remote sockets via the chunnel.
+        let mut forwarder = ForwarderSession::new(folkestone.into(), coquelles.into());
+
+        // Put some traffic in from London.
+        let greeting = b"hello";
+        london.write_all(greeting).unwrap();
+
+        // Expect forwarding from the local end not to have reached EOF.
+        assert!(!forwarder.forward_from_local().unwrap());
+        let mut salutation = [0u8; 8];
+        let count = paris.read(&mut salutation).unwrap();
+        assert_eq!(greeting.len(), count);
+        assert_eq!(greeting, &salutation[..count]);
+
+        // Shut the local socket down. The forwarder should detect this and perform a shutdown,
+        // which will manifest as an EOF when reading.
+        london.shutdown(Shutdown::Write).unwrap();
+        assert!(forwarder.forward_from_local().unwrap());
+        assert_eq!(paris.read(&mut salutation).unwrap(), 0);
+
+        // Don't consider the forwarder shut down until both ends are.
+        assert!(!forwarder.is_shut_down());
+
+        // Forward traffic from the remote end.
+        let salutation = b"bonjour";
+        paris.write_all(salutation).unwrap();
+
+        // Expect forwarding from the remote end not to have reached EOF.
+        assert!(!forwarder.forward_from_remote().unwrap());
+        let mut greeting = [0u8; 8];
+        let count = london.read(&mut greeting).unwrap();
+        assert_eq!(salutation.len(), count);
+        assert_eq!(salutation, &greeting[..count]);
+
+        // Shut the remote socket down. The forwarder should detect this and perform a shutdown,
+        // which will manifest as an EOF when reading.
+        paris.shutdown(Shutdown::Write).unwrap();
+        assert!(forwarder.forward_from_remote().unwrap());
+        assert_eq!(london.read(&mut greeting).unwrap(), 0);
+
+        // The forwarder should now be considered shut down.
+        assert!(forwarder.is_shut_down());
+    }
+}
diff --git a/libs/libforwarder/src/lib.rs b/libs/libforwarder/src/lib.rs
new file mode 100644
index 0000000..bcce689
--- /dev/null
+++ b/libs/libforwarder/src/lib.rs
@@ -0,0 +1,21 @@
+// Copyright 2024 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Copied from ChromiumOS with relicensing:
+// src/platform2/vm_tools/chunnel/src/lib.rs
+
+//! Library for stream socket forwarding.
+
+pub mod forwarder;
+pub mod stream;
diff --git a/libs/libforwarder/src/stream.rs b/libs/libforwarder/src/stream.rs
new file mode 100644
index 0000000..d8c7f51
--- /dev/null
+++ b/libs/libforwarder/src/stream.rs
@@ -0,0 +1,263 @@
+// Copyright 2024 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Copied from ChromiumOS with relicensing:
+// src/platform2/vm_tools/chunnel/src/stream.rs
+
+//! This module provides abstraction of various stream socket type.
+
+use std::fmt;
+use std::io;
+use std::net::TcpStream;
+use std::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd, RawFd};
+use std::os::unix::net::UnixStream;
+use std::result;
+
+use libc::{self, c_void, shutdown, EPIPE, SHUT_WR};
+use vsock::VsockAddr;
+use vsock::VsockStream;
+
+/// Parse a vsock SocketAddr from a string. vsock socket addresses are of the form
+/// "vsock:cid:port".
+pub fn parse_vsock_addr(addr: &str) -> result::Result<VsockAddr, io::Error> {
+    let components: Vec<&str> = addr.split(':').collect();
+    if components.len() != 3 || components[0] != "vsock" {
+        return Err(io::Error::from_raw_os_error(libc::EINVAL));
+    }
+
+    Ok(VsockAddr::new(
+        components[1].parse().map_err(|_| io::Error::from_raw_os_error(libc::EINVAL))?,
+        components[2].parse().map_err(|_| io::Error::from_raw_os_error(libc::EINVAL))?,
+    ))
+}
+
+/// StreamSocket provides a generic abstraction around any connection-oriented stream socket.
+/// The socket will be closed when StreamSocket is dropped, but writes to the socket can also
+/// be shut down manually.
+pub struct StreamSocket {
+    fd: RawFd,
+    shut_down: bool,
+}
+
+impl StreamSocket {
+    /// Connects to the given socket address. Supported socket types are vsock, unix, and TCP.
+    pub fn connect(sockaddr: &str) -> result::Result<StreamSocket, StreamSocketError> {
+        const UNIX_PREFIX: &str = "unix:";
+        const VSOCK_PREFIX: &str = "vsock:";
+
+        if sockaddr.starts_with(VSOCK_PREFIX) {
+            let addr = parse_vsock_addr(sockaddr)
+                .map_err(|e| StreamSocketError::ConnectVsock(sockaddr.to_string(), e))?;
+            let vsock_stream = VsockStream::connect(&addr)
+                .map_err(|e| StreamSocketError::ConnectVsock(sockaddr.to_string(), e))?;
+            Ok(vsock_stream.into())
+        } else if sockaddr.starts_with(UNIX_PREFIX) {
+            let (_prefix, sock_path) = sockaddr.split_at(UNIX_PREFIX.len());
+            let unix_stream = UnixStream::connect(sock_path)
+                .map_err(|e| StreamSocketError::ConnectUnix(sockaddr.to_string(), e))?;
+            Ok(unix_stream.into())
+        } else {
+            // Assume this is a TCP stream.
+            let tcp_stream = TcpStream::connect(sockaddr)
+                .map_err(|e| StreamSocketError::ConnectTcp(sockaddr.to_string(), e))?;
+            Ok(tcp_stream.into())
+        }
+    }
+
+    /// Shuts down writes to the socket using shutdown(2).
+    pub fn shut_down_write(&mut self) -> io::Result<()> {
+        // SAFETY:
+        // Safe because no memory is modified and the return value is checked.
+        let ret = unsafe { shutdown(self.fd, SHUT_WR) };
+        if ret < 0 {
+            return Err(io::Error::last_os_error());
+        }
+
+        self.shut_down = true;
+        Ok(())
+    }
+
+    /// Returns true if the socket has been shut down for writes, false otherwise.
+    pub fn is_shut_down(&self) -> bool {
+        self.shut_down
+    }
+}
+
+impl io::Read for StreamSocket {
+    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        // SAFETY:
+        // Safe because this will only modify the contents of |buf| and we check the return value.
+        let ret = unsafe { libc::read(self.fd, buf.as_mut_ptr() as *mut c_void, buf.len()) };
+        if ret < 0 {
+            return Err(io::Error::last_os_error());
+        }
+
+        Ok(ret as usize)
+    }
+}
+
+impl io::Write for StreamSocket {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        // SAFETY:
+        // Safe because this doesn't modify any memory and we check the return value.
+        let ret = unsafe { libc::write(self.fd, buf.as_ptr() as *const c_void, buf.len()) };
+        if ret < 0 {
+            // If a write causes EPIPE then the socket is shut down for writes.
+            let err = io::Error::last_os_error();
+            if let Some(errno) = err.raw_os_error() {
+                if errno == EPIPE {
+                    self.shut_down = true
+                }
+            }
+
+            return Err(err);
+        }
+
+        Ok(ret as usize)
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        // No buffered data so nothing to do.
+        Ok(())
+    }
+}
+
+impl AsRawFd for StreamSocket {
+    fn as_raw_fd(&self) -> RawFd {
+        self.fd
+    }
+}
+
+impl From<TcpStream> for StreamSocket {
+    fn from(stream: TcpStream) -> Self {
+        StreamSocket { fd: stream.into_raw_fd(), shut_down: false }
+    }
+}
+
+impl From<UnixStream> for StreamSocket {
+    fn from(stream: UnixStream) -> Self {
+        StreamSocket { fd: stream.into_raw_fd(), shut_down: false }
+    }
+}
+
+impl From<VsockStream> for StreamSocket {
+    fn from(stream: VsockStream) -> Self {
+        StreamSocket { fd: stream.into_raw_fd(), shut_down: false }
+    }
+}
+
+impl FromRawFd for StreamSocket {
+    unsafe fn from_raw_fd(fd: RawFd) -> Self {
+        StreamSocket { fd, shut_down: false }
+    }
+}
+
+impl Drop for StreamSocket {
+    fn drop(&mut self) {
+        // SAFETY:
+        // Safe because this doesn't modify any memory and we are the only
+        // owner of the file descriptor.
+        unsafe { libc::close(self.fd) };
+    }
+}
+
+/// Error enums for StreamSocket.
+#[remain::sorted]
+#[derive(Debug)]
+pub enum StreamSocketError {
+    /// Error on connecting TCP socket.
+    ConnectTcp(String, io::Error),
+    /// Error on connecting unix socket.
+    ConnectUnix(String, io::Error),
+    /// Error on connecting vsock socket.
+    ConnectVsock(String, io::Error),
+}
+
+impl fmt::Display for StreamSocketError {
+    #[remain::check]
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        use self::StreamSocketError::*;
+
+        #[remain::sorted]
+        match self {
+            ConnectTcp(sockaddr, e) => {
+                write!(f, "failed to connect to TCP sockaddr {}: {}", sockaddr, e)
+            }
+            ConnectUnix(sockaddr, e) => {
+                write!(f, "failed to connect to unix sockaddr {}: {}", sockaddr, e)
+            }
+            ConnectVsock(sockaddr, e) => {
+                write!(f, "failed to connect to vsock sockaddr {}: {}", sockaddr, e)
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::io::{Read, Write};
+    use std::net::TcpListener;
+    use std::os::unix::net::{UnixListener, UnixStream};
+    use tempfile::TempDir;
+
+    #[test]
+    fn sock_connect_tcp() {
+        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
+        let sockaddr = format!("127.0.0.1:{}", listener.local_addr().unwrap().port());
+
+        let _stream = StreamSocket::connect(&sockaddr).unwrap();
+    }
+
+    #[test]
+    fn sock_connect_unix() {
+        let tempdir = TempDir::new().unwrap();
+        let path = tempdir.path().to_owned().join("test.sock");
+        let _listener = UnixListener::bind(&path).unwrap();
+
+        let unix_addr = format!("unix:{}", path.to_str().unwrap());
+        let _stream = StreamSocket::connect(&unix_addr).unwrap();
+    }
+
+    #[test]
+    fn invalid_sockaddr() {
+        assert!(StreamSocket::connect("this is not a valid sockaddr").is_err());
+    }
+
+    #[test]
+    fn shut_down_write() {
+        let (unix_stream, _dummy) = UnixStream::pair().unwrap();
+        let mut stream: StreamSocket = unix_stream.into();
+
+        stream.write_all(b"hello").unwrap();
+
+        stream.shut_down_write().unwrap();
+
+        assert!(stream.is_shut_down());
+        assert!(stream.write(b"goodbye").is_err());
+    }
+
+    #[test]
+    fn read_from_shut_down_sock() {
+        let (unix_stream1, unix_stream2) = UnixStream::pair().unwrap();
+        let mut stream1: StreamSocket = unix_stream1.into();
+        let mut stream2: StreamSocket = unix_stream2.into();
+
+        stream1.shut_down_write().unwrap();
+
+        // Reads from the other end of the socket should now return EOF.
+        let mut buf = Vec::new();
+        assert_eq!(stream2.read_to_end(&mut buf).unwrap(), 0);
+    }
+}
diff --git a/tests/hostside/Android.bp b/tests/hostside/Android.bp
index 2eca2fa..d0838a6 100644
--- a/tests/hostside/Android.bp
+++ b/tests/hostside/Android.bp
@@ -16,7 +16,6 @@
     static_libs: [
         "MicrodroidHostTestHelper",
         "compatibility-host-util",
-        "cts-host-utils",
         "cts-statsd-atom-host-test-utils",
         "microdroid_payload_metadata",
     ],
diff --git a/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java b/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
index 2fb4483..2d55d66 100644
--- a/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
+++ b/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
@@ -32,8 +32,6 @@
 
 import static java.util.stream.Collectors.toList;
 
-import android.cts.host.utils.DeviceJUnit4ClassRunnerWithParameters;
-import android.cts.host.utils.DeviceJUnit4Parameterized;
 import android.cts.statsdatom.lib.ConfigUtils;
 import android.cts.statsdatom.lib.ReportUtils;
 
@@ -49,6 +47,7 @@
 import com.android.tradefed.device.ITestDevice;
 import com.android.tradefed.device.TestDevice;
 import com.android.tradefed.testtype.DeviceJUnit4ClassRunner.TestMetrics;
+import com.android.tradefed.testtype.junit4.DeviceParameterizedRunner;
 import com.android.tradefed.util.CommandResult;
 import com.android.tradefed.util.CommandStatus;
 import com.android.tradefed.util.FileUtil;
@@ -56,6 +55,9 @@
 import com.android.tradefed.util.xml.AbstractXmlParser;
 import com.android.virt.PayloadMetadata;
 
+import junitparams.Parameters;
+import junitparams.naming.TestCaseName;
+
 import org.json.JSONArray;
 import org.json.JSONObject;
 import org.junit.After;
@@ -65,8 +67,6 @@
 import org.junit.Test;
 import org.junit.rules.TestName;
 import org.junit.runner.RunWith;
-import org.junit.runners.Parameterized;
-import org.junit.runners.Parameterized.UseParametersRunnerFactory;
 import org.xml.sax.Attributes;
 import org.xml.sax.helpers.DefaultHandler;
 
@@ -76,8 +76,8 @@
 import java.io.PipedOutputStream;
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collection;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -88,8 +88,7 @@
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 
-@RunWith(DeviceJUnit4Parameterized.class)
-@UseParametersRunnerFactory(DeviceJUnit4ClassRunnerWithParameters.RunnerFactory.class)
+@RunWith(DeviceParameterizedRunner.class)
 public class MicrodroidHostTests extends MicrodroidHostTestCaseBase {
     private static final String APK_NAME = "MicrodroidTestApp.apk";
     private static final String APK_UPDATED_NAME = "MicrodroidTestAppUpdated.apk";
@@ -112,26 +111,38 @@
         }
     }
 
-    @Parameterized.Parameters(name = "protectedVm={0},gki={1}")
-    public static Collection<Object[]> params() {
-        List<Object[]> ret = new ArrayList<>();
-        ret.add(new Object[] {true /* protectedVm */, null /* use microdroid kernel */});
-        ret.add(new Object[] {false /* protectedVm */, null /* use microdroid kernel */});
+    // This map is needed because the parameterizer `DeviceParameterizedRunner` doesn't support "-"
+    // in test names. The key is the test name, while the value is the actual kernel version.
+    private static HashMap<String, String> sGkiVersions = new HashMap<>();
+
+    private static void initGkiVersions() {
+        if (!sGkiVersions.isEmpty()) {
+            return;
+        }
+        sGkiVersions.put("null", null); /* use microdroid kernel */
         // TODO(b/302465542): run only the latest GKI on presubmit to reduce running time
         for (String gki : SUPPORTED_GKI_VERSIONS) {
-            ret.add(new Object[] {true /* protectedVm */, gki});
-            ret.add(new Object[] {false /* protectedVm */, gki});
+            String key = gki.split("-")[0];
+            assertThat(sGkiVersions.containsKey(key)).isFalse();
+            sGkiVersions.put(key, gki);
+        }
+    }
+
+    public static List<Object[]> params() {
+        List<Object[]> ret = new ArrayList<>();
+        for (Object[] gki : gkiVersions()) {
+            ret.add(new Object[] {true /* protectedVm */, gki[0]});
+            ret.add(new Object[] {false /* protectedVm */, gki[0]});
         }
         return ret;
     }
 
-    @Parameterized.Parameter(0)
-    public boolean mProtectedVm;
-
-    @Parameterized.Parameter(1)
-    public String mGki;
-
-    private String mOs;
+    public static List<Object[]> gkiVersions() {
+        initGkiVersions();
+        return sGkiVersions.keySet().stream()
+                .map(gki -> new Object[] {gki})
+                .collect(Collectors.toList());
+    }
 
     @Rule public TestLogData mTestLogs = new TestLogData();
     @Rule public TestName mTestName = new TestName();
@@ -291,9 +302,11 @@
             File key,
             Map<String, File> keyOverrides,
             boolean isProtected,
-            boolean updateBootconfigs)
+            boolean updateBootconfigs,
+            String gki)
             throws Exception {
         CommandRunner android = new CommandRunner(getDevice());
+        gki = sGkiVersions.get(gki);
 
         File virtApexDir = FileUtil.createTempDir("virt_apex");
 
@@ -345,7 +358,8 @@
         //   - its idsig
 
         // Load etc/microdroid.json
-        File microdroidConfigFile = new File(virtApexEtcDir, mOs + ".json");
+        final String os = (gki == null) ? "microdroid" : "microdroid_gki-" + gki;
+        File microdroidConfigFile = new File(virtApexEtcDir, os + ".json");
         JSONObject config = new JSONObject(FileUtil.readStringFromFile(microdroidConfigFile));
 
         // Replace paths so that the config uses re-signed images from TEST_ROOT
@@ -361,7 +375,7 @@
         }
 
         // Add partitions to the second disk
-        final String initrdPath = TEST_ROOT + "etc/" + mOs + "_initrd_debuggable.img";
+        final String initrdPath = TEST_ROOT + "etc/" + os + "_initrd_debuggable.img";
         config.put("initrd", initrdPath);
         // Add instance image as a partition in disks[1]
         disks.put(
@@ -409,6 +423,9 @@
                         "--console " + CONSOLE_PATH,
                         "--log " + LOG_PATH,
                         configPath);
+        if (gki != null) {
+            args.add("--gki " + gki);
+        }
 
         PipedInputStream pis = new PipedInputStream();
         Process process = createRunUtil().runCmdInBackground(args, new PipedOutputStream(pis));
@@ -419,28 +436,32 @@
     @CddTest
     @VsrTest(requirements = {"VSR-7.1-001.008"})
     public void UpgradedPackageIsAcceptedWithSecretkeeper() throws Exception {
+        // Preconditions
+        assumeVmTypeSupported(true);
         assumeUpdatableVmSupported();
+
         getDevice().uninstallPackage(PACKAGE_NAME);
         getDevice().installPackage(findTestFile(APK_NAME), /* reinstall= */ true);
-        ensureMicrodroidBootsSuccessfully(INSTANCE_ID_FILE, INSTANCE_IMG);
+        ensureProtectedMicrodroidBootsSuccessfully(INSTANCE_ID_FILE, INSTANCE_IMG);
 
         getDevice().uninstallPackage(PACKAGE_NAME);
         cleanUpVirtualizationTestSetup(getDevice());
         // Install the updated version of app (versionCode 6)
         getDevice().installPackage(findTestFile(APK_UPDATED_NAME), /* reinstall= */ true);
-        ensureMicrodroidBootsSuccessfully(INSTANCE_ID_FILE, INSTANCE_IMG);
+        ensureProtectedMicrodroidBootsSuccessfully(INSTANCE_ID_FILE, INSTANCE_IMG);
     }
 
     @Test
     @CddTest
     @VsrTest(requirements = {"VSR-7.1-001.008"})
     public void DowngradedPackageIsRejectedProtectedVm() throws Exception {
-        assumeProtectedVm(); // Rollback protection is provided only for protected VM.
+        // Preconditions: Rollback protection is provided only for protected VM.
+        assumeVmTypeSupported(true);
 
         // Install the upgraded version (v6)
         getDevice().uninstallPackage(PACKAGE_NAME);
         getDevice().installPackage(findTestFile(APK_UPDATED_NAME), /* reinstall= */ true);
-        ensureMicrodroidBootsSuccessfully(INSTANCE_ID_FILE, INSTANCE_IMG);
+        ensureProtectedMicrodroidBootsSuccessfully(INSTANCE_ID_FILE, INSTANCE_IMG);
 
         getDevice().uninstallPackage(PACKAGE_NAME);
         cleanUpVirtualizationTestSetup(getDevice());
@@ -450,11 +471,11 @@
         assertThrows(
                 "pVM must fail to boot with downgraded payload apk",
                 DeviceRuntimeException.class,
-                () -> ensureMicrodroidBootsSuccessfully(INSTANCE_ID_FILE, INSTANCE_IMG));
+                () -> ensureProtectedMicrodroidBootsSuccessfully(INSTANCE_ID_FILE, INSTANCE_IMG));
     }
 
-    private void ensureMicrodroidBootsSuccessfully(String instanceIdPath, String instanceImgPath)
-            throws DeviceNotAvailableException {
+    private void ensureProtectedMicrodroidBootsSuccessfully(
+            String instanceIdPath, String instanceImgPath) throws DeviceNotAvailableException {
         final String configPath = "assets/vm_config.json";
         ITestDevice microdroid = null;
         int timeout = 30000; // 30 seconds
@@ -464,7 +485,7 @@
                             .debugLevel("full")
                             .memoryMib(minMemorySize())
                             .cpuTopology("match_host")
-                            .protectedVm(mProtectedVm)
+                            .protectedVm(true)
                             .instanceIdFile(instanceIdPath)
                             .instanceImgFile(instanceImgPath)
                             .setAdbConnectTimeoutMs(timeout)
@@ -479,10 +500,13 @@
     }
 
     @Test
+    @Parameters(method = "gkiVersions")
+    @TestCaseName("{method}_gki_{0}")
     @CddTest(requirements = {"9.17/C-2-1", "9.17/C-2-2", "9.17/C-2-6"})
-    public void protectedVmRunsPvmfw() throws Exception {
+    public void protectedVmRunsPvmfw(String gki) throws Exception {
         // Arrange
-        assumeProtectedVm();
+        assumeKernelSupported(gki);
+        assumeVmTypeSupported(true);
         final String configPath = "assets/vm_config_apex.json";
 
         // Act
@@ -492,7 +516,7 @@
                         .memoryMib(minMemorySize())
                         .cpuTopology("match_host")
                         .protectedVm(true)
-                        .gki(mGki)
+                        .gki(sGkiVersions.get(gki))
                         .name("protected_vm_runs_pvmfw")
                         .build(getAndroidDevice());
 
@@ -509,10 +533,13 @@
     }
 
     @Test
+    @Parameters(method = "gkiVersions")
+    @TestCaseName("{method}_gki_{0}")
     @CddTest(requirements = {"9.17/C-2-1", "9.17/C-2-2", "9.17/C-2-5", "9.17/C-2-6"})
-    public void protectedVmWithImageSignedWithDifferentKeyFailsToVerifyPayload() throws Exception {
-        // Arrange
-        assumeProtectedVm();
+    public void protectedVmWithImageSignedWithDifferentKeyFailsToVerifyPayload(String gki)
+            throws Exception {
+        assumeKernelSupported(gki);
+        assumeVmTypeSupported(true);
         File key = findTestFile("test.com.android.virt.pem");
 
         // Act
@@ -521,7 +548,8 @@
                         key,
                         /* keyOverrides= */ Map.of(),
                         /* isProtected= */ true,
-                        /* updateBootconfigs= */ true);
+                        /* updateBootconfigs= */ true,
+                        gki);
 
         // Assert
         vmInfo.mProcess.waitFor(5L, TimeUnit.SECONDS);
@@ -534,15 +562,23 @@
     }
 
     @Test
+    @Parameters(method = "gkiVersions")
+    @TestCaseName("{method}_gki_{0}")
     @CddTest(requirements = {"9.17/C-2-2", "9.17/C-2-6"})
-    public void testBootSucceedsWhenNonProtectedVmStartsWithImagesSignedWithDifferentKey()
+    public void testBootSucceedsWhenNonProtectedVmStartsWithImagesSignedWithDifferentKey(String gki)
             throws Exception {
-        assumeNonProtectedVm();
+        // Preconditions
+        assumeKernelSupported(gki);
+
         File key = findTestFile("test.com.android.virt.pem");
         Map<String, File> keyOverrides = Map.of();
         VmInfo vmInfo =
                 runMicrodroidWithResignedImages(
-                        key, keyOverrides, /* isProtected= */ false, /* updateBootconfigs= */ true);
+                        key,
+                        keyOverrides,
+                        /* isProtected= */ false,
+                        /* updateBootconfigs= */ true,
+                        gki);
         assertThatEventually(
                 100000,
                 () ->
@@ -554,16 +590,23 @@
     }
 
     @Test
+    @Parameters(method = "gkiVersions")
+    @TestCaseName("{method}_gki_{0}")
     @CddTest(requirements = {"9.17/C-2-2", "9.17/C-2-5", "9.17/C-2-6"})
-    public void testBootFailsWhenVbMetaDigestDoesNotMatchBootconfig() throws Exception {
+    public void testBootFailsWhenVbMetaDigestDoesNotMatchBootconfig(String gki) throws Exception {
         // protectedVmWithImageSignedWithDifferentKeyRunsPvmfw() is the protected case.
-        assumeNonProtectedVm();
+        assumeKernelSupported(gki);
+
         // Sign everything with key1 except vbmeta
         File key = findTestFile("test.com.android.virt.pem");
         // To be able to stop it, it should be a daemon.
         VmInfo vmInfo =
                 runMicrodroidWithResignedImages(
-                        key, Map.of(), /* isProtected= */ false, /* updateBootconfigs= */ false);
+                        key,
+                        Map.of(),
+                        /* isProtected= */ false,
+                        /* updateBootconfigs= */ false,
+                        gki);
         // Wait so that init can print errors to console (time in cuttlefish >> in real device)
         assertThatEventually(
                 100000,
@@ -611,7 +654,8 @@
     }
 
     private boolean isTombstoneGeneratedWithCmd(
-            boolean protectedVm, String configPath, String... crashCommand) throws Exception {
+            boolean protectedVm, String gki, String configPath, String... crashCommand)
+            throws Exception {
         CommandRunner android = new CommandRunner(getDevice());
         String testStartTime = android.runWithTimeout(1000, "date", "'+%Y-%m-%d %H:%M:%S.%N'");
 
@@ -621,7 +665,7 @@
                         .memoryMib(minMemorySize())
                         .cpuTopology("match_host")
                         .protectedVm(protectedVm)
-                        .gki(mGki)
+                        .gki(sGkiVersions.get(gki))
                         .build(getAndroidDevice());
         mMicrodroidDevice.waitForBootComplete(BOOT_COMPLETE_TIMEOUT);
         mMicrodroidDevice.enableAdbRoot();
@@ -637,12 +681,19 @@
     }
 
     @Test
-    public void testTombstonesAreGeneratedUponUserspaceCrash() throws Exception {
+    @Parameters(method = "params")
+    @TestCaseName("{method}_protectedVm_{0}_gki_{1}")
+    public void testTombstonesAreGeneratedUponUserspaceCrash(boolean protectedVm, String gki)
+            throws Exception {
+        // Preconditions
+        assumeKernelSupported(gki);
+        assumeVmTypeSupported(protectedVm);
         // TODO(b/291867858): tombstones are failing in HWASAN enabled Microdroid.
         assumeFalse("tombstones are failing in HWASAN enabled Microdroid.", isHwasan());
         assertThat(
                         isTombstoneGeneratedWithCmd(
-                                mProtectedVm,
+                                protectedVm,
+                                gki,
                                 "assets/vm_config.json",
                                 "kill",
                                 "-SIGSEGV",
@@ -651,12 +702,19 @@
     }
 
     @Test
-    public void testTombstonesAreNotGeneratedIfNotExportedUponUserspaceCrash() throws Exception {
+    @Parameters(method = "params")
+    @TestCaseName("{method}_protectedVm_{0}_gki_{1}")
+    public void testTombstonesAreNotGeneratedIfNotExportedUponUserspaceCrash(
+            boolean protectedVm, String gki) throws Exception {
+        // Preconditions
+        assumeKernelSupported(gki);
+        assumeVmTypeSupported(protectedVm);
         // TODO(b/291867858): tombstones are failing in HWASAN enabled Microdroid.
         assumeFalse("tombstones are failing in HWASAN enabled Microdroid.", isHwasan());
         assertThat(
                         isTombstoneGeneratedWithCmd(
-                                mProtectedVm,
+                                protectedVm,
+                                gki,
                                 "assets/vm_config_no_tombstone.json",
                                 "kill",
                                 "-SIGSEGV",
@@ -665,13 +723,22 @@
     }
 
     @Test
+    @Parameters(method = "params")
+    @TestCaseName("{method}_protectedVm_{0}_gki_{1}")
     @Ignore("b/341087884") // TODO(b/341087884): fix & re-enable
-    public void testTombstonesAreGeneratedUponKernelCrash() throws Exception {
+    public void testTombstonesAreGeneratedUponKernelCrash(boolean protectedVm, String gki)
+            throws Exception {
+        // Preconditions
+        assumeKernelSupported(gki);
+        assumeVmTypeSupported(protectedVm);
         assumeFalse("Cuttlefish is not supported", isCuttlefish());
         assumeFalse("Skipping test because ramdump is disabled on user build", isUserBuild());
+
+        // Act
         assertThat(
                         isTombstoneGeneratedWithCmd(
-                                mProtectedVm,
+                                protectedVm,
+                                gki,
                                 "assets/vm_config.json",
                                 "echo",
                                 "c",
@@ -681,10 +748,12 @@
     }
 
     private boolean isTombstoneGeneratedWithVmRunApp(
-            boolean protectedVm, boolean debuggable, String... additionalArgs) throws Exception {
+            boolean protectedVm, String gki, boolean debuggable, String... additionalArgs)
+            throws Exception {
         // we can't use microdroid builder as it wants ADB connection (debuggable)
         CommandRunner android = new CommandRunner(getDevice());
         String testStartTime = android.runWithTimeout(1000, "date", "'+%Y-%m-%d %H:%M:%S.%N'");
+        gki = sGkiVersions.get(gki);
 
         android.run("rm", "-rf", TEST_ROOT + "*");
         android.run("mkdir", "-p", TEST_ROOT + "*");
@@ -711,9 +780,9 @@
         if (protectedVm) {
             cmd.add("--protected");
         }
-        if (mGki != null) {
+        if (gki != null) {
             cmd.add("--gki");
-            cmd.add(mGki);
+            cmd.add(gki);
         }
         Collections.addAll(cmd, additionalArgs);
 
@@ -721,52 +790,89 @@
         return isTombstoneReceivedFromHostLogcat(testStartTime);
     }
 
-    private boolean isTombstoneGeneratedWithCrashPayload(boolean protectedVm, boolean debuggable)
-            throws Exception {
+    private boolean isTombstoneGeneratedWithCrashPayload(
+            boolean protectedVm, String gki, boolean debuggable) throws Exception {
         return isTombstoneGeneratedWithVmRunApp(
-                protectedVm, debuggable, "--payload-binary-name", "MicrodroidCrashNativeLib.so");
+                protectedVm,
+                gki,
+                debuggable,
+                "--payload-binary-name",
+                "MicrodroidCrashNativeLib.so");
     }
 
     @Test
-    public void testTombstonesAreGeneratedWithCrashPayload() throws Exception {
+    @Parameters(method = "params")
+    @TestCaseName("{method}_protectedVm_{0}_gki_{1}")
+    public void testTombstonesAreGeneratedWithCrashPayload(boolean protectedVm, String gki)
+            throws Exception {
+        // Preconditions
         // TODO(b/291867858): tombstones are failing in HWASAN enabled Microdroid.
         assumeFalse("tombstones are failing in HWASAN enabled Microdroid.", isHwasan());
-        assertThat(isTombstoneGeneratedWithCrashPayload(mProtectedVm, /* debuggable= */ true))
+        assumeKernelSupported(gki);
+        assumeVmTypeSupported(protectedVm);
+
+        // Act
+        assertThat(isTombstoneGeneratedWithCrashPayload(protectedVm, gki, /* debuggable= */ true))
                 .isTrue();
     }
 
     @Test
-    public void testTombstonesAreNotGeneratedWithCrashPayloadWhenNonDebuggable() throws Exception {
+    @Parameters(method = "params")
+    @TestCaseName("{method}_protectedVm_{0}_gki_{1}")
+    public void testTombstonesAreNotGeneratedWithCrashPayloadWhenNonDebuggable(
+            boolean protectedVm, String gki) throws Exception {
+        // Preconditions
         // TODO(b/291867858): tombstones are failing in HWASAN enabled Microdroid.
         assumeFalse("tombstones are failing in HWASAN enabled Microdroid.", isHwasan());
-        assertThat(isTombstoneGeneratedWithCrashPayload(mProtectedVm, /* debuggable= */ false))
+        assumeKernelSupported(gki);
+        assumeVmTypeSupported(protectedVm);
+
+        // Act
+        assertThat(isTombstoneGeneratedWithCrashPayload(protectedVm, gki, /* debuggable= */ false))
                 .isFalse();
     }
 
-    private boolean isTombstoneGeneratedWithCrashConfig(boolean protectedVm, boolean debuggable)
-            throws Exception {
+    private boolean isTombstoneGeneratedWithCrashConfig(
+            boolean protectedVm, String gki, boolean debuggable) throws Exception {
         return isTombstoneGeneratedWithVmRunApp(
-                protectedVm, debuggable, "--config-path", "assets/vm_config_crash.json");
+                protectedVm, gki, debuggable, "--config-path", "assets/vm_config_crash.json");
     }
 
     @Test
-    public void testTombstonesAreGeneratedWithCrashConfig() throws Exception {
+    @Parameters(method = "params")
+    @TestCaseName("{method}_protectedVm_{0}_gki_{1}")
+    public void testTombstonesAreGeneratedWithCrashConfig(boolean protectedVm, String gki)
+            throws Exception {
+        // Preconditions
         // TODO(b/291867858): tombstones are failing in HWASAN enabled Microdroid.
         assumeFalse("tombstones are failing in HWASAN enabled Microdroid.", isHwasan());
-        assertThat(isTombstoneGeneratedWithCrashConfig(mProtectedVm, /* debuggable= */ true))
+        assumeKernelSupported(gki);
+        assumeVmTypeSupported(protectedVm);
+
+        // Act
+        assertThat(isTombstoneGeneratedWithCrashConfig(protectedVm, gki, /* debuggable= */ true))
                 .isTrue();
     }
 
     @Test
-    public void testTombstonesAreNotGeneratedWithCrashConfigWhenNonDebuggable() throws Exception {
+    @Parameters(method = "params")
+    @TestCaseName("{method}_protectedVm_{0}_gki_{1}")
+    public void testTombstonesAreNotGeneratedWithCrashConfigWhenNonDebuggable(
+            boolean protectedVm, String gki) throws Exception {
         // TODO(b/291867858): tombstones are failing in HWASAN enabled Microdroid.
         assumeFalse("tombstones are failing in HWASAN enabled Microdroid.", isHwasan());
-        assertThat(isTombstoneGeneratedWithCrashConfig(mProtectedVm, /* debuggable= */ false))
+        assumeKernelSupported(gki);
+        assumeVmTypeSupported(protectedVm);
+        assertThat(isTombstoneGeneratedWithCrashConfig(protectedVm, gki, /* debuggable= */ false))
                 .isFalse();
     }
 
     @Test
-    public void testTelemetryPushedAtoms() throws Exception {
+    @Parameters(method = "params")
+    @TestCaseName("{method}_protectedVm_{0}_gki_{1}")
+    public void testTelemetryPushedAtoms(boolean protectedVm, String gki) throws Exception {
+        assumeKernelSupported(gki);
+        assumeVmTypeSupported(protectedVm);
         // Reset statsd config and report before the test
         ConfigUtils.removeConfig(getDevice());
         ReportUtils.clearReports(getDevice());
@@ -787,8 +893,8 @@
                         .debugLevel("full")
                         .memoryMib(minMemorySize())
                         .cpuTopology("match_host")
-                        .protectedVm(mProtectedVm)
-                        .gki(mGki)
+                        .protectedVm(protectedVm)
+                        .gki(sGkiVersions.get(gki))
                         .name("test_telemetry_pushed_atoms")
                         .build(device);
         microdroid.waitForBootComplete(BOOT_COMPLETE_TIMEOUT);
@@ -818,7 +924,7 @@
             assertThat(atomVmCreationRequested.getHypervisor())
                     .isEqualTo(AtomsProto.VmCreationRequested.Hypervisor.PKVM);
         }
-        assertThat(atomVmCreationRequested.getIsProtected()).isEqualTo(mProtectedVm);
+        assertThat(atomVmCreationRequested.getIsProtected()).isEqualTo(protectedVm);
         assertThat(atomVmCreationRequested.getCreationSucceeded()).isTrue();
         assertThat(atomVmCreationRequested.getBinderExceptionCode()).isEqualTo(0);
         assertThat(atomVmCreationRequested.getVmIdentifier())
@@ -924,29 +1030,41 @@
     }
 
     @Test
-    @CddTest(requirements = {"9.17/C-1-1", "9.17/C-1-2", "9.17/C-1-3"})
-    public void testMicrodroidBoots() throws Exception {
+    @Parameters(method = "params")
+    @TestCaseName("{method}_protectedVm_{0}_gki_{1}")
+    @CddTest(requirements = {"9.17/C-1-1", "9.17/C-1-2", "9.17/C/1-3"})
+    public void testMicrodroidBoots(boolean protectedVm, String gki) throws Exception {
+        // Preconditions
+        assumeKernelSupported(gki);
+        assumeVmTypeSupported(protectedVm);
+
         final String configPath = "assets/vm_config.json"; // path inside the APK
         testMicrodroidBootsWithBuilder(
                 MicrodroidBuilder.fromDevicePath(getPathForPackage(PACKAGE_NAME), configPath)
                         .debugLevel("full")
                         .memoryMib(minMemorySize())
                         .cpuTopology("match_host")
-                        .protectedVm(mProtectedVm)
+                        .protectedVm(protectedVm)
                         .name("test_microdroid_boots")
-                        .gki(mGki));
+                        .gki(sGkiVersions.get(gki)));
     }
 
     @Test
-    public void testMicrodroidRamUsage() throws Exception {
+    @Parameters(method = "params")
+    @TestCaseName("{method}_protectedVm_{0}_gki_{1}")
+    public void testMicrodroidRamUsage(boolean protectedVm, String gki) throws Exception {
+        // Preconditions
+        assumeKernelSupported(gki);
+        assumeVmTypeSupported(protectedVm);
+
         final String configPath = "assets/vm_config.json";
         mMicrodroidDevice =
                 MicrodroidBuilder.fromDevicePath(getPathForPackage(PACKAGE_NAME), configPath)
                         .debugLevel("full")
                         .memoryMib(minMemorySize())
                         .cpuTopology("match_host")
-                        .protectedVm(mProtectedVm)
-                        .gki(mGki)
+                        .protectedVm(protectedVm)
+                        .gki(sGkiVersions.get(gki))
                         .name("test_microdroid_ram_usage")
                         .build(getAndroidDevice());
         mMicrodroidDevice.waitForBootComplete(BOOT_COMPLETE_TIMEOUT);
@@ -1128,8 +1246,12 @@
     }
 
     @Test
-    public void testDeviceAssignment() throws Exception {
-        // Check for preconditions
+    @Parameters(method = "params")
+    @TestCaseName("{method}_protectedVm_{0}_gki_{1}")
+    public void testDeviceAssignment(boolean protectedVm, String gki) throws Exception {
+        // Preconditions
+        assumeKernelSupported(gki);
+        assumeVmTypeSupported(protectedVm);
         assumeVfioPlatformSupported();
 
         List<AssignableDevice> devices = getAssignableDevices();
@@ -1139,7 +1261,7 @@
 
         // Try assign devices one by one
         for (AssignableDevice device : devices) {
-            launchWithDeviceAssignment(device.node);
+            launchWithDeviceAssignment(device.node, protectedVm, gki);
 
             String dtPath =
                     new CommandRunner(mMicrodroidDevice)
@@ -1161,7 +1283,8 @@
         }
     }
 
-    private void launchWithDeviceAssignment(String device) throws Exception {
+    private void launchWithDeviceAssignment(String device, boolean protectedVm, String gki)
+            throws Exception {
         Objects.requireNonNull(device);
         final String configPath = "assets/vm_config.json";
 
@@ -1170,8 +1293,8 @@
                         .debugLevel("full")
                         .memoryMib(minMemorySize())
                         .cpuTopology("match_host")
-                        .protectedVm(mProtectedVm)
-                        .gki(mGki)
+                        .protectedVm(protectedVm)
+                        .gki(sGkiVersions.get(gki))
                         .addAssignableDevice(device)
                         .build(getAndroidDevice());
 
@@ -1189,7 +1312,13 @@
     }
 
     @Test
-    public void testHugePages() throws Exception {
+    @Parameters(method = "params")
+    @TestCaseName("{method}_protectedVm_{0}_gki_{1}")
+    public void testHugePages(boolean protectedVm, String gki) throws Exception {
+        // Preconditions
+        assumeKernelSupported(gki);
+        assumeVmTypeSupported(protectedVm);
+
         ITestDevice device = getDevice();
         boolean disableRoot = !device.isAdbRoot();
         CommandRunner android = new CommandRunner(device);
@@ -1210,8 +1339,8 @@
                         .debugLevel("full")
                         .memoryMib(minMemorySize())
                         .cpuTopology("match_host")
-                        .protectedVm(mProtectedVm)
-                        .gki(mGki)
+                        .protectedVm(protectedVm)
+                        .gki(sGkiVersions.get(gki))
                         .hugePages(true)
                         .name("test_huge_pages")
                         .build(getAndroidDevice());
@@ -1233,19 +1362,6 @@
 
         getDevice().installPackage(findTestFile(APK_NAME), /* reinstall= */ false);
 
-        // Skip test if given device doesn't support protected or non-protected VM.
-        assumeTrue(
-                "Microdroid is not supported for specific VM protection type",
-                getAndroidDevice().supportsMicrodroid(mProtectedVm));
-
-        if (mGki != null) {
-            assumeTrue(
-                    "GKI version \"" + mGki + "\" is not supported on this device",
-                    getSupportedGKIVersions().contains(mGki));
-        }
-
-        mOs = (mGki != null) ? "microdroid_gki-" + mGki : "microdroid";
-
         new CommandRunner(getDevice())
                 .tryRun(
                         "pm",
@@ -1268,14 +1384,6 @@
         getDevice().uninstallPackage(PACKAGE_NAME);
     }
 
-    private void assumeProtectedVm() {
-        assumeTrue("This test is only for protected VM.", mProtectedVm);
-    }
-
-    private void assumeNonProtectedVm() {
-        assumeFalse("This test is only for non-protected VM.", mProtectedVm);
-    }
-
     private void assumeVfioPlatformSupported() throws Exception {
         TestDevice device = getAndroidDevice();
         assumeTrue(
@@ -1305,4 +1413,19 @@
         runUtil.unsetEnvVariable("LD_LIBRARY_PATH");
         return runUtil;
     }
+
+    private void assumeKernelSupported(String gki) throws Exception {
+        String gkiVersion = sGkiVersions.get(gki);
+        if (gkiVersion != null) {
+            assumeTrue(
+                    "Skipping test as the GKI is not supported: " + gkiVersion,
+                    getSupportedGKIVersions().contains(gkiVersion));
+        }
+    }
+
+    private void assumeVmTypeSupported(boolean protectedVm) throws Exception {
+        assumeTrue(
+                "Microdroid is not supported for specific VM protection type",
+                getAndroidDevice().supportsMicrodroid(protectedVm));
+    }
 }