From c44199a8af792652c9cff1e618baa883c4545cfb Mon Sep 17 00:00:00 2001
From: Sai Diliyaer <sai.diliyaer@gmail.com>
Date: Wed, 9 Oct 2024 18:50:52 -0400
Subject: [PATCH] fix: update controller's device list in AssignController

---
 object/virtual_device_list.go      |  4 +++-
 object/virtual_device_list_test.go |  8 +++++++-
 simulator/virtual_machine.go       | 14 ++++++++++++++
 3 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/object/virtual_device_list.go b/object/virtual_device_list.go
index 295171af7..4d4335eac 100644
--- a/object/virtual_device_list.go
+++ b/object/virtual_device_list.go
@@ -1,5 +1,5 @@
 /*
-Copyright (c) 2015-2017 VMware, Inc. All Rights Reserved.
+Copyright (c) 2015-2024 VMware, Inc. All Rights Reserved.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -545,6 +545,8 @@ func (l VirtualDeviceList) AssignController(device types.BaseVirtualDevice, c ty
 	if d.Key == 0 {
 		d.Key = l.newRandomKey()
 	}
+
+	c.GetVirtualController().Device = append(c.GetVirtualController().Device, d.Key)
 }
 
 // newRandomKey returns a random negative device key.
diff --git a/object/virtual_device_list_test.go b/object/virtual_device_list_test.go
index 466fb02d9..9cfc32736 100644
--- a/object/virtual_device_list_test.go
+++ b/object/virtual_device_list_test.go
@@ -1,5 +1,5 @@
 /*
-Copyright (c) 2015 VMware, Inc. All Rights Reserved.
+Copyright (c) 2015-2024 VMware, Inc. All Rights Reserved.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -744,6 +744,12 @@ func TestAssignController(t *testing.T) {
 	if disk.Key >= 0 {
 		t.Errorf("device key %d should be negative", disk.Key)
 	}
+
+	// AssignController should add the disk to the controller's device list.
+	cd := scsi.(*types.VirtualLsiLogicController).Device[0]
+	if cd != disk.Key {
+		t.Errorf("expected controller device key: %d, got: %d\n", disk.Key, cd)
+	}
 }
 
 func TestCreateSCSIController(t *testing.T) {
diff --git a/simulator/virtual_machine.go b/simulator/virtual_machine.go
index a596798d6..470b59c81 100644
--- a/simulator/virtual_machine.go
+++ b/simulator/virtual_machine.go
@@ -1279,12 +1279,26 @@ func (vm *VirtualMachine) configureDevice(
 	d := device.GetVirtualDevice()
 	var controller types.BaseVirtualController
 
+	key := d.Key
 	if d.Key <= 0 {
 		// Keys can't be negative; Key 0 is reserved
 		d.Key = devices.NewKey()
 		d.Key *= -1
 	}
 
+	// Update device controller's key reference
+	if key != d.Key {
+		if device := devices.FindByKey(d.ControllerKey); device != nil {
+			c := device.(types.BaseVirtualController).GetVirtualController()
+			for i := range c.Device {
+				if c.Device[i] == key {
+					c.Device[i] = d.Key
+					break
+				}
+			}
+		}
+	}
+
 	// Choose a unique key
 	for {
 		if devices.FindByKey(d.Key) == nil {