feat: add machine gateways

valyentdev · Jan 11, 2025 · 2f85eba · 2f85eba
1 parent 4f0bd66
commit 2f85eba
Show file tree

Hide file tree

Showing 27 changed files with 495 additions and 192 deletions.
diff --git a/agent/client/machines.go b/agent/client/machines.go
@@ -105,3 +105,21 @@ func (a *AgentClient) GetMachineLogsRaw(ctx context.Context, id string, follow b
 	}
 	return a.client.RawGet(ctx, path)
 }
+
+func (a *AgentClient) DisableMachineGateway(ctx context.Context, id string) error {
+	path := "/machines/" + id + "/gateway/disable"
+	err := a.client.Post(ctx, path, nil)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func (a *AgentClient) EnableMachineGateway(ctx context.Context, id string) error {
+	path := "/machines/" + id + "/gateway/enable"
+	err := a.client.Post(ctx, path, nil)
+	if err != nil {
+		return err
+	}
+	return nil
+}
diff --git a/agent/machine.go b/agent/machine.go
@@ -65,10 +65,11 @@ func (a *Agent) PutMachine(ctx context.Context, opt cluster.PutMachineOptions) (
 		Machine: opt.Machine,
 		Version: opt.Version,
 		State: structs.MachineInstanceState{
-			DesiredStatus: desiredStatus,
-			Status:        api.MachineStatusCreated,
-			CreatedAt:     time.Now(),
-			UpdatedAt:     time.Now(),
+			DesiredStatus:         desiredStatus,
+			Status:                api.MachineStatusCreated,
+			CreatedAt:             time.Now(),
+			UpdatedAt:             time.Now(),
+			MachineGatewayEnabled: opt.EnableGateway,
 		},
 	}
 
@@ -139,3 +140,21 @@ func (d *Agent) MachineExec(ctx context.Context, id string, cmd []string, timeou
 
 	return machine.Exec(ctx, cmd, timeout)
 }
+
+func (d *Agent) EnableMachineGateway(ctx context.Context, id string) error {
+	machine, err := d.machines.GetMachine(id)
+	if err != nil {
+		return err
+	}
+
+	return machine.EnableGateway(ctx)
+}
+
+func (d *Agent) DisableMachineGateway(ctx context.Context, id string) error {
+	machine, err := d.machines.GetMachine(id)
+	if err != nil {
+		return err
+	}
+
+	return machine.DisableGateway(ctx)
+}
diff --git a/agent/machinerunner/machine_gateway.go b/agent/machinerunner/machine_gateway.go
@@ -0,0 +1,33 @@
+package machinerunner
+
+import "context"
+
+func (m *MachineRunner) EnableGateway(ctx context.Context) error {
+	m.mutex.Lock()
+	defer m.mutex.Unlock()
+
+	if m.state.MachineInstance().State.MachineGatewayEnabled {
+		return nil
+	}
+
+	if err := m.state.PushGatewayEvent(true); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (m *MachineRunner) DisableGateway(ctx context.Context) error {
+	m.mutex.Lock()
+	defer m.mutex.Unlock()
+
+	if !m.state.MachineInstance().State.MachineGatewayEnabled {
+		return nil
+	}
+
+	if err := m.state.PushGatewayEvent(false); err != nil {
+		return err
+	}
+
+	return nil
+}
diff --git a/agent/machinerunner/state/events.go b/agent/machinerunner/state/events.go
@@ -323,3 +323,27 @@ func (s *MachineInstanceState) PushDestroyedEvent(ctx context.Context) (err erro
 
 	return nil
 }
+
+func (s *MachineInstanceState) PushGatewayEvent(enabled bool) (err error) {
+	s.mutex.Lock()
+	defer s.mutex.Unlock()
+
+	event := s.newEvent(
+		api.MachineGateway,
+		api.OriginUser,
+		s.mi.State.Status,
+		api.MachineEventPayload{
+			Gateway: &api.MachineGatewayEventPayload{
+				Enabled: enabled,
+			},
+		},
+	)
+
+	s.mi.State.MachineGatewayEnabled = enabled
+
+	if err = s.persistStateAndEvent(event); err != nil {
+		return
+	}
+
+	return nil
+}
diff --git a/agent/server/machines.go b/agent/server/machines.go
@@ -130,3 +130,35 @@ func (s *AgentServer) getMachineLogs(ctx context.Context, req *GetMachineLogsReq
 	}
 	return &GetMachineLogsResponse{Body: logs}, nil
 }
+
+type EnableMachineGatewayRequest struct {
+	Id string `path:"id"`
+}
+
+type EnableMachineGatewayResponse struct {
+}
+
+func (s *AgentServer) enableMachineGateway(ctx context.Context, req *EnableMachineGatewayRequest) (*EnableMachineGatewayResponse, error) {
+	err := s.agent.EnableMachineGateway(ctx, req.Id)
+	if err != nil {
+		s.log("Failed to enable machine gateway", err)
+		return nil, err
+	}
+	return &EnableMachineGatewayResponse{}, nil
+}
+
+type DisableMachineGatewayRequest struct {
+	Id string `path:"id"`
+}
+
+type DisableMachineGatewayResponse struct {
+}
+
+func (s *AgentServer) disableMachineGateway(ctx context.Context, req *DisableMachineGatewayRequest) (*DisableMachineGatewayResponse, error) {
+	err := s.agent.DisableMachineGateway(ctx, req.Id)
+	if err != nil {
+		s.log("Failed to disable machine gateway", err)
+		return nil, err
+	}
+	return &DisableMachineGatewayResponse{}, nil
+}
diff --git a/agent/server/routes.go b/agent/server/routes.go
@@ -73,4 +73,16 @@ func (s AgentServer) registerEndpoints(mux humago.Mux) {
 		Method:      http.MethodDelete,
 	}, s.destroyMachine)
 
+	huma.Register(api, huma.Operation{
+		OperationID: "enableMachineGateway",
+		Path:        "/machines/{id}/gateway/enable",
+		Method:      http.MethodPost,
+	}, s.enableMachineGateway)
+
+	huma.Register(api, huma.Operation{
+		OperationID: "disableMachineGateway",
+		Path:        "/machines/{id}/gateway/disable",
+		Method:      http.MethodPost,
+	}, s.disableMachineGateway)
+
 }
diff --git a/agent/structs/machine_instance.go b/agent/structs/machine_instance.go
@@ -9,13 +9,14 @@ import (
 )
 
 type MachineInstanceState struct {
-	DesiredStatus api.MachineStatus  `json:"desired_status"`
-	Status        api.MachineStatus  `json:"status"`
-	Restarts      int                `json:"restarts"`
-	CreatedAt     time.Time          `json:"created_at"`
-	UpdatedAt     time.Time          `json:"updated_at"`
-	LocalIPV4     string             `json:"local_ipv4"`
-	LastEvents    []api.MachineEvent `json:"last_events"`
+	DesiredStatus         api.MachineStatus  `json:"desired_status"`
+	Status                api.MachineStatus  `json:"status"`
+	Restarts              int                `json:"restarts"`
+	CreatedAt             time.Time          `json:"created_at"`
+	UpdatedAt             time.Time          `json:"updated_at"`
+	LocalIPV4             string             `json:"local_ipv4"`
+	LastEvents            []api.MachineEvent `json:"last_events"`
+	MachineGatewayEnabled bool               `json:"machine_gateway_enabled"`
 }
 
 type MachineInstance struct {
@@ -47,15 +48,16 @@ func (mi *MachineInstance) InstanceOptions() instance.InstanceOptions {
 
 func (mi *MachineInstance) ClusterInstance() cluster.MachineInstance {
 	return cluster.MachineInstance{
-		Id:             mi.Machine.InstanceId,
-		Node:           mi.Machine.Node,
-		Namespace:      mi.Machine.Namespace,
-		MachineId:      mi.Machine.Id,
-		MachineVersion: mi.Version.Id,
-		Events:         mi.State.LastEvents,
-		Status:         mi.State.Status,
-		LocalIPV4:      mi.State.LocalIPV4,
-		CreatedAt:      mi.State.CreatedAt,
-		UpdatedAt:      mi.State.UpdatedAt,
+		Id:                   mi.Machine.InstanceId,
+		Node:                 mi.Machine.Node,
+		Namespace:            mi.Machine.Namespace,
+		MachineId:            mi.Machine.Id,
+		MachineVersion:       mi.Version.Id,
+		Events:               mi.State.LastEvents,
+		Status:               mi.State.Status,
+		LocalIPV4:            mi.State.LocalIPV4,
+		CreatedAt:            mi.State.CreatedAt,
+		UpdatedAt:            mi.State.UpdatedAt,
+		EnableMachineGateway: mi.State.MachineGatewayEnabled,
 	}
 }
diff --git a/api/machines.go b/api/machines.go
@@ -29,6 +29,7 @@ type Machine struct {
 	UpdatedAt      time.Time      `json:"updated_at"`
 	Events         []MachineEvent `json:"events"`
 	Status         MachineStatus  `json:"state"`
+	GatewayEnabled bool           `json:"gateway_enabled"`
 }
 
 type MachineStatus string
@@ -69,7 +70,7 @@ type (
 	MachineConfig struct {
 		Image       string      `json:"image"`
 		Guest       GuestConfig `json:"guest"`
-		Workload    Workload    `json:"workload"`
+		Workload    Workload    `json:"workload,omitempty"`
 		StopConfig  *StopConfig `json:"stop_config,omitempty"`
 		AutoDestroy bool        `json:"auto_destroy,omitempty"`
 	}
@@ -120,12 +121,14 @@ const (
 	MachineExited        MachineEventType = "machine.exited"
 	MachineDestroy       MachineEventType = "machine.destroy"
 	MachineDestroyed     MachineEventType = "machine.destroyed"
+	MachineGateway       MachineEventType = "machine.gateway"
 )
 
 type CreateMachinePayload struct {
-	Region    string        `json:"region"`
-	Config    MachineConfig `json:"config"`
-	SkipStart bool          `json:"skip_start"`
+	Region               string        `json:"region"`
+	Config               MachineConfig `json:"config"`
+	SkipStart            bool          `json:"skip_start,omitempty"`
+	EnableMachineGateway bool          `json:"enable_machine_gateway,omitempty"`
 }
 
 type MachineStartEventPayload struct {
@@ -159,6 +162,10 @@ type MachineDestroyEventPayload struct {
 	Force       bool   `json:"force"`
 }
 
+type MachineGatewayEventPayload struct {
+	Enabled bool `json:"enabled"`
+}
+
 type MachineEventPayload struct {
 	PrepareFailed *MachinePrepareFailedEventPayload `json:"prepare_failed,omitempty"`
 	Stop          *MachineStopEventPayload          `json:"stop,omitempty"`
@@ -167,6 +174,7 @@ type MachineEventPayload struct {
 	Started       *MachineStartedEventPayload       `json:"started,omitempty"`
 	Exited        *MachineExitedEventPayload        `json:"stopped,omitempty"`
 	Destroy       *MachineDestroyEventPayload       `json:"destroy,omitempty"`
+	Gateway       *MachineGatewayEventPayload       `json:"gateway,omitempty"`
 }
 
 type Origin string

diff --git a/core/cluster/agent.go b/core/cluster/agent.go
@@ -8,10 +8,11 @@ import (
 )
 
 type PutMachineOptions struct {
-	Machine      Machine            `json:"machine"`
-	Version      api.MachineVersion `json:"version"`
-	AllocationId string             `json:"allocation_id"`
-	Start        bool               `json:"start"`
+	Machine       Machine            `json:"machine"`
+	Version       api.MachineVersion `json:"version"`
+	AllocationId  string             `json:"allocation_id"`
+	Start         bool               `json:"start"`
+	EnableGateway bool               `json:"enable_gateway"`
 }
 
 type Agent interface {
@@ -26,4 +27,7 @@ type Agent interface {
 	DestroyMachine(ctx context.Context, machineId string, force bool) error
 	SubscribeToMachineLogs(ctx context.Context, id string) ([]*api.LogEntry, <-chan *api.LogEntry, error)
 	GetMachineLogs(ctx context.Context, id string) ([]*api.LogEntry, error)
+
+	EnableMachineGateway(ctx context.Context, id string) error
+	DisableMachineGateway(ctx context.Context, id string) error
 }
diff --git a/core/cluster/cluster_state.go b/core/cluster/cluster_state.go
@@ -8,6 +8,8 @@ import (
 )
 
 type Queries interface {
+
+	/* Used by the API server */
 	CreateGateway(ctx context.Context, gateway api.Gateway) error
 	DeleteGateway(ctx context.Context, id string) error
 	DeleteFleetGateways(ctx context.Context, fleetId string) error
@@ -17,17 +19,17 @@ type Queries interface {
 	UpdateMachine(ctx context.Context, m Machine) error
 	DestroyMachine(ctx context.Context, id string) error
 
-	UpsertInstance(ctx context.Context, i MachineInstance) error
+	ListAPIMachines(ctx context.Context, namespace string, fleetId string, includeDestroyed bool) ([]api.Machine, error)
+	GetAPIMachine(ctx context.Context, namespace string, fleetId string, machineId string) (*api.Machine, error)
+	DestroyNamespaceData(ctx context.Context, namespace string) error
 
 	GetNode(ctx context.Context, id string) (api.Node, error)
-	UpsertNode(ctx context.Context, node api.Node) error
 	ListNodesInRegion(ctx context.Context, region string) ([]api.Node, error)
 	ListNodes(ctx context.Context) ([]api.Node, error)
 
-	ListAPIMachines(ctx context.Context, namespace string, fleetId string, includeDestroyed bool) ([]api.Machine, error)
-	GetAPIMachine(ctx context.Context, namespace string, fleetId string, machineId string) (*api.Machine, error)
-
-	DestroyNamespaceData(ctx context.Context, namespace string) error
+	/* Used by raveld */
+	UpsertNode(ctx context.Context, node api.Node) error
+	UpsertInstance(ctx context.Context, i MachineInstance) error
 }
 
 type ClusterState interface {

diff --git a/core/cluster/corrosion/instances.go b/core/cluster/corrosion/instances.go
@@ -9,22 +9,36 @@ import (
 	"github.com/valyentdev/ravel/core/instance"
 )
 
+type corroBool uint8
+
+func (b corroBool) Bool() bool {
+	return b == 1
+}
+
+func fromBool(b bool) corroBool {
+	if b {
+		return 1
+	}
+	return 0
+}
+
 func (c *Queries) UpsertInstance(ctx context.Context, i cluster.MachineInstance) error {
 	eventsBytes, err := json.Marshal(i.Events)
 	if err != nil {
 		return err
 	}
 	_, err = c.dbtx.Exec(ctx,
-		`INSERT INTO instances (id, node, machine_id, machine_version, status, created_at, updated_at, local_ipv4, events, namespace)
-					VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
+		`INSERT INTO instances (id, node, machine_id, machine_version, status, created_at, updated_at, local_ipv4, events, namespace, enable_machine_gateway)
+					VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
 					ON CONFLICT (id, machine_id) DO UPDATE SET
 						status = $5,
 						updated_at = $7,
 						local_ipv4 = $8,
-						events = $9
+						events = $9,
+						enable_machine_gateway = $11
 						`,
 
-		i.Id, i.Node, i.MachineId, i.MachineVersion, i.Status, i.CreatedAt.Unix(), i.UpdatedAt.Unix(), i.LocalIPV4, string(eventsBytes), i.Namespace,
+		i.Id, i.Node, i.MachineId, i.MachineVersion, i.Status, i.CreatedAt.Unix(), i.UpdatedAt.Unix(), i.LocalIPV4, string(eventsBytes), i.Namespace, fromBool(i.EnableMachineGateway),
 	)
 	if err != nil {
 		return err