portfwd: close ClosableListeners on exit

HostSocket will be deleted on close.

Signed-off-by: Norio Nomura <[email protected]>

Author: norio-nomura

pkg/hostagent/hostagent.go

@@ -788,6 +788,7 @@ func (a *HostAgent) processGuestAgentEvents(ctx context.Context, client *guestag
 			a.grpcPortForwarder.OnEvent(ctx, client, ev)
 		}
 	}
+	defer a.grpcPortForwarder.Close()
 
 	if err := client.Events(ctx, onEvent); err != nil {
 		if status.Code(err) == codes.Canceled {

pkg/portfwd/forward.go

@@ -34,6 +34,10 @@ func NewPortForwarder(rules []limatype.PortForward, ignoreTCP, ignoreUDP bool) *
 	}
 }
 
+func (fw *Forwarder) Close() {
+	fw.closableListeners.Close()
+}
+
 func (fw *Forwarder) OnEvent(ctx context.Context, client *guestagentclient.GuestAgentClient, ev *api.Event) {
 	for _, f := range ev.AddedLocalPorts {
 		// Before forwarding, check if any static rule matches this port otherwise it will be forwarded twice and cause a port conflict

pkg/portfwd/listener.go

@@ -35,6 +35,21 @@ func NewClosableListener() *ClosableListeners {
 	}
 }
 
+func (p *ClosableListeners) Close() {
+	p.listenersRW.Lock()
+	defer p.listenersRW.Unlock()
+	for _, listener := range p.listeners {
+		listener.Close()
+	}
+	clear(p.listeners)
+	p.udpListenersRW.Lock()
+	defer p.udpListenersRW.Unlock()
+	for _, listener := range p.udpListeners {
+		listener.Close()
+	}
+	clear(p.udpListeners)
+}
+
 func (p *ClosableListeners) Forward(ctx context.Context, client *guestagentclient.GuestAgentClient,
 	protocol string, hostAddress string, guestAddress string,
 ) {
@@ -90,6 +105,9 @@ func (p *ClosableListeners) forwardTCP(ctx context.Context, client *guestagentcl
 	for {
 		conn, err := tcpLis.Accept()
 		if err != nil {
+			if opErr, ok := err.(*net.OpError); ok && opErr.Err.Error() == "use of closed network connection" {
+				return
+			}
 			logrus.Errorf("failed to accept TCP connection: %v", err)
 			if strings.Contains(err.Error(), "pseudoloopback") {
 				// don't stop forwarding because the forwarder has rejected a non-local address