Add support for forcing socket connection types for plugins.

PiperOrigin-RevId: 874856247

Author: drewhli

internal/cfg/cfg.go

@@ -137,6 +137,7 @@ timeout_in_seconds = 60
 socket_connections_dir = {{.socketConnectionsDir}}
 state_dir = {{.baseStateDir}}
 local_plugin_dir = {{.localPluginDir}}
+connection_type =
 
 [ACS]
 endpoint =
@@ -439,6 +440,10 @@ type Plugin struct {
 	StateDir string `ini:"state_dir,omitempty"`
 	// LocalPluginDir defines the directory path where local plugins are installed.
 	LocalPluginDir string `ini:"local_plugin_dir,omitempty"`
+	// ConnectionType defines the connection type of the plugin. Only "tcp" and
+	// "uds" are supported; anything else will be ignored and the default
+	// behavior will be used.
+	ConnectionType string `ini:"connection_type,omitempty"`
 }
 
 // Unstable contains the configurations of Unstable section. No long term

internal/plugin/manager/pluginlauncher.go

@@ -252,10 +252,10 @@ func address(ctx context.Context, protocol, id string, policy retry.Policy) (str
 	return retry.RunWithResponse(ctx, policy, f)
 }
 
-// isUDSSupported returns true if UDS is supported on Windows. Instead of going
-// by version to figure out if UDS is supported, try listening on test address
-// using UDS, if it gets listener successfully consider UDS is supported.
-func isUDSSupported() bool {
+// defaultIsUDSSupported returns true if UDS is supported on Windows. Instead of
+// going by version to figure out if UDS is supported, try listening on test
+// address using UDS, if it gets listener successfully consider UDS is supported.
+func defaultIsUDSSupported() bool {
 	if runtime.GOOS == "linux" {
 		return true
 	}

internal/plugin/manager/pluginmanager.go

@@ -62,6 +62,10 @@ const (
 var (
 	// pluginManager is the instance of plugin manager.
 	pluginManager *PluginManager
+
+	// isUDSSupported checks if UDS is supported on the host. This is stubbed out
+	// for testing.
+	isUDSSupported = defaultIsUDSSupported
 )
 
 // PluginManager struct represents the plugins that plugin manager manages.
@@ -262,8 +266,21 @@ func InitPluginManager(ctx context.Context, instanceID string) (*PluginManager,
 	pluginManager.pendingPluginRevisionsMu.Unlock()
 	wg.Wait()
 
+	// TCP is always supported, so if the forced connection type is TCP, we can
+	// return early.
+	connType := cfg.Retrieve().Plugin.ConnectionType
+	if connType == tcpProtocol {
+		pluginManager.protocol = tcpProtocol
+		return pluginManager, nil
+	}
+
+	// If the connection type is not TCP, check if UDS is supported and set the
+	// protocol to UDS if it is. Otherwise, fallback to TCP.
 	if isUDSSupported() {
 		pluginManager.protocol = udsProtocol
+	} else {
+		galog.Debugf("UDS is not supported, fallback to TCP")
+		pluginManager.protocol = tcpProtocol
 	}
 	return pluginManager, nil
 }

internal/plugin/manager/pluginmanager_test.go

@@ -422,6 +422,83 @@ func TestInitPluginManager(t *testing.T) {
 	}
 }
 
+func TestInitPluginManagerConfig(t *testing.T) {
+	ctx := context.WithValue(context.Background(), client.OverrideConnection, &fakeACS{})
+	stateDir := t.TempDir()
+	addr := filepath.Join(t.TempDir(), "pluginA_revisionA.sock")
+
+	tests := []struct {
+		name             string
+		isUDSSupported   bool
+		connectionType   string
+		expectedProtocol string
+	}{
+		{
+			name:             "uds_connection_type",
+			isUDSSupported:   true,
+			connectionType:   udsProtocol,
+			expectedProtocol: udsProtocol,
+		},
+		{
+			name:             "tcp_connection_type",
+			isUDSSupported:   true,
+			connectionType:   tcpProtocol,
+			expectedProtocol: tcpProtocol,
+		},
+		{
+			name:             "invalid_connection_type",
+			isUDSSupported:   true,
+			connectionType:   "invalid",
+			expectedProtocol: udsProtocol,
+		},
+		{
+			name:             "uds_not_supported",
+			isUDSSupported:   false,
+			connectionType:   udsProtocol,
+			expectedProtocol: tcpProtocol,
+		},
+		{
+			name:             "uds_not_supported_tcp_connection_type",
+			isUDSSupported:   false,
+			connectionType:   tcpProtocol,
+			expectedProtocol: tcpProtocol,
+		},
+		{
+			name:             "uds_not_supported_invalid_connection_type",
+			isUDSSupported:   false,
+			connectionType:   "invalid",
+			expectedProtocol: tcpProtocol,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			config := fmt.Sprintf("[PluginConfig]\nstate_dir = %s\nconnection_type = %s\n[Core]\nacs_client = false\nsocket_connections_dir = %s", stateDir, tc.connectionType, filepath.Dir(addr))
+			if err := cfg.Load([]byte(config)); err != nil {
+				t.Fatalf("cfg.Load(nil) failed unexpectedly with error: %v", err)
+			}
+
+			oldIsUDSSupported := isUDSSupported
+			t.Cleanup(func() { isUDSSupported = oldIsUDSSupported })
+			isUDSSupported = func() bool { return tc.isUDSSupported }
+
+			pm, err := InitPluginManager(ctx, "1234567890")
+			if err != nil {
+				t.Fatalf("InitPluginManager(ctx) failed unexpectedly with error: %v", err)
+			}
+			if pm.protocol != tc.expectedProtocol {
+				t.Errorf("InitPluginManager(ctx) = protocol %q, want %q", pm.protocol, tc.expectedProtocol)
+			}
+
+			t.Cleanup(func() {
+				if err := command.CurrentMonitor().UnregisterHandler(VMEventCmd); err != nil {
+					t.Fatalf("command.CurrentMonitor().UnregisterHandler(VMEventCmd) failed unexpectedly with error: %v", err)
+				}
+			})
+		})
+	}
+}
+
 func TestConfigurePluginStates(t *testing.T) {
 	if err := cfg.Load(nil); err != nil {
 		t.Fatalf("Failed to load config: %v", err)