From b5d78d47f540fe10899195bb578b5c50671f2dba Mon Sep 17 00:00:00 2001
From: Adrien Delorme <adrien.delorme@icloud.com>
Date: Fri, 15 Mar 2019 16:37:44 +0100
Subject: [PATCH] vmware StepConfigureVNC: use common/net pkg to find open port

---
 builder/vmware/common/run_config.go         |  4 +-
 builder/vmware/common/step_configure_vnc.go | 62 ++++++++-------------
 2 files changed, 26 insertions(+), 40 deletions(-)

diff --git a/builder/vmware/common/run_config.go b/builder/vmware/common/run_config.go
index c4bfc3413..e8b6a3d37 100644
--- a/builder/vmware/common/run_config.go
+++ b/builder/vmware/common/run_config.go
@@ -10,8 +10,8 @@ type RunConfig struct {
 	Headless bool `mapstructure:"headless"`
 
 	VNCBindAddress     string `mapstructure:"vnc_bind_address"`
-	VNCPortMin         uint   `mapstructure:"vnc_port_min"`
-	VNCPortMax         uint   `mapstructure:"vnc_port_max"`
+	VNCPortMin         int    `mapstructure:"vnc_port_min"`
+	VNCPortMax         int    `mapstructure:"vnc_port_max"`
 	VNCDisablePassword bool   `mapstructure:"vnc_disable_password"`
 }
 
diff --git a/builder/vmware/common/step_configure_vnc.go b/builder/vmware/common/step_configure_vnc.go
index 54f5ac182..284c58edd 100644
--- a/builder/vmware/common/step_configure_vnc.go
+++ b/builder/vmware/common/step_configure_vnc.go
@@ -5,8 +5,8 @@ import (
 	"fmt"
 	"log"
 	"math/rand"
-	"net"
 
+	"github.com/hashicorp/packer/common/net"
 	"github.com/hashicorp/packer/helper/multistep"
 	"github.com/hashicorp/packer/packer"
 )
@@ -18,43 +18,20 @@ import (
 //   vmx_path string
 //
 // Produces:
-//   vnc_port uint - The port that VNC is configured to listen on.
+//   vnc_port int - The port that VNC is configured to listen on.
 type StepConfigureVNC struct {
 	Enabled            bool
 	VNCBindAddress     string
-	VNCPortMin         uint
-	VNCPortMax         uint
+	VNCPortMin         int
+	VNCPortMax         int
 	VNCDisablePassword bool
+
+	l *net.Listener
 }
 
 type VNCAddressFinder interface {
-	VNCAddress(string, uint, uint) (string, uint, error)
-
 	// UpdateVMX, sets driver specific VNC values to VMX data.
-	UpdateVMX(vncAddress, vncPassword string, vncPort uint, vmxData map[string]string)
-}
-
-func (StepConfigureVNC) VNCAddress(vncBindAddress string, portMin, portMax uint) (string, uint, error) {
-	// Find an open VNC port. Note that this can still fail later on
-	// because we have to release the port at some point. But this does its
-	// best.
-	var vncPort uint
-	portRange := int(portMax - portMin)
-	for {
-		if portRange > 0 {
-			vncPort = uint(rand.Intn(portRange)) + portMin
-		} else {
-			vncPort = portMin
-		}
-
-		log.Printf("Trying port: %d", vncPort)
-		l, err := net.Listen("tcp", fmt.Sprintf("%s:%d", vncBindAddress, vncPort))
-		if err == nil {
-			defer l.Close()
-			break
-		}
-	}
-	return vncBindAddress, vncPort, nil
+	UpdateVMX(vncAddress, vncPassword string, vncPort int, vmxData map[string]string)
 }
 
 func VNCPassword(skipPassword bool) string {
@@ -75,7 +52,7 @@ func VNCPassword(skipPassword bool) string {
 	return string(password)
 }
 
-func (s *StepConfigureVNC) Run(_ context.Context, state multistep.StateBag) multistep.StepAction {
+func (s *StepConfigureVNC) Run(ctx context.Context, state multistep.StateBag) multistep.StepAction {
 	if !s.Enabled {
 		log.Println("Skipping VNC configuration step...")
 		return multistep.ActionContinue
@@ -99,19 +76,27 @@ func (s *StepConfigureVNC) Run(_ context.Context, state multistep.StateBag) mult
 	} else {
 		vncFinder = s
 	}
+
 	log.Printf("Looking for available port between %d and %d", s.VNCPortMin, s.VNCPortMax)
-	vncBindAddress, vncPort, err := vncFinder.VNCAddress(s.VNCBindAddress, s.VNCPortMin, s.VNCPortMax)
+	s.l, err = net.ListenRangeConfig{
+		Addr:    s.VNCBindAddress,
+		Min:     s.VNCPortMin,
+		Max:     s.VNCPortMax,
+		Network: "tcp",
+	}.Listen(ctx)
+
 	if err != nil {
 		state.Put("error", err)
 		ui.Error(err.Error())
 		return multistep.ActionHalt
 	}
+	s.l.Listener.Close() // free port, but don't unlock lock file
 
 	vncPassword := VNCPassword(s.VNCDisablePassword)
 
-	log.Printf("Found available VNC port: %d", vncPort)
+	log.Printf("Found available VNC port: %v", s.l)
 
-	vncFinder.UpdateVMX(vncBindAddress, vncPassword, vncPort, vmxData)
+	vncFinder.UpdateVMX(s.l.Address, vncPassword, s.l.Port, vmxData)
 
 	if err := WriteVMX(vmxPath, vmxData); err != nil {
 		err := fmt.Errorf("Error writing VMX data: %s", err)
@@ -120,14 +105,12 @@ func (s *StepConfigureVNC) Run(_ context.Context, state multistep.StateBag) mult
 		return multistep.ActionHalt
 	}
 
-	state.Put("vnc_port", vncPort)
-	state.Put("vnc_ip", vncBindAddress)
 	state.Put("vnc_password", vncPassword)
 
 	return multistep.ActionContinue
 }
 
-func (StepConfigureVNC) UpdateVMX(address, password string, port uint, data map[string]string) {
+func (*StepConfigureVNC) UpdateVMX(address, password string, port int, data map[string]string) {
 	data["remotedisplay.vnc.enabled"] = "TRUE"
 	data["remotedisplay.vnc.port"] = fmt.Sprintf("%d", port)
 	data["remotedisplay.vnc.ip"] = address
@@ -136,5 +119,8 @@ func (StepConfigureVNC) UpdateVMX(address, password string, port uint, data map[
 	}
 }
 
-func (StepConfigureVNC) Cleanup(multistep.StateBag) {
+func (s *StepConfigureVNC) Cleanup(multistep.StateBag) {
+	if err := s.l.Close(); err != nil {
+		log.Printf("failed to unlock port lockfile: %v", err)
+	}
 }