fix: delete sockets on shim exit

Signed-off-by: Henry Wang <henwang@amazon.com>

pkg/shim/shim.go

@@ -411,15 +411,14 @@ func run(ctx context.Context, manager Manager, config Config) error {
 
 	if err := serve(ctx, server, signals, sd.Shutdown); err != nil {
 		if !errors.Is(err, shutdown.ErrShutdown) {
+			cleanupSockets(ctx)
 			return err
 		}
 	}
 
 	// NOTE: If the shim server is down(like oom killer), the address
 	// socket might be leaking.
-	if address, err := ReadAddress("address"); err == nil {
-		_ = RemoveSocket(address)
-	}
+	cleanupSockets(ctx)
 
 	select {
 	case <-sd.Done():

pkg/shim/util_unix.go

@@ -279,3 +279,14 @@ func dialHybridVsock(address string, timeout time.Duration) (net.Conn, error) {
 	}
 	return hybridVsockDialer(addr, port, timeout)
 }
+
+func cleanupSockets(ctx context.Context) {
+	if address, err := ReadAddress("address"); err == nil {
+		_ = RemoveSocket(address)
+	}
+	if len(socketFlag) > 0 {
+		_ = RemoveSocket("unix://" + socketFlag)
+	} else if address, err := SocketAddress(ctx, addressFlag, id); err == nil {
+		_ = RemoveSocket(address)
+	}
+}

pkg/shim/util_windows.go

@@ -85,3 +85,9 @@ func AnonDialer(address string, timeout time.Duration) (net.Conn, error) {
 func RemoveSocket(address string) error {
 	return nil
 }
+
+func cleanupSockets(context.Context) {
+	if address, err := ReadAddress("address"); err == nil {
+		_ = RemoveSocket(address)
+	}
+}