diff --git a/rialto/src/main.rs b/rialto/src/main.rs
index 42d39c4..b34b9de 100644
--- a/rialto/src/main.rs
+++ b/rialto/src/main.rs
@@ -34,7 +34,7 @@
 use libfdt::FdtError;
 use log::{debug, error, info};
 use virtio_drivers::{
-    device::socket::VsockAddr,
+    device::socket::{VsockAddr, VMADDR_CID_HOST},
     transport::{pci::bus::PciRoot, DeviceType, Transport},
     Hal,
 };
@@ -52,12 +52,7 @@
 };
 
 fn host_addr() -> VsockAddr {
-    const PROTECTED_VM_PORT: u32 = 5679;
-    const NON_PROTECTED_VM_PORT: u32 = 5680;
-    const VMADDR_CID_HOST: u64 = 2;
-
-    let port = if is_protected_vm() { PROTECTED_VM_PORT } else { NON_PROTECTED_VM_PORT };
-    VsockAddr { cid: VMADDR_CID_HOST, port }
+    VsockAddr { cid: VMADDR_CID_HOST, port: service_vm_comm::host_port(is_protected_vm()) }
 }
 
 fn is_protected_vm() -> bool {
diff --git a/rialto/tests/test.rs b/rialto/tests/test.rs
index 2bd8968..e9bdab6 100644
--- a/rialto/tests/test.rs
+++ b/rialto/tests/test.rs
@@ -24,7 +24,7 @@
 };
 use anyhow::{anyhow, bail, Context, Result};
 use log::info;
-use service_vm_comm::{Request, Response};
+use service_vm_comm::{host_port, Request, Response};
 use std::fs::File;
 use std::io::{self, BufRead, BufReader, BufWriter, Write};
 use std::os::unix::io::FromRawFd;
@@ -34,11 +34,6 @@
 use vmclient::{DeathReason, VmInstance};
 use vsock::{VsockListener, VMADDR_CID_HOST};
 
-// TODO(b/291732060): Move the port numbers to the common library shared between the host
-// and rialto.
-const PROTECTED_VM_PORT: u32 = 5679;
-const NON_PROTECTED_VM_PORT: u32 = 5680;
-
 const SIGNED_RIALTO_PATH: &str = "/data/local/tmp/rialto_test/arm64/rialto.bin";
 const UNSIGNED_RIALTO_PATH: &str = "/data/local/tmp/rialto_test/arm64/rialto_unsigned.bin";
 const INSTANCE_IMG_PATH: &str = "/data/local/tmp/rialto_test/arm64/instance.img";
@@ -131,7 +126,7 @@
     )
     .context("Failed to create VM")?;
 
-    let port = if protected_vm { PROTECTED_VM_PORT } else { NON_PROTECTED_VM_PORT };
+    let port = host_port(protected_vm);
     let check_socket_handle = thread::spawn(move || try_check_socket_connection(port).unwrap());
 
     vm.start().context("Failed to start VM")?;