update network peers when new machine added

internal/cli/cli.go

@@ -9,6 +9,7 @@ import (
 	"uncloud/internal/cli/client"
 	"uncloud/internal/cli/client/connector"
 	"uncloud/internal/cli/config"
+	"uncloud/internal/machine"
 	"uncloud/internal/machine/api/pb"
 	"uncloud/internal/secret"
 	"uncloud/internal/sshexec"
@@ -225,31 +226,39 @@ func (cli *CLI) AddMachine(ctx context.Context, remoteMachine RemoteMachine, clu
 		return fmt.Errorf("uncloudd binary not found on the remote machine: %w", err)
 	}
 
-	resp, err := machineClient.Token(ctx, &emptypb.Empty{})
+	tokenResp, err := machineClient.Token(ctx, &emptypb.Empty{})
 	if err != nil {
 		return fmt.Errorf("get remote machine token: %w", err)
 	}
-	token := resp.Token
+	token, err := machine.ParseToken(tokenResp.Token)
+	if err != nil {
+		return fmt.Errorf("parse remote machine token: %w", err)
+	}
 
-	fmt.Println("Token:", token)
+	endpoints := make([]*pb.IPPort, len(token.Endpoints))
+	for i, addrPort := range token.Endpoints {
+		endpoints[i] = pb.NewIPPort(addrPort)
+	}
+	addReq := &pb.AddMachineRequest{
+		Name: machineName,
+		Network: &pb.NetworkConfig{
+			Endpoints: endpoints,
+			PublicKey: token.PublicKey,
+		},
+	}
+	addResp, err := c.AddMachine(ctx, addReq)
+	if err != nil {
+		return fmt.Errorf("add machine to cluster: %w", err)
+	}
 
-	//req := &pb.AddMachineRequest{
+	fmt.Println("Machine added to cluster", addResp.Machine)
-	//	Name: machineName,
-	//	Network: &pb.NetworkConfig{
-	//
-	//	}
-	//}
-	//resp, err := c.AddMachine(ctx, req)
-	//if err != nil {
-	//	return fmt.Errorf("add machine to cluster: %w", err)
-	//}
 
 	// TODO:
 	// --1. Establish a client connection to the remote machine.
 	// --2. Check if the machine is already provisioned and ask the user to reset it first.
 	// --3. Download and install the latest uncloudd binary by running the install shell script from GitHub.
-	// 4. Request token from the remote machine.
+	// --4. Request token from the remote machine.
-	// 5. Add the machine to the cluster using its token and receive a configuration token.
+	// 5. Add the machine to the cluster using its token and receive the added machine info.
 	// 6. Request the machine to join the cluster using the configuration token.
 	// 7. Save the machine's SSH connection details in the cluster config.
 

internal/machine/api/pb/cluster.go

@@ -0,0 +1,49 @@
+package pb
+
+import (
+	"fmt"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
+// KeyLen is the expected key length for a WireGuard public or private key.
+const KeyLen = 32
+
+func (c *NetworkConfig) Validate() error {
+	if c.Subnet != nil {
+		_, err := c.Subnet.ToPrefix()
+		if err != nil {
+			return status.Errorf(codes.InvalidArgument, "invalid subnet: %v", err)
+		}
+	}
+	if c.ManagementIp != nil {
+		_, err := c.ManagementIp.ToAddr()
+		if err != nil {
+			return status.Errorf(codes.InvalidArgument, "invalid management IP: %v", err)
+		}
+	}
+	for _, ep := range c.Endpoints {
+		if ep == nil {
+			return status.Error(codes.InvalidArgument, "endpoint not set")
+		}
+		_, err := ep.ToAddrPort()
+		if err != nil {
+			return status.Errorf(codes.InvalidArgument, "invalid endpoint: %v", err)
+		}
+	}
+	if c.PublicKey == nil {
+		return status.Error(codes.InvalidArgument, "public key not set")
+	}
+	if len(c.PublicKey) != KeyLen {
+		return status.Errorf(codes.InvalidArgument, "invalid public key length: %d", len(c.PublicKey))
+	}
+
+	return nil
+}
+
+func (r *AddMachineRequest) Validate() error {
+	if r.Network == nil {
+		return fmt.Errorf("network not set")
+	}
+	return r.Network.Validate()
+}

internal/machine/api/pb/common.go

@@ -17,6 +17,9 @@ func (ip *IP) ToAddr() (netip.Addr, error) {
 	if err := addr.UnmarshalBinary(ip.Ip); err != nil {
 		return netip.Addr{}, fmt.Errorf("unmarshal IP: %w", err)
 	}
+	if !addr.IsValid() {
+		return netip.Addr{}, fmt.Errorf("invalid IP")
+	}
 	return addr, nil
 }
 
@@ -45,5 +48,9 @@ func (p *IPPrefix) ToPrefix() (netip.Prefix, error) {
 	if err != nil {
 		return netip.Prefix{}, err
 	}
-	return netip.PrefixFrom(addr, int(p.Bits)), nil
+	prefix := netip.PrefixFrom(addr, int(p.Bits))
+	if !prefix.IsValid() {
+		return netip.Prefix{}, fmt.Errorf("invalid prefix")
+	}
+	return prefix, nil
 }

internal/machine/cluster/cluster.go

@@ -7,20 +7,26 @@ import (
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/proto"
+	"log/slog"
 	"net/netip"
 	"uncloud/internal/machine/api/pb"
 	"uncloud/internal/machine/network"
+	"uncloud/internal/secret"
 )
 
 type Cluster struct {
 	pb.UnimplementedClusterServer
 
 	state *State
+
+	// TODO: temporary channel until the state is replaced with networkDB.
+	newMachinesCh chan *pb.MachineInfo
 }
 
 func NewCluster(state *State) *Cluster {
 	return &Cluster{
-		state: state,
+		state:         state,
+		newMachinesCh: make(chan *pb.MachineInfo),
 	}
 }
 
@@ -54,15 +60,22 @@ func (c *Cluster) SetNetwork(network *pb.IPPrefix) error {
 	return nil
 }
 
+// TODO: this is a temporary watcher for PoC until the state is state is replaced with networkDB.
+func (c *Cluster) WatchNewMachines() <-chan *pb.MachineInfo {
+	return c.newMachinesCh
+}
+
 // AddMachine adds a machine to the cluster.
 func (c *Cluster) AddMachine(ctx context.Context, req *pb.AddMachineRequest) (*pb.AddMachineResponse, error) {
 	if c.state == nil {
 		return nil, status.Error(codes.FailedPrecondition, "cluster is not initialized")
 	}
 
-	// TODO: replace errors with gRPC status.Error(f), e.g. status.Error(codes.InvalidArgument, "management IP not set")
+	if err := req.Validate(); err != nil {
-	if req.Network.PublicKey == nil {
+		return nil, err
-		return nil, fmt.Errorf("public key not set")
+	}
+	if len(req.Network.Endpoints) == 0 {
+		return nil, status.Error(codes.InvalidArgument, "endpoints not set")
 	}
 
 	machines := c.state.State.Machines
@@ -71,13 +84,16 @@ func (c *Cluster) AddMachine(ctx context.Context, req *pb.AddMachineRequest) (*p
 	i := 0
 	for _, m := range machines {
 		if req.Name != "" && m.Name == req.Name {
-			return nil, fmt.Errorf("machine with name %q already exists", req.Name)
+			return nil, status.Errorf(codes.AlreadyExists, "machine with name %q already exists", req.Name)
 		}
-		if m.Network.ManagementIp != nil && m.Network.ManagementIp.Equal(req.Network.ManagementIp) {
+		if m.Network.ManagementIp != nil && req.Network.ManagementIp != nil &&
-			return nil, fmt.Errorf("machine with management IP %q already exists", req.Network.ManagementIp)
+			m.Network.ManagementIp.Equal(req.Network.ManagementIp) {
+			manageIP, _ := req.Network.ManagementIp.ToAddr()
+			return nil, status.Errorf(codes.AlreadyExists, "machine with management IP %q already exists", manageIP)
 		}
 		if bytes.Equal(m.Network.PublicKey, req.Network.PublicKey) {
-			return nil, fmt.Errorf("machine with public key %q already exists", req.Network.PublicKey)
+			publicKey := secret.Secret(m.Network.PublicKey)
+			return nil, status.Errorf(codes.AlreadyExists, "machine with public key %q already exists", publicKey)
 		}
 
 		if allocatedSubnets[i], err = m.Network.Subnet.ToPrefix(); err != nil {
@@ -88,7 +104,7 @@ func (c *Cluster) AddMachine(ctx context.Context, req *pb.AddMachineRequest) (*p
 
 	mid, err := NewMachineID()
 	if err != nil {
-		return nil, fmt.Errorf("generate machine ID: %w", err)
+		return nil, status.Errorf(codes.Internal, "generate machine ID: %v", err)
 	}
 	manageIP := req.Network.ManagementIp
 	if manageIP == nil {
@@ -105,39 +121,44 @@ func (c *Cluster) AddMachine(ctx context.Context, req *pb.AddMachineRequest) (*p
 	if m.Name == "" {
 		m.Name, err = NewRandomMachineName()
 		if err != nil {
-			return nil, fmt.Errorf("generate machine name: %w", err)
+			return nil, status.Errorf(codes.Internal, "generate machine name: %v", err)
 		}
 	}
 
 	clusterNetwork, err := c.Network()
 	if err != nil {
-		return nil, fmt.Errorf("get cluster network: %w", err)
+		return nil, status.Errorf(codes.Internal, "get cluster network: %v", err)
 	}
 	ipam, err := NewIPAMWithAllocated(clusterNetwork, allocatedSubnets)
 	if err != nil {
-		return nil, fmt.Errorf("create IPAM manager: %w", err)
+		return nil, status.Errorf(codes.Internal, "create IPAM manager: %v", err)
 	}
 	subnet, err := ipam.AllocateSubnetLen(network.DefaultSubnetBits)
 	if err != nil {
-		return nil, fmt.Errorf("allocate subnet: %w", err)
+		return nil, status.Errorf(codes.Internal, "allocate subnet for machine: %v", err)
 	}
 	m.Network.Subnet = pb.NewIPPrefix(subnet)
 
 	// TODO: announce the new machine to the cluster network using Serf and achieve consensus.
-	mState := proto.Clone(m).(*pb.MachineInfo)
+	mCopy := proto.Clone(m).(*pb.MachineInfo)
-	c.state.State.Machines[mState.Id] = mState
+	c.state.State.Machines[mCopy.Id] = mCopy
 	// Store machine endpoints in a separate collection to allow modifying them with limited permissions.
-	c.state.State.Endpoints[mState.Id] = &pb.MachineEndpoints{
+	c.state.State.Endpoints[mCopy.Id] = &pb.MachineEndpoints{
-		Id:        mState.Id,
+		Id:        mCopy.Id,
 		Endpoints: req.Network.Endpoints,
 	}
 	if err = c.state.Save(); err != nil {
 		return nil, status.Errorf(codes.Internal, "save state: %v", err)
 	}
+	slog.Info("Machine added to the cluster.", "id", m.Id, "name", m.Name)
 
 	// Include the machine endpoints in the response.
 	m.Network.Endpoints = req.Network.Endpoints
 
+	// TODO: notify all cluster machines about the new machine so they can update their peers config.
+	//  In PoC we just notify the local machine.
+	c.newMachinesCh <- m
+
 	resp := &pb.AddMachineResponse{Machine: m}
 	return resp, nil
 }

internal/machine/machine.go

@@ -44,8 +44,9 @@ type Machine struct {
 	localServer   *grpc.Server
 	networkServer *grpc.Server
 
-	clusterState *cluster.State
+	clusterState  *cluster.State
-	cluster      *cluster.Cluster
+	cluster       *cluster.Cluster
+	newMachinesCh <-chan *pb.MachineInfo
 }
 
 func NewMachine(config *Config) (*Machine, error) {
@@ -101,6 +102,7 @@ func NewMachine(config *Config) (*Machine, error) {
 	}
 	pb.RegisterClusterServer(m.localServer, m.cluster)
 	pb.RegisterClusterServer(m.networkServer, m.cluster)
+	m.newMachinesCh = m.cluster.WatchNewMachines()
 
 	if m.IsInitialised() {
 		m.initialised <- struct{}{}
@@ -199,6 +201,51 @@ func (m *Machine) Run(ctx context.Context) error {
 		},
 	)
 
+	// Handle new machines added to the cluster.
+	errGroup.Go(
+		func() error {
+			for {
+				select {
+				case machineInfo := <-m.newMachinesCh:
+					slog.Info("Handling new machine added to the cluster.", "name", machineInfo.Name)
+					if err := machineInfo.Network.Validate(); err != nil {
+						slog.Error("Invalid machine network configuration.", "err", err)
+						continue
+					}
+					// Ignore errors as they are already validated.
+					subnet, _ := machineInfo.Network.Subnet.ToPrefix()
+					manageIP, _ := machineInfo.Network.ManagementIp.ToAddr()
+					endpoints := make([]netip.AddrPort, len(machineInfo.Network.Endpoints))
+					for i, ep := range machineInfo.Network.Endpoints {
+						addrPort, _ := ep.ToAddrPort()
+						endpoints[i] = addrPort
+					}
+
+					peer := network.PeerConfig{
+						Subnet:       &subnet,
+						ManagementIP: manageIP,
+						AllEndpoints: endpoints,
+						PublicKey:    machineInfo.Network.PublicKey,
+					}
+					if len(endpoints) > 0 {
+						peer.Endpoint = &endpoints[0]
+					}
+
+					m.state.Network.Peers = append(m.state.Network.Peers, peer)
+					if err := m.state.Save(); err != nil {
+						return fmt.Errorf("save machine state: %w", err)
+					}
+
+					if err := m.configureNetwork(); err != nil {
+						return fmt.Errorf("configure network with new peer: %w", err)
+					}
+				case <-ctx.Done():
+					return nil
+				}
+			}
+		},
+	)
+
 	// Shutdown goroutine.
 	errGroup.Go(
 		func() error {