From d03ee5820accf53b6b690743298112999e9a4338 Mon Sep 17 00:00:00 2001
From: Chris Coulson <chris.coulson@canonical.com>
Date: Mon, 20 Jan 2025 13:52:28 +0000
Subject: [PATCH] mssim: Use uint16 for representing ports everywhere

---
 mssim/device.go | 12 ++++++------
 mssim/mssim.go  |  2 +-
 testutil/tpm.go | 15 +++++++++++----
 3 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/mssim/device.go b/mssim/device.go
index b947bc5..8e67cf1 100644
--- a/mssim/device.go
+++ b/mssim/device.go
@@ -17,7 +17,7 @@ const (
 	// DefaultPort is the default IP port that the TPM channel of
 	// the simulator runs on. The platform port is normally this + 1,
 	// but can be customized by WithPlatformPort.
-	DefaultPort uint = 2321
+	DefaultPort uint16 = 2321
 )
 
 var (
@@ -31,7 +31,7 @@ var (
 
 type deviceAddr struct {
 	Host string
-	Port uint
+	Port uint16
 }
 
 func (a deviceAddr) Network() string {
@@ -95,7 +95,7 @@ func WithHost(host string) DeviceOption {
 // WithPort is used to customize the TCP ports on which the TPM and platform
 // channels for the simulator are accessed. It sets the platform channel port
 // to the TPM channel port + 1.
-func WithPort(port uint) DeviceOption {
+func WithPort(port uint16) DeviceOption {
 	return func(d *Device) {
 		d.tpm.Port = port
 		d.platform.Port = port + 1
@@ -105,7 +105,7 @@ func WithPort(port uint) DeviceOption {
 // WithTPMPort is used to customize the TCP port on which the TPM channel
 // for the simulator is accessed. It doesn't modify the port for the platform
 // channel.
-func WithTPMPort(port uint) DeviceOption {
+func WithTPMPort(port uint16) DeviceOption {
 	return func(d *Device) {
 		d.tpm.Port = port
 	}
@@ -114,7 +114,7 @@ func WithTPMPort(port uint) DeviceOption {
 // WithPlatformPort is used to customize the TCP port on which the platform
 // chanel for the simulator is accessed. It doesn't modify the port for the
 // TPM channel
-func WithPlatformPort(port uint) DeviceOption {
+func WithPlatformPort(port uint16) DeviceOption {
 	return func(d *Device) {
 		d.platform.Port = port
 	}
@@ -178,7 +178,7 @@ func (d *Device) Host() string {
 //
 // Deprecated: Use [Device.TPMAddr] or [Device.PlatformAddr] instead.
 func (d *Device) Port() uint {
-	return d.tpm.Port
+	return uint(d.tpm.Port)
 }
 
 func (d *Device) openInternal() (transport *Transport, err error) {
diff --git a/mssim/mssim.go b/mssim/mssim.go
index 22820f6..16a1c21 100644
--- a/mssim/mssim.go
+++ b/mssim/mssim.go
@@ -572,6 +572,6 @@ func (t *tpmMainTransport) Close() error {
 //
 // Deprecated: Use [NewDevice], [NewLocalDevice] or [DefaultDevice].
 func OpenConnection(host string, port uint) (*Transport, error) {
-	device := NewDevice(WithHost(host), WithPort(port))
+	device := NewDevice(WithHost(host), WithPort(uint16(port)))
 	return device.openInternal()
 }
diff --git a/testutil/tpm.go b/testutil/tpm.go
index 8d44fd9..8445247 100644
--- a/testutil/tpm.go
+++ b/testutil/tpm.go
@@ -169,7 +169,7 @@ var (
 	TPMDevicePath string = "/dev/tpm0"
 
 	// MssimPort defines the port number of the TPM simulator command port where TPMBackend is TPMBackendMssim.
-	MssimPort uint = mssim.DefaultPort
+	MssimPort uint16 = mssim.DefaultPort
 
 	wrapMssimTransport = WrapTransport
 
@@ -212,11 +212,18 @@ func AddCommandLineFlags() {
 	flag.Var(&PermittedTPMFeatures, "tpm-permitted-features", "Comma-separated list of features that tests can use on a TPM character device")
 
 	flag.StringVar(&TPMDevicePath, "tpm-path", "/dev/tpm0", "The path of the TPM character device to use for testing (default: /dev/tpm0)")
-	flag.UintVar(&MssimPort, "mssim-port", 2321, "The port number of the TPM simulator command channel (default: 2321)")
+	flag.Func("mssim-port", "The port number of the TPM simulator command channel (default: 2321)", func(opt string) error {
+		port, err := strconv.ParseUint(opt, 10, 16)
+		if err != nil {
+			return err
+		}
+		MssimPort = uint16(port)
+		return nil
+	})
 }
 
 type tpmSimulatorLaunchContext struct {
-	port               uint
+	port               uint16
 	persistentSavePath string
 	workDir            string
 	keepWorkDir        bool
@@ -512,7 +519,7 @@ Loop:
 type TPMSimulatorOptions struct {
 	// Port is the TCP port to use for the command channel. This port + 1 will also be used for the
 	// platform channel. If this is zero, then the value of [MssimPort] will be used.
-	Port uint
+	Port uint16
 
 	SourcePath     string    // Path for the source persistent data file
 	Manufacture    bool      // Indicates that the simulator should be executed in re-manufacture mode