Regarding being able to handle a reboot from within a bootstrap script, the short answer is that it doesn't. You can trigger a reboot from the bootstrap script, but it needs to have a short delay (I used 2 seconds) to allow the script to terminate successfully before the reboot, otherwise the extension reports an error. There are then a couple of ways to run the rest of the script after the reboot. One is to simply execute the script directly from terraform using a local-exec provisioner and az vm run-command, and the other is to run it as a scheduled job on the VM, created to run on startup by the initial bootstrap script (before the reboot). I chose the latter as I wanted to keep it compatible with the original implementation (without the reboot) due do the same module being used by several systems. For testing I added a check to make sure the scheduled task was executed successfully (again using local-exec), but critically this is not actually required, although obviously good practice.