refactor cluster connection in CLI

cmd/uncloud/machine/add.go

@@ -1,15 +1,14 @@
 package machine
 
 import (
-	"context"
+	"fmt"
 	"github.com/spf13/cobra"
 	"uncloud/internal/cli"
+	"uncloud/internal/cli/config"
 )
 
 type addOptions struct {
 	name    string
-	user    string
-	port    int
 	sshKey  string
 	cluster string
 }
@@ -17,25 +16,34 @@ type addOptions struct {
 func NewAddCommand() *cobra.Command {
 	opts := addOptions{}
 	cmd := &cobra.Command{
-		// TODO: add support for [USER@]HOST[:PORT] syntax
+		Use:   "add [USER@]HOST[:PORT]",
-		Use:   "add HOST",
+		Short: "Add a remote machine to a cluster.",
-		Short: "Add a new machine to a cluster.",
 		Args:  cobra.ExactArgs(1),
 		RunE: func(cmd *cobra.Command, args []string) error {
 			uncli := cmd.Context().Value("cli").(*cli.CLI)
-			return add(cmd.Context(), uncli, args[0], opts)
+
+			user, host, port, err := config.SSHDestination(args[0]).Parse()
+			if err != nil {
+				return fmt.Errorf("parse remote machine: %w", err)
+			}
+			remoteMachine := cli.RemoteMachine{
+				User:    user,
+				Host:    host,
+				Port:    port,
+				KeyPath: opts.sshKey,
+			}
+
+			return uncli.AddMachine(cmd.Context(), remoteMachine, opts.cluster, opts.name)
 		},
 	}
 	cmd.Flags().StringVarP(&opts.name, "name", "n", "", "Assign a name to the machine")
-	cmd.Flags().StringVarP(&opts.user, "user", "u", "root", "Username for SSH remote login")
+	cmd.Flags().StringVarP(
-	cmd.Flags().IntVarP(&opts.port, "port", "p", 22, "Port for SSH remote login")
+		&opts.sshKey, "ssh-key", "i", "",
-	cmd.Flags().StringVarP(&opts.sshKey, "ssh-key", "i", "",
+		"path to SSH private key for SSH remote login (default ~/.ssh/id_*)",
-		"path to SSH private key for SSH remote login (default ~/.ssh/id_*)")
+	)
-	cmd.Flags().StringVarP(&opts.cluster, "cluster", "c", "",
+	cmd.Flags().StringVarP(
-		"Name of the cluster to add the machine to (default is the current cluster)")
+		&opts.cluster, "cluster", "c", "",
+		"Name of the cluster to add the machine to (default is the current cluster)",
+	)
 	return cmd
 }
-
-func add(ctx context.Context, uncli *cli.CLI, host string, opts addOptions) error {
-	return uncli.AddMachine(ctx, opts.cluster, opts.name, opts.user, host, opts.port, opts.sshKey)
-}

internal/cli/cli.go

@@ -32,14 +32,14 @@ func New(configPath string) (*CLI, error) {
 	}, nil
 }
 
-func (cli *CLI) CreateCluster(name string, userPrivateKey secret.Secret) (*client.ClusterClient, error) {
+func (cli *CLI) CreateCluster(name string, userPrivateKey secret.Secret) error {
 	if _, ok := cli.config.Clusters[name]; ok {
-		return nil, fmt.Errorf("cluster %q already exists", name)
+		return fmt.Errorf("cluster %q already exists", name)
 	}
 	if userPrivateKey == nil {
 		user, err := client.NewUser(nil)
 		if err != nil {
-			return nil, fmt.Errorf("generate user: %w", err)
+			return fmt.Errorf("generate user: %w", err)
 		}
 		userPrivateKey = user.PrivateKey()
 	}
@@ -48,42 +48,14 @@ func (cli *CLI) CreateCluster(name string, userPrivateKey secret.Secret) (*clien
 		Name:           name,
 		UserPrivateKey: userPrivateKey,
 	}
-	if err := cli.config.Save(); err != nil {
+	return cli.config.Save()
-		return nil, err
-	}
-
-	return cli.GetCluster(name)
 }
 
-func (cli *CLI) CreateDefaultCluster() (*client.ClusterClient, error) {
+func (cli *CLI) CreateDefaultCluster() error {
-	c, err := cli.CreateCluster("default", nil)
+	if err := cli.CreateCluster(defaultClusterName, nil); err != nil {
-	if err != nil {
+		return err
-		return nil, err
 	}
-	if err = cli.SetCurrentCluster(c.Name()); err != nil {
+	return cli.SetCurrentCluster(defaultClusterName)
-		return nil, err
-	}
-	return c, nil
-}
-
-func (cli *CLI) GetCluster(name string) (*client.ClusterClient, error) {
-	clusterCfg, ok := cli.config.Clusters[name]
-	if !ok {
-		return nil, ErrNotFound
-	}
-	clusterCfg.Name = name
-
-	user, err := client.NewUser(clusterCfg.UserPrivateKey)
-	if err != nil {
-		return nil, fmt.Errorf("create user: %w", err)
-	}
-	wgConnector := connector.NewWireGuardConnector(user, clusterCfg.Connections)
-
-	return client.NewClusterClient(clusterCfg, wgConnector)
-}
-
-func (cli *CLI) GetCurrentCluster() (*client.ClusterClient, error) {
-	return cli.GetCluster(cli.config.CurrentCluster)
 }
 
 func (cli *CLI) SetCurrentCluster(name string) error {
@@ -94,18 +66,6 @@ func (cli *CLI) SetCurrentCluster(name string) error {
 	return cli.config.Save()
 }
 
-func (cli *CLI) ListClusters() ([]*client.ClusterClient, error) {
-	var clusters []*client.ClusterClient
-	for name := range cli.config.Clusters {
-		c, err := cli.GetCluster(name)
-		if err != nil {
-			return nil, fmt.Errorf("get cluster %q: %w", name, err)
-		}
-		clusters = append(clusters, c)
-	}
-	return clusters, nil
-}
-
 func (cli *CLI) InitCluster(
 	ctx context.Context, remoteMachine *RemoteMachine, clusterName, machineName string, netPrefix netip.Prefix,
 ) error {
@@ -164,8 +124,7 @@ func (cli *CLI) initRemoteMachine(
 	}
 	fmt.Printf("Cluster %q initialised with machine %q\n", clusterName, resp.Machine.Name)
 
-	_, err = cli.CreateCluster(clusterName, user.PrivateKey())
+	if err = cli.CreateCluster(clusterName, user.PrivateKey()); err != nil {
-	if err != nil {
 		return fmt.Errorf("save cluster to config: %w", err)
 	}
 	// Set the current cluster to the just created one if it is the only cluster in the config.
@@ -185,59 +144,71 @@ func (cli *CLI) initRemoteMachine(
 	return nil
 }
 
-func (cli *CLI) AddMachine(
+func (cli *CLI) ConnectCluster(ctx context.Context, clusterName string) (*client.Client, error) {
-	ctx context.Context, clusterName, machineName, user, host string, port int, sshKeyPath string,
+	if len(cli.config.Clusters) == 0 {
-) error {
+		return nil, errors.New(
-	var (
+			"no clusters found in the Uncloud config. " +
-		cluster *client.ClusterClient
+				"Please initialise a cluster with `uncloud machine init` first",
-		err     error
 		)
-	if clusterName == "" {
-		// If the cluster is not specified, use the current cluster. If there are no clusters, create a default one.
-		cluster, err = cli.GetCurrentCluster()
-		if err != nil {
-			if errors.Is(err, ErrNotFound) {
-				// Do not create a default cluster if there are already clusters but the current cluster is not set.
-				clusters, cErr := cli.ListClusters()
-				if cErr != nil {
-					return fmt.Errorf("list clusters: %w", cErr)
 	}
-				if len(clusters) > 0 {
+	if clusterName == "" {
-					return errors.New(
+		// If the cluster is not specified, use the current cluster if set.
+		if cli.config.CurrentCluster == "" {
+			return nil, errors.New(
 				"the current cluster is not set in the Uncloud config. " +
 					"Please specify a cluster with the --cluster flag or set current_cluster in the config",
 			)
 		}
+		if _, ok := cli.config.Clusters[cli.config.CurrentCluster]; !ok {
+			return nil, fmt.Errorf(
+				"current cluster %q not found in the config. "+
+					"Please specify a cluster with the --cluster flag or update current_cluster in the config",
+				cli.config.CurrentCluster,
+			)
+		}
+		clusterName = cli.config.CurrentCluster
+	}
 
-				cluster, err = cli.CreateDefaultCluster()
+	cfg, ok := cli.config.Clusters[clusterName]
-				if err != nil {
+	if !ok {
-					return fmt.Errorf("create default cluster: %w", err)
+		return nil, fmt.Errorf("cluster %q not found in the config", clusterName)
 	}
-				fmt.Printf("Created %q cluster\n", cluster.Name())
+	if len(cfg.Connections) == 0 {
-			} else {
+		return nil, fmt.Errorf("no connection configurations found for cluster %q in the config", clusterName)
-				return fmt.Errorf("get current cluster: %w", err)
 	}
-		}
-	} else {
-		cluster, err = cli.GetCluster(clusterName)
-		if err != nil {
-			return fmt.Errorf("get cluster %q: %w", clusterName, err)
-		}
-	}
-	defer func() {
-		_ = cluster.Close()
-	}()
 
-	name, connCfg, err := cluster.AddMachine(ctx, machineName, user, host, port, sshKeyPath)
+	// TODO: iterate over all connections and try to connect to the cluster using the first successful connection.
+	conn := cfg.Connections[0]
+	user, host, port, err := conn.SSH.Parse()
 	if err != nil {
-		return fmt.Errorf("add machine to cluster %q: %w", cluster.Name(), err)
+		return nil, fmt.Errorf("parse SSH connection %q: %w", conn.SSH, err)
 	}
-	fmt.Printf("Machine %q added to cluster %q\n", name, cluster.Name())
+	sshConfig := &connector.SSHConnectorConfig{
+		User: user,
+		Host: host,
+		Port: port,
+	}
+	return client.New(ctx, connector.NewSSHConnector(sshConfig))
+}
 
-	cli.config.Clusters[cluster.Name()].Connections = append(cli.config.Clusters[cluster.Name()].Connections, connCfg)
+func (cli *CLI) AddMachine(ctx context.Context, remoteMachine RemoteMachine, clusterName, machineName string) error {
-	if err = cli.config.Save(); err != nil {
+	c, err := cli.ConnectCluster(ctx, clusterName)
-		return fmt.Errorf("save config: %w", err)
+	if err != nil {
+		return fmt.Errorf("connect to cluster: %w", err)
 	}
 
+	fmt.Println("Adding machine to cluster...", c)
+	// TODO
+	//name, connCfg, err := c.AddMachine(ctx, machineName, user, host, port, sshKeyPath)
+	//if err != nil {
+	//	return fmt.Errorf("add machine to cluster %q: %w", cluster.Name(), err)
+	//}
+	//fmt.Printf("Machine %q added to cluster %q\n", name, cluster.Name())
+	//
+	//cli.config.Clusters[cluster.Name()].Connections = append(cli.config.Clusters[cluster.Name()].Connections, connCfg)
+	//if err = cli.config.Save(); err != nil {
+	//	return fmt.Errorf("save config: %w", err)
+	//}
+
 	return nil
 }

internal/cli/client/client.go

@@ -8,11 +8,13 @@ import (
 	"uncloud/internal/machine/api/pb"
 )
 
+// Client is a client for the machine API.
 type Client struct {
 	connector Connector
 	conn      *grpc.ClientConn
 
 	pb.MachineClient
+	pb.ClusterClient
 }
 
 // Connector is an interface for establishing a connection to the machine API.
@@ -21,6 +23,8 @@ type Connector interface {
 	Close() error
 }
 
+// New creates a new client for the machine API. The connector is used to establish the connection
+// either locally or remotely. The client is responsible for closing the connector.
 func New(ctx context.Context, connector Connector) (*Client, error) {
 	c := &Client{
 		connector: connector,
@@ -32,6 +36,7 @@ func New(ctx context.Context, connector Connector) (*Client, error) {
 	}
 
 	c.MachineClient = pb.NewMachineClient(c.conn)
+	c.ClusterClient = pb.NewClusterClient(c.conn)
 	return c, nil
 }