From 673ca4bcb2448b3c252fccff0d243932c97fd893 Mon Sep 17 00:00:00 2001
From: Dmitriy Matrenichev <dmitry.matrenichev@siderolabs.com>
Date: Thu, 30 Jan 2025 16:32:51 +0300
Subject: [PATCH] fix: ensure proper closure of client.Client.conn with
 finalizer

Always close the connection, even if we forgot to do this ourselves.

Signed-off-by: Dmitriy Matrenichev <dmitry.matrenichev@siderolabs.com>
---
 pkg/machinery/client/client.go      |  8 ++++
 pkg/machinery/client/client_test.go | 61 +++++++++++++++++++++++++++++
 pkg/machinery/client/export_test.go |  7 ++++
 3 files changed, 76 insertions(+)
 create mode 100644 pkg/machinery/client/client_test.go

diff --git a/pkg/machinery/client/client.go b/pkg/machinery/client/client.go
index 878691fbf1..8520a115fb 100644
--- a/pkg/machinery/client/client.go
+++ b/pkg/machinery/client/client.go
@@ -13,6 +13,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"runtime"
 	"time"
 
 	cosiv1alpha1 "github.com/cosi-project/runtime/api/v1alpha1"
@@ -173,11 +174,16 @@ func New(_ context.Context, opts ...OptionFunc) (c *Client, err error) {
 	c.Inspect = &InspectClient{c.InspectClient}
 	c.COSI = state.WrapCore(client.NewAdapter(cosiv1alpha1.NewStateClient(c.conn)))
 
+	// Just like Go team does in their os.File internals, lets ensure that we always close the connection.
+	runtime.SetFinalizer(c.conn, clientFinalizer)
+
 	return c, nil
 }
 
 // Close shuts down client protocol.
 func (c *Client) Close() error {
+	runtime.SetFinalizer(c.conn, nil) // remove so it doesn't get closed twice
+
 	return c.conn.Close()
 }
 
@@ -1011,3 +1017,5 @@ func (c *Client) BlockDeviceWipe(ctx context.Context, req *storageapi.BlockDevic
 
 	return err
 }
+
+var clientFinalizer = io.Closer.Close
diff --git a/pkg/machinery/client/client_test.go b/pkg/machinery/client/client_test.go
new file mode 100644
index 0000000000..ede49b7409
--- /dev/null
+++ b/pkg/machinery/client/client_test.go
@@ -0,0 +1,61 @@
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+package client_test
+
+import (
+	"context"
+	"io"
+	"runtime"
+	"testing"
+	"time"
+
+	"github.com/siderolabs/gen/xtesting/must"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+
+	"github.com/siderolabs/talos/pkg/machinery/client"
+)
+
+func TestNew(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	fCalled := make(chan struct{})
+
+	prev := client.SetClientFinalizer(func(closer io.Closer) error {
+		defer close(fCalled)
+
+		require.NoError(t, closer.Close())
+
+		return nil
+	})
+	defer client.SetClientFinalizer(prev)
+
+	c := must.Value(
+		client.New(
+			ctx,
+			client.WithUnixSocket("/path/to/socket"),
+			client.WithGRPCDialOptions(
+				grpc.WithTransportCredentials(insecure.NewCredentials()),
+			),
+		),
+	)(t)
+
+	runtime.SetFinalizer(c, func(c *client.Client) { t.Log("client finalizer set and called") })
+
+	for {
+		select {
+		case <-fCalled:
+			t.Log("client conn finalized")
+
+			return
+		case <-ctx.Done():
+			require.Fail(t, "client finalizer not called")
+		default:
+			runtime.GC()
+		}
+	}
+}
diff --git a/pkg/machinery/client/export_test.go b/pkg/machinery/client/export_test.go
index c005dfc80f..728ffa1fa4 100644
--- a/pkg/machinery/client/export_test.go
+++ b/pkg/machinery/client/export_test.go
@@ -6,6 +6,7 @@ package client
 
 import (
 	"crypto/tls"
+	"io"
 
 	clientconfig "github.com/siderolabs/talos/pkg/machinery/client/config"
 )
@@ -17,3 +18,9 @@ func ReduceURLsToAddresses(endpoints []string) []string {
 func BuildTLSConfig(configContext *clientconfig.Context) (*tls.Config, error) {
 	return buildTLSConfig(configContext)
 }
+
+func SetClientFinalizer(fn func(io.Closer) error) (old func(io.Closer) error) {
+	old, clientFinalizer = clientFinalizer, fn
+
+	return old
+}