From e2478a0ca004451a33a41ba0fa3d6fb05b4577cc Mon Sep 17 00:00:00 2001
From: Black-Hole1 <bh@bugs.cc>
Date: Mon, 8 Jan 2024 11:24:22 +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 | 14 ++++++++++++--
 pkg/vf/vsock.go   | 13 +++++++------
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/cmd/vfkit/main.go b/cmd/vfkit/main.go
index ea02dc02..1c32b222 100644
--- a/cmd/vfkit/main.go
+++ b/cmd/vfkit/main.go
@@ -22,6 +22,7 @@ package main
 import (
 	"errors"
 	"fmt"
+	"io"
 	"os"
 	"os/signal"
 	"runtime"
@@ -156,7 +157,9 @@ func runVirtualMachine(vmConfig *config.VirtualMachine, vm *vz.VirtualMachine) e
 	}
 	log.Infof("virtual machine is running")
 
-	for _, vsock := range vmConfig.VirtioVsockDevices() {
+	vsockDevs := vmConfig.VirtioVsockDevices()
+	closes := make([]io.Closer, 0, len(vsockDevs))
+	for _, vsock := range vsockDevs {
 		port := vsock.Port
 		socketURL := vsock.SocketURL
 		if socketURL == "" {
@@ -168,10 +171,17 @@ func runVirtualMachine(vmConfig *config.VirtualMachine, vm *vz.VirtualMachine) e
 			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 {
+		if closer, err := vf.ExposeVsock(vm, port, socketURL, vsock.Listen); err != nil {
 			log.Warnf("error exposing vsock port %d: %v", port, err)
+		} else {
+			closes = append(closes, closer)
 		}
 	}
+	defer func() {
+		for _, c := range closes {
+			_ = c.Close()
+		}
+	}()
 
 	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..47ccf01d 100644
--- a/pkg/vf/vsock.go
+++ b/pkg/vf/vsock.go
@@ -3,6 +3,7 @@ package vf
 import (
 	"context"
 	"fmt"
+	"io"
 	"net"
 	"net/url"
 	"strconv"
@@ -11,7 +12,7 @@ import (
 	"inet.af/tcpproxy"
 )
 
-func ExposeVsock(vm *vz.VirtualMachine, port uint, vsockPath string, listen bool) error {
+func ExposeVsock(vm *vz.VirtualMachine, port uint, vsockPath string, listen bool) (io.Closer, error) {
 	if listen {
 		return listenVsock(vm, port, vsockPath)
 	}
@@ -36,7 +37,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) (io.Closer, error) {
 
 	var proxy tcpproxy.Proxy
 	// listen for connections on the host unix socket
@@ -70,12 +71,12 @@ func connectVsock(vm *vz.VirtualMachine, port uint, vsockPath string) error {
 			}
 		},
 	})
-	return proxy.Start()
+	return &proxy, 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) (io.Closer, error) {
 	var proxy tcpproxy.Proxy
 	// listen for connections on the vsock port
 	proxy.ListenFunc = func(_, laddr string) (net.Listener, error) {
@@ -116,6 +117,6 @@ func listenVsock(vm *vz.VirtualMachine, port uint, vsockPath string) error {
 			}
 		},
 	})
-	// FIXME: defer proxy.Close()
-	return proxy.Start()
+
+	return &proxy, proxy.Start()
 }