Add private-join and private-leave commands for Private AKS cluster

Add current node into the private cluster.

  - Add private-join command to join Private AKS cluster via Gateway
  - Add private-leave command with --mode=local|full cleanup options
  - Add shell scripts for installation and uninstallation
  - Add documentation for creating Private AKS cluster
  - Readme.md and create_private_clsuter.md shows howto

Author: weiliu2

commands.go

@@ -14,6 +14,7 @@ import (
 	"go.goms.io/aks/AKSFlexNode/pkg/bootstrapper"
 	"go.goms.io/aks/AKSFlexNode/pkg/config"
 	"go.goms.io/aks/AKSFlexNode/pkg/logger"
+	"go.goms.io/aks/AKSFlexNode/pkg/privatecluster"
 	"go.goms.io/aks/AKSFlexNode/pkg/status"
 )
 
@@ -118,6 +119,107 @@ func runVersion() {
 	fmt.Printf("Build Time: %s\n", BuildTime)
 }
 
+// Private cluster command variables
+var (
+	aksResourceID   string
+	cleanupModeFlag string
+)
+
+// NewPrivateJoinCommand creates a new private-join command
+func NewPrivateJoinCommand() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "private-join",
+		Short: "Join a Private AKS cluster (requires sudo)",
+		Long: `Join a Private AKS cluster.
+
+Prerequisites:
+  1. A Private AKS cluster must exist with AAD and Azure RBAC enabled
+     See: pkg/privatecluster/create_private_cluster.md
+
+  2. Current user must have the following roles on the cluster:
+     - Azure Kubernetes Service Cluster Admin Role
+     - Azure Kubernetes Service RBAC Cluster Admin
+
+  3. Current user must be logged in via 'sudo az login'
+
+The full resource ID of the Private AKS cluster is required as the --aks-resource-id parameter.
+This same resource ID can be used later with the private-leave command.`,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			return runPrivateJoin(cmd.Context())
+		},
+	}
+
+	cmd.Flags().StringVar(&aksResourceID, "aks-resource-id", "", "AKS cluster resource ID (required)")
+	cmd.MarkFlagRequired("aks-resource-id")
+
+	return cmd
+}
+
+// NewPrivateLeaveCommand creates a new private-leave command
+func NewPrivateLeaveCommand() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "private-leave",
+		Short: "Leave a Private AKS cluster (--mode=local|full, requires sudo)",
+		Long: `Remove this edge node from a Private AKS cluster.
+
+Cleanup modes:
+  --local  Local cleanup only (default):
+           - Remove node from AKS cluster
+           - Run aks-flex-node unbootstrap
+           - Remove Arc Agent
+           - Stop VPN and remove client config
+           - Keep Gateway for other nodes
+
+  --full   Full cleanup (requires --aks-resource-id):
+           - All local cleanup steps
+           - Delete Gateway VM
+           - Delete Gateway subnet, NSG, Public IP
+           - Delete SSH keys
+
+This command requires the current user to be logged in via 'sudo az login'.`,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			return runPrivateLeave(cmd.Context())
+		},
+	}
+
+	cmd.Flags().StringVar(&cleanupModeFlag, "mode", "local", "Cleanup mode: 'local' (keep Gateway) or 'full' (remove all Azure resources)")
+	cmd.Flags().StringVar(&aksResourceID, "aks-resource-id", "", "AKS cluster resource ID (required for --mode=full)")
+
+	return cmd
+}
+
+// runPrivateJoin executes the private cluster join process
+func runPrivateJoin(ctx context.Context) error {
+	if os.Getuid() != 0 {
+		return fmt.Errorf("this command requires root privileges, please run with 'sudo'")
+	}
+	runner := privatecluster.NewScriptRunner("")
+	return runner.RunPrivateInstall(ctx, aksResourceID)
+}
+
+// runPrivateLeave executes the private cluster leave process
+func runPrivateLeave(ctx context.Context) error {
+	if os.Getuid() != 0 {
+		return fmt.Errorf("this command requires root privileges, please run with 'sudo'")
+	}
+	// Validate cleanup mode
+	var mode privatecluster.CleanupMode
+	switch cleanupModeFlag {
+	case "local":
+		mode = privatecluster.CleanupModeLocal
+	case "full":
+		mode = privatecluster.CleanupModeFull
+		if aksResourceID == "" {
+			return fmt.Errorf("--aks-resource-id is required for full cleanup mode")
+		}
+	default:
+		return fmt.Errorf("invalid cleanup mode: %s (use 'local' or 'full')", cleanupModeFlag)
+	}
+
+	runner := privatecluster.NewScriptRunner("")
+	return runner.RunPrivateUninstall(ctx, mode, aksResourceID)
+}
+
 // runDaemonLoop runs the periodic status collection and bootstrap monitoring daemon
 func runDaemonLoop(ctx context.Context, cfg *config.Config) error {
 	logger := logger.GetLoggerFromContext(ctx)

go.mod

@@ -15,6 +15,8 @@ require (
 	github.com/sirupsen/logrus v1.9.3
 	github.com/spf13/cobra v1.8.0
 	github.com/spf13/viper v1.18.2
+	github.com/stretchr/testify v1.11.1
+	golang.org/x/crypto v0.45.0
 	k8s.io/client-go v0.26.0
 )
 
@@ -40,6 +42,7 @@ require (
 	github.com/modern-go/reflect2 v1.0.3-0.20250322232337-35a7c28c31ee // indirect
 	github.com/pelletier/go-toml/v2 v2.1.0 // indirect
 	github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c // indirect
+	github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
 	github.com/sagikazarmark/locafero v0.4.0 // indirect
 	github.com/sagikazarmark/slog-shim v0.1.0 // indirect
 	github.com/sourcegraph/conc v0.3.0 // indirect
@@ -50,7 +53,6 @@ require (
 	go.uber.org/atomic v1.9.0 // indirect
 	go.uber.org/multierr v1.9.0 // indirect
 	go.yaml.in/yaml/v2 v2.4.3 // indirect
-	golang.org/x/crypto v0.45.0 // indirect
 	golang.org/x/exp v0.0.0-20240325151524-a685a6edb6d8 // indirect
 	golang.org/x/net v0.47.0 // indirect
 	golang.org/x/oauth2 v0.30.0 // indirect

main.go

@@ -25,13 +25,16 @@ func main() {
 	}
 
 	// Add global flags for configuration
-	rootCmd.PersistentFlags().StringVar(&configPath, "config", "", "Path to configuration JSON file (required)")
+	rootCmd.PersistentFlags().StringVar(&configPath, "config", "", "Path to configuration JSON file (required for agent/unbootstrap)")
+	rootCmd.PersistentFlags().MarkHidden("config") // Hide from global help, shown in agent/unbootstrap help
 	// Don't mark as required globally - we'll check in PersistentPreRunE for commands that need it
 
 	// Add commands
 	rootCmd.AddCommand(NewAgentCommand())
 	rootCmd.AddCommand(NewUnbootstrapCommand())
 	rootCmd.AddCommand(NewVersionCommand())
+	rootCmd.AddCommand(NewPrivateJoinCommand())
+	rootCmd.AddCommand(NewPrivateLeaveCommand())
 
 	// Set up context with signal handling
 	ctx, cancel := context.WithCancel(context.Background())
@@ -49,8 +52,9 @@ func main() {
 
 	// Set up persistent pre-run to initialize config and logger
 	rootCmd.PersistentPreRunE = func(cmd *cobra.Command, args []string) error {
-		// Skip config loading for version command
-		if cmd.Name() == "version" {
+		// Skip config loading for commands that don't need it
+		switch cmd.Name() {
+		case "version", "private-join", "private-leave":
 			return nil
 		}
 

pkg/bootstrapper/bootstrapper.go

@@ -32,6 +32,7 @@ func New(cfg *config.Config, logger *logrus.Logger) *Bootstrapper {
 // Bootstrap executes all bootstrap steps sequentially
 func (b *Bootstrapper) Bootstrap(ctx context.Context) (*ExecutionResult, error) {
 	// Define the bootstrap steps in order - using modules directly
+	// Note: Network (WireGuard VPN) setup is handled by "aks-flex-node private-join" before this runs
 	steps := []Executor{
 		arc.NewInstaller(b.logger),                  // Setup Arc
 		services.NewUnInstaller(b.logger),           // Stop kubelet before setup
@@ -60,6 +61,7 @@ func (b *Bootstrapper) Unbootstrap(ctx context.Context) (*ExecutionResult, error
 		runc.NewUnInstaller(b.logger),                 // Uninstall runc binary
 		system_configuration.NewUnInstaller(b.logger), // Clean system settings
 		arc.NewUnInstaller(b.logger),                  // Uninstall Arc (after cleanup)
+		// Note: Network (WireGuard VPN) cleanup is handled by "aks-flex-node private-leave"
 	}
 
 	return b.ExecuteSteps(ctx, steps, "unbootstrap")

pkg/components/arc/arc_installer.go

@@ -91,7 +91,7 @@ func (i *Installer) Execute(ctx context.Context) error {
 
 // IsCompleted checks if Arc setup has been completed
 // This can be used by bootstrap steps to verify completion status
-// Uses the same reliable logic as status collector for consistency
+// Checks both Arc agent connection status AND RBAC role assignments
 func (i *Installer) IsCompleted(ctx context.Context) bool {
 	i.logger.Debug("Checking Arc setup completion status")
 
@@ -113,26 +113,60 @@ func (i *Installer) IsCompleted(ctx context.Context) bool {
 	}
 
 	// Parse output to check if agent is connected (same logic as status collector)
+	isConnected := false
 	lines := strings.Split(strings.TrimSpace(string(output)), "\n")
 	for _, line := range lines {
 		line = strings.TrimSpace(line)
 		if strings.Contains(line, "Agent Status") && strings.Contains(line, ":") {
 			parts := strings.SplitN(line, ":", 2)
 			if len(parts) == 2 {
 				status := strings.TrimSpace(parts[1])
-				isConnected := strings.ToLower(status) == "connected"
-				if isConnected {
-					i.logger.Debug("Arc setup appears to be completed - agent is connected")
-				} else {
+				isConnected = strings.ToLower(status) == "connected"
+				if !isConnected {
 					i.logger.Debugf("Arc agent status is '%s' - not ready", status)
+					return false
 				}
-				return isConnected
 			}
 		}
 	}
 
-	i.logger.Debug("Could not find Agent Status in azcmagent show output - Arc not ready")
-	return false
+	if !isConnected {
+		i.logger.Debug("Could not find Agent Status in azcmagent show output - Arc not ready")
+		return false
+	}
+
+	// Arc is connected, now check if RBAC roles are assigned
+	i.logger.Debug("Arc agent is connected, checking RBAC role assignments...")
+	if err := i.setUpClients(ctx); err != nil {
+		i.logger.Debugf("Failed to set up clients for RBAC check: %v", err)
+		return false
+	}
+
+	arcMachine, err := i.getArcMachine(ctx)
+	if err != nil || arcMachine == nil {
+		i.logger.Debugf("Failed to get Arc machine: %v", err)
+		return false
+	}
+
+	managedIdentityID := getArcMachineIdentityID(arcMachine)
+	if managedIdentityID == "" {
+		i.logger.Debug("Arc machine has no managed identity ID")
+		return false
+	}
+
+	hasPermissions, err := i.checkRequiredPermissions(ctx, managedIdentityID)
+	if err != nil {
+		i.logger.Debugf("Failed to check RBAC permissions: %v", err)
+		return false
+	}
+
+	if !hasPermissions {
+		i.logger.Debug("Arc is connected but RBAC roles are not assigned")
+		return false
+	}
+
+	i.logger.Debug("Arc setup is complete - agent connected and RBAC roles assigned")
+	return true
 }
 
 // registerArcMachine registers the machine with Azure Arc using the Arc agent

pkg/config/config.go

@@ -240,9 +240,70 @@ func (c *Config) Validate() error {
 		return fmt.Errorf("invalid agent.logLevel: %s. Valid values are: debug, info, warning, error", c.Agent.LogLevel)
 	}
 
+	// Validate network configuration if present
+	if c.Network != nil {
+		if err := c.Network.Validate(); err != nil {
+			return fmt.Errorf("invalid network config: %w", err)
+		}
+	}
+
+	return nil
+}
+
+// Validate validates the network configuration
+func (n *NetworkConfig) Validate() error {
+	switch n.Mode {
+	case "", "direct":
+		// No additional validation needed for direct mode
+		return nil
+	case "wireguard":
+		if n.Gateway == nil {
+			return fmt.Errorf("gateway configuration is required when mode is 'wireguard'")
+		}
+		return n.Gateway.Validate()
+	default:
+		return fmt.Errorf("unsupported network mode: %s. Valid values are: direct, wireguard", n.Mode)
+	}
+}
+
+// Validate validates the Gateway configuration
+func (w *GatewayConfig) Validate() error {
+	if w.ServerEndpoint == "" {
+		return fmt.Errorf("serverEndpoint is required")
+	}
+	if w.ServerPublicKey == "" {
+		return fmt.Errorf("serverPublicKey is required")
+	}
+	if w.ClientAddress == "" {
+		return fmt.Errorf("clientAddress is required")
+	}
+	if len(w.AllowedIPs) == 0 {
+		return fmt.Errorf("allowedIPs is required and must not be empty")
+	}
 	return nil
 }
 
+// GetInterfaceName returns the Gateway interface name, defaulting to "wg-aks"
+func (w *GatewayConfig) GetInterfaceName() string {
+	if w.InterfaceName != "" {
+		return w.InterfaceName
+	}
+	return "wg-aks"
+}
+
+// GetPersistentKeepalive returns the keepalive interval, defaulting to 25 seconds
+func (w *GatewayConfig) GetPersistentKeepalive() int {
+	if w.PersistentKeepalive > 0 {
+		return w.PersistentKeepalive
+	}
+	return 25
+}
+
+// IsGatewayEnabled returns true if Gateway VPN is configured
+func (c *Config) IsGatewayEnabled() bool {
+	return c.Network != nil && c.Network.Mode == "wireguard" && c.Network.Gateway != nil
+}
+
 // populateTargetClusterInfoFromConfig extracts cluster information from the resource ID
 // This function should only be called after validateAzureResourceID confirms the format is correct
 func populateTargetClusterInfoFromConfig(cfg *Config) {

pkg/config/structs.go

@@ -14,6 +14,48 @@ type Config struct {
 	Node       NodeConfig       `json:"node"`
 	Paths      PathsConfig      `json:"paths"`
 	Npd        NPDConfig        `json:"npd"`
+	Network    *NetworkConfig   `json:"network,omitempty"` // Network configuration for Private Cluster access
+}
+
+// NetworkConfig holds network configuration for accessing Private AKS clusters.
+type NetworkConfig struct {
+	// Mode specifies the network mode: "direct" (default) or "wireguard"
+	Mode string `json:"mode,omitempty"`
+
+	// Gateway configuration (required when mode is "wireguard")
+	// Note: JSON/mapstructure field name kept as "wireguard" for backward compatibility
+	Gateway *GatewayConfig `json:"wireguard,omitempty" mapstructure:"wireguard"`
+}
+
+// GatewayConfig holds Gateway VPN configuration for Private Cluster access.
+type GatewayConfig struct {
+	// ServerEndpoint is the Gateway server endpoint (host:port)
+	ServerEndpoint string `json:"serverEndpoint"`
+
+	// ServerPublicKey is the server's Gateway public key (base64 encoded)
+	ServerPublicKey string `json:"serverPublicKey"`
+
+	// ClientPrivateKey is the client's Gateway private key (optional, auto-generated if empty)
+	ClientPrivateKey string `json:"clientPrivateKey,omitempty"`
+
+	// ClientAddress is the client's VPN address with CIDR (e.g., "172.16.0.2/24")
+	ClientAddress string `json:"clientAddress"`
+
+	// AllowedIPs are the IP ranges to route through the VPN tunnel
+	AllowedIPs []string `json:"allowedIPs"`
+
+	// DNS servers to use through the VPN (optional)
+	DNS []string `json:"dns,omitempty"`
+
+	// PersistentKeepalive interval in seconds (default: 25)
+	PersistentKeepalive int `json:"persistentKeepalive,omitempty"`
+
+	// InterfaceName is the Gateway interface name (default: wg-aks)
+	InterfaceName string `json:"interfaceName,omitempty"`
+
+	// TestEndpoint is an endpoint to test connectivity after VPN is established (optional)
+	// If not specified, will try to test connectivity to the AKS API Server
+	TestEndpoint string `json:"testEndpoint,omitempty"`
 }
 
 // AzureConfig holds Azure-specific configuration required for connecting to Azure services.