From b6f6404a815227aee2cb8f37e8536eb5f63cffd0 Mon Sep 17 00:00:00 2001
From: Black-Hole1 <bh@bugs.cc>
Date: Tue, 26 Dec 2023 17:03:53 +0800
Subject: [PATCH] refactor(vsock): call proxy.Close when vm stop

When other projects use the `vf.ExposeVsock` method, there will be unexpected issues due to the absence of a close proxy.

Signed-off-by: Black-Hole1 <bh@bugs.cc>
---
 cmd/vfkit/main.go | 19 +++----------------
 pkg/vf/vsock.go   | 46 ++++++++++++++++++++++++++++++++++++++++------
 2 files changed, 43 insertions(+), 22 deletions(-)

diff --git a/cmd/vfkit/main.go b/cmd/vfkit/main.go
index ea02dc02..0bed7731 100644
--- a/cmd/vfkit/main.go
+++ b/cmd/vfkit/main.go
@@ -156,22 +156,9 @@ func runVirtualMachine(vmConfig *config.VirtualMachine, vm *vz.VirtualMachine) e
 	}
 	log.Infof("virtual machine is running")
 
-	for _, vsock := range vmConfig.VirtioVsockDevices() {
-		port := vsock.Port
-		socketURL := vsock.SocketURL
-		if socketURL == "" {
-			// the timesync code adds a vsock device without an associated URL.
-			continue
-		}
-		var listenStr string
-		if vsock.Listen {
-			listenStr = " (listening)"
-		}
-		log.Infof("Exposing vsock port %d on %s%s", port, socketURL, listenStr)
-		if err := vf.ExposeVsock(vm, port, socketURL, vsock.Listen); err != nil {
-			log.Warnf("error exposing vsock port %d: %v", port, err)
-		}
-	}
+	vsockDevs := vmConfig.VirtioVsockDevices()
+	releases := vf.ExposeVsocks(vm, vsockDevs)
+	defer releases()
 
 	if err := setupGuestTimeSync(vm, vmConfig.TimeSync()); err != nil {
 		log.Warnf("Error configuring guest time synchronization")
diff --git a/pkg/vf/vsock.go b/pkg/vf/vsock.go
index 4a8e93af..cd8ad178 100644
--- a/pkg/vf/vsock.go
+++ b/pkg/vf/vsock.go
@@ -8,10 +8,40 @@ import (
 	"strconv"
 
 	"github.com/Code-Hex/vz/v3"
+	"github.com/crc-org/vfkit/pkg/config"
+	log "github.com/sirupsen/logrus"
 	"inet.af/tcpproxy"
 )
 
-func ExposeVsock(vm *vz.VirtualMachine, port uint, vsockPath string, listen bool) error {
+func ExposeVsocks(vm *vz.VirtualMachine, devs []*config.VirtioVsock) (release func()) {
+	releases := make([]func(), 0, len(devs))
+	for _, vsock := range devs {
+		port := vsock.Port
+		socketURL := vsock.SocketURL
+		if socketURL == "" {
+			// the timesync code adds a vsock device without an associated URL.
+			continue
+		}
+		var listenStr string
+		if vsock.Listen {
+			listenStr = " (listening)"
+		}
+		log.Infof("Exposing vsock port %d on %s%s", port, socketURL, listenStr)
+		if release, err := ExposeVsock(vm, port, socketURL, vsock.Listen); err != nil {
+			log.Warnf("error exposing vsock port %d: %v", port, err)
+		} else {
+			releases = append(releases, release)
+		}
+	}
+
+	return func() {
+		for _, r := range releases {
+			r()
+		}
+	}
+}
+
+func ExposeVsock(vm *vz.VirtualMachine, port uint, vsockPath string, listen bool) (release func(), err error) {
 	if listen {
 		return listenVsock(vm, port, vsockPath)
 	}
@@ -36,7 +66,7 @@ func ConnectVsockSync(vm *vz.VirtualMachine, port uint) (net.Conn, error) {
 
 // connectVsock proxies connections from a host unix socket to a vsock port
 // This allows the host to initiate connections to the guest over vsock
-func connectVsock(vm *vz.VirtualMachine, port uint, vsockPath string) error {
+func connectVsock(vm *vz.VirtualMachine, port uint, vsockPath string) (release func(), err error) {
 
 	var proxy tcpproxy.Proxy
 	// listen for connections on the host unix socket
@@ -70,12 +100,15 @@ func connectVsock(vm *vz.VirtualMachine, port uint, vsockPath string) error {
 			}
 		},
 	})
-	return proxy.Start()
+
+	return func() {
+		_ = proxy.Close()
+	}, proxy.Start()
 }
 
 // listenVsock proxies connections from a vsock port to a host unix socket.
 // This allows the guest to initiate connections to the host over vsock
-func listenVsock(vm *vz.VirtualMachine, port uint, vsockPath string) error {
+func listenVsock(vm *vz.VirtualMachine, port uint, vsockPath string) (release func(), err error) {
 	var proxy tcpproxy.Proxy
 	// listen for connections on the vsock port
 	proxy.ListenFunc = func(_, laddr string) (net.Listener, error) {
@@ -116,6 +149,7 @@ func listenVsock(vm *vz.VirtualMachine, port uint, vsockPath string) error {
 			}
 		},
 	})
-	// FIXME: defer proxy.Close()
-	return proxy.Start()
+	return func() {
+		_ = proxy.Close()
+	}, proxy.Start()
 }