fix: delete sockets on shim exit

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

integration/client/container_linux_test.go

@@ -39,6 +39,7 @@ import (
 	"github.com/containerd/containerd/v2/core/containers"
 	"github.com/containerd/containerd/v2/pkg/cio"
 	"github.com/containerd/containerd/v2/pkg/oci"
+	"github.com/containerd/containerd/v2/pkg/shim"
 	"github.com/containerd/containerd/v2/pkg/sys"
 	"github.com/containerd/containerd/v2/plugins"
 	"github.com/containerd/errdefs"
@@ -312,6 +313,67 @@ func TestShimDoesNotLeakPipes(t *testing.T) {
 	}
 }
 
+func TestShimDoesNotLeakSockets(t *testing.T) {
+	client, err := newClient(t, address)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer client.Close()
+
+	var (
+		image       Image
+		ctx, cancel = testContext(t)
+		id          = t.Name()
+	)
+	defer cancel()
+
+	image, err = client.GetImage(ctx, testImage)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	container, err := client.NewContainer(ctx, id, WithNewSnapshot(id, image), WithNewSpec(oci.WithImageConfig(image), withProcessArgs("sleep", "30")))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	task, err := container.NewTask(ctx, empty())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	exitChannel, err := task.Wait(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if err := task.Start(ctx); err != nil {
+		t.Fatal(err)
+	}
+
+	if err := task.Kill(ctx, syscall.SIGKILL); err != nil {
+		t.Fatal(err)
+	}
+
+	<-exitChannel
+
+	if _, err := task.Delete(ctx); err != nil {
+		t.Fatal(err)
+	}
+
+	if err := container.Delete(ctx, WithSnapshotCleanup); err != nil {
+		t.Fatal(err)
+	}
+
+	s, err := shim.SocketAddress(ctx, address, id)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if _, err = os.Stat(strings.TrimPrefix(s, "unix://")); err == nil || !os.IsNotExist(err) {
+		t.Errorf("Shim sockets have leaked after container has been deleted.")
+	}
+}
+
 func numPipes(pid int) (int, error) {
 	cmd := exec.Command("sh", "-c", fmt.Sprintf("lsof -p %d | grep FIFO", pid))
 

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():
@@ -479,3 +478,14 @@ func dumpStacks(logger *log.Entry) {
 	buf = buf[:stackSize]
 	logger.Infof("=== BEGIN goroutine stack dump ===\n%s\n=== END goroutine stack dump ===", buf)
 }
+
+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)
+	}
+}