From 3db3db3bf23a15735e17491dc6c98b0680df4bba Mon Sep 17 00:00:00 2001
From: Markus Rudy <mr@edgeless.systems>
Date: Tue, 25 Jun 2024 10:11:57 +0200
Subject: [PATCH] operator: fix gRPC dialing over UDS (#3201)

* operator: add test for gRPC connection over UDS
---
 .../internal/upgrade/BUILD.bazel              | 13 +++++
 .../internal/upgrade/upgrade.go               | 15 +++---
 .../internal/upgrade/upgrade_test.go          | 52 +++++++++++++++++++
 3 files changed, 73 insertions(+), 7 deletions(-)
 create mode 100644 operators/constellation-node-operator/internal/upgrade/upgrade_test.go

diff --git a/operators/constellation-node-operator/internal/upgrade/BUILD.bazel b/operators/constellation-node-operator/internal/upgrade/BUILD.bazel
index ec1a62223..cf1588a80 100644
--- a/operators/constellation-node-operator/internal/upgrade/BUILD.bazel
+++ b/operators/constellation-node-operator/internal/upgrade/BUILD.bazel
@@ -1,4 +1,5 @@
 load("@io_bazel_rules_go//go:def.bzl", "go_library")
+load("//bazel/go:go_test.bzl", "go_test")
 
 go_library(
     name = "upgrade",
@@ -13,3 +14,15 @@ go_library(
         "@org_golang_google_grpc//credentials/insecure",
     ],
 )
+
+go_test(
+    name = "upgrade_test",
+    srcs = ["upgrade_test.go"],
+    embed = [":upgrade"],
+    deps = [
+        "//internal/versions/components",
+        "//upgrade-agent/upgradeproto",
+        "@com_github_stretchr_testify//require",
+        "@org_golang_google_grpc//:grpc",
+    ],
+)
diff --git a/operators/constellation-node-operator/internal/upgrade/upgrade.go b/operators/constellation-node-operator/internal/upgrade/upgrade.go
index 615e92dac..7204f473d 100644
--- a/operators/constellation-node-operator/internal/upgrade/upgrade.go
+++ b/operators/constellation-node-operator/internal/upgrade/upgrade.go
@@ -21,24 +21,25 @@ import (
 
 // Client is a client for the upgrade agent.
 type Client struct {
+	addr   string
 	dialer Dialer
 }
 
-// NewClient creates a new upgrade agent client.
+// NewClient creates a new upgrade agent client connecting to the default upgrade-agent Unix socket.
 func NewClient() *Client {
+	return newClientWithAddress(mainconstants.UpgradeAgentMountPath)
+}
+
+func newClientWithAddress(addr string) *Client {
 	return &Client{
+		addr:   "unix:" + addr,
 		dialer: &net.Dialer{},
 	}
 }
 
 // Upgrade upgrades the Constellation node to the given Kubernetes version.
 func (c *Client) Upgrade(ctx context.Context, kubernetesComponents components.Components, WantedKubernetesVersion string) error {
-	conn, err := grpc.NewClient(mainconstants.UpgradeAgentMountPath, grpc.WithTransportCredentials(insecure.NewCredentials()),
-		grpc.WithContextDialer(
-			func(ctx context.Context, addr string) (net.Conn, error) {
-				return c.dialer.DialContext(ctx, "unix", addr)
-			},
-		))
+	conn, err := grpc.NewClient(c.addr, grpc.WithTransportCredentials(insecure.NewCredentials()))
 	if err != nil {
 		return fmt.Errorf("failed to dial: %w", err)
 	}
diff --git a/operators/constellation-node-operator/internal/upgrade/upgrade_test.go b/operators/constellation-node-operator/internal/upgrade/upgrade_test.go
new file mode 100644
index 000000000..394ffc54d
--- /dev/null
+++ b/operators/constellation-node-operator/internal/upgrade/upgrade_test.go
@@ -0,0 +1,52 @@
+/*
+Copyright (c) Edgeless Systems GmbH
+
+SPDX-License-Identifier: AGPL-3.0-only
+*/
+
+package upgrade
+
+import (
+	"context"
+	"net"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/edgelesssys/constellation/v2/internal/versions/components"
+	"github.com/edgelesssys/constellation/v2/upgrade-agent/upgradeproto"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc"
+)
+
+// TestGRPCDialer is a regression test to ensure the upgrade client can connect to a UDS.
+func TestGRPCDialer(t *testing.T) {
+	require := require.New(t)
+
+	dir := t.TempDir()
+	sockAddr := filepath.Join(dir, "test.socket")
+
+	upgradeAgent := &fakeUpgradeAgent{}
+	grpcServer := grpc.NewServer()
+	upgradeproto.RegisterUpdateServer(grpcServer, upgradeAgent)
+
+	listener, err := net.Listen("unix", sockAddr)
+	require.NoError(err)
+	go grpcServer.Serve(listener)
+	t.Cleanup(grpcServer.Stop)
+
+	fileInfo, err := os.Stat(sockAddr)
+	require.NoError(err)
+	require.Equal(os.ModeSocket, fileInfo.Mode()&os.ModeType)
+
+	upgradeClient := newClientWithAddress(sockAddr)
+	require.NoError(upgradeClient.Upgrade(context.Background(), []*components.Component{}, "v1.29.6"))
+}
+
+type fakeUpgradeAgent struct {
+	upgradeproto.UnimplementedUpdateServer
+}
+
+func (s *fakeUpgradeAgent) ExecuteUpdate(_ context.Context, _ *upgradeproto.ExecuteUpdateRequest) (*upgradeproto.ExecuteUpdateResponse, error) {
+	return &upgradeproto.ExecuteUpdateResponse{}, nil
+}