feat: reboot instances (#68)

2026-10-16 01:22:00 +00:00 · 2026-03-10 10:48:58 +10:00
parent 015dbbc272
commit ee3df227a7
9 changed files with 442 additions and 1 deletions
@@ -88,6 +88,8 @@ func (d *Daemon) handleEC2Events(msg *nats.Msg) {
 		d.handleDetachVolume(msg, command, instance)
 	case command.Attributes.StartInstance:
 		d.handleStartInstance(msg, command, instance)
+	case command.Attributes.RebootInstance:
+		d.handleRebootInstance(msg, command, instance)
 	case command.Attributes.StopInstance, command.Attributes.TerminateInstance:
 		d.handleStopOrTerminateInstance(msg, command, instance)
 	default:
@@ -11,6 +11,7 @@ import (
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/ec2"
 	"github.com/mulgadc/hive/hive/awserrors"
+	"github.com/mulgadc/hive/hive/qmp"
 	"github.com/mulgadc/hive/hive/types"
 	"github.com/mulgadc/hive/hive/utils"
 	"github.com/mulgadc/hive/hive/vm"
@@ -305,6 +306,33 @@ func (d *Daemon) handleEC2RunInstances(msg *nats.Msg) {
 	slog.Info("handleEC2RunInstances completed", "requested", launchCount, "created", len(instances), "launched", successCount)
 }

+func (d *Daemon) handleRebootInstance(msg *nats.Msg, command types.EC2InstanceCommand, instance *vm.VM) {
+	slog.Info("Rebooting instance", "id", command.ID)
+
+	d.Instances.Mu.Lock()
+	status := instance.Status
+	d.Instances.Mu.Unlock()
+
+	if status != vm.StateRunning {
+		slog.Error("RebootInstance: instance not in running state", "instanceId", command.ID, "status", status)
+		respondWithError(msg, awserrors.ErrorIncorrectInstanceState)
+		return
+	}
+
+	_, err := d.SendQMPCommand(instance.QMPClient, qmp.QMPCommand{Execute: "system_reset"}, command.ID)
+	if err != nil {
+		slog.Error("RebootInstance: QMP system_reset failed", "instanceId", command.ID, "err", err)
+		respondWithError(msg, awserrors.ErrorServerInternal)
+		return
+	}
+
+	slog.Info("Instance rebooted", "instanceId", command.ID)
+
+	if err := msg.Respond([]byte(`{}`)); err != nil {
+		slog.Error("Failed to respond to NATS request", "err", err)
+	}
+}
+
 func (d *Daemon) handleStartInstance(msg *nats.Msg, command types.EC2InstanceCommand, instance *vm.VM) {
 	slog.Info("Starting instance", "id", command.ID)

@@ -623,6 +623,142 @@ func TestHandleEC2Events_TerminateInstance(t *testing.T) {
 	assert.Equal(t, vm.StateShuttingDown, status)
 }

+func TestHandleEC2Events_RebootRunningInstance(t *testing.T) {
+	natsURL := sharedNATSURL
+
+	daemon := createTestDaemon(t, natsURL)
+
+	instanceID := "i-test-reboot-001"
+	daemon.Instances.VMS[instanceID] = &vm.VM{
+		ID:           instanceID,
+		InstanceType: getTestInstanceType(),
+		Status:       vm.StateRunning,
+		Instance:     &ec2.Instance{},
+		QMPClient:    &qmp.QMPClient{},
+		AccountID:    testAccountID,
+	}
+
+	sub, err := daemon.natsConn.Subscribe(
+		fmt.Sprintf("ec2.cmd.%s", instanceID),
+		daemon.handleEC2Events,
+	)
+	require.NoError(t, err)
+	defer sub.Unsubscribe()
+
+	cmd := types.EC2InstanceCommand{
+		ID:         instanceID,
+		Attributes: types.EC2CommandAttributes{RebootInstance: true},
+	}
+	cmdData, _ := json.Marshal(cmd)
+
+	reply, err := natsRequest(daemon.natsConn,
+		fmt.Sprintf("ec2.cmd.%s", instanceID),
+		cmdData,
+		5*time.Second,
+	)
+	require.NoError(t, err)
+	require.NotNil(t, reply)
+
+	// With nil QMPClient encoder/decoder, SendQMPCommand returns error,
+	// so we expect an error response (ServerInternal).
+	var errResp map[string]any
+	err = json.Unmarshal(reply.Data, &errResp)
+	require.NoError(t, err)
+	assert.Contains(t, errResp, "Code")
+
+	// Instance should remain in running state (reboot doesn't change state)
+	daemon.Instances.Mu.Lock()
+	status := daemon.Instances.VMS[instanceID].Status
+	daemon.Instances.Mu.Unlock()
+	assert.Equal(t, vm.StateRunning, status)
+}
+
+func TestHandleEC2Events_RebootStoppedInstance(t *testing.T) {
+	natsURL := sharedNATSURL
+
+	daemon := createTestDaemon(t, natsURL)
+
+	instanceID := "i-test-reboot-stopped"
+	daemon.Instances.VMS[instanceID] = &vm.VM{
+		ID:           instanceID,
+		InstanceType: getTestInstanceType(),
+		Status:       vm.StateStopped,
+		Instance:     &ec2.Instance{},
+		QMPClient:    &qmp.QMPClient{},
+		AccountID:    testAccountID,
+	}
+
+	sub, err := daemon.natsConn.Subscribe(
+		fmt.Sprintf("ec2.cmd.%s", instanceID),
+		daemon.handleEC2Events,
+	)
+	require.NoError(t, err)
+	defer sub.Unsubscribe()
+
+	cmd := types.EC2InstanceCommand{
+		ID:         instanceID,
+		Attributes: types.EC2CommandAttributes{RebootInstance: true},
+	}
+	cmdData, _ := json.Marshal(cmd)
+
+	reply, err := natsRequest(daemon.natsConn,
+		fmt.Sprintf("ec2.cmd.%s", instanceID),
+		cmdData,
+		5*time.Second,
+	)
+	require.NoError(t, err)
+	require.NotNil(t, reply)
+
+	// Should get IncorrectInstanceState error
+	var errResp map[string]any
+	err = json.Unmarshal(reply.Data, &errResp)
+	require.NoError(t, err)
+	assert.Contains(t, errResp, "Code")
+}
+
+func TestHandleEC2Events_RebootTerminatedInstance(t *testing.T) {
+	natsURL := sharedNATSURL
+
+	daemon := createTestDaemon(t, natsURL)
+
+	instanceID := "i-test-reboot-terminated"
+	daemon.Instances.VMS[instanceID] = &vm.VM{
+		ID:           instanceID,
+		InstanceType: getTestInstanceType(),
+		Status:       vm.StateTerminated,
+		Instance:     &ec2.Instance{},
+		QMPClient:    &qmp.QMPClient{},
+		AccountID:    testAccountID,
+	}
+
+	sub, err := daemon.natsConn.Subscribe(
+		fmt.Sprintf("ec2.cmd.%s", instanceID),
+		daemon.handleEC2Events,
+	)
+	require.NoError(t, err)
+	defer sub.Unsubscribe()
+
+	cmd := types.EC2InstanceCommand{
+		ID:         instanceID,
+		Attributes: types.EC2CommandAttributes{RebootInstance: true},
+	}
+	cmdData, _ := json.Marshal(cmd)
+
+	reply, err := natsRequest(daemon.natsConn,
+		fmt.Sprintf("ec2.cmd.%s", instanceID),
+		cmdData,
+		5*time.Second,
+	)
+	require.NoError(t, err)
+	require.NotNil(t, reply)
+
+	// Should get IncorrectInstanceState error
+	var errResp map[string]any
+	err = json.Unmarshal(reply.Data, &errResp)
+	require.NoError(t, err)
+	assert.Contains(t, errResp, "Code")
+}
+
 func TestHandleEC2Events_InstanceNotFound(t *testing.T) {
 	natsURL := sharedNATSURL

@@ -65,6 +65,9 @@ var ec2Actions = map[string]EC2Handler{
 	"StopInstances": ec2Handler(func(input *ec2.StopInstancesInput, gw *GatewayConfig, accountID string) (any, error) {
 		return gateway_ec2_instance.StopInstances(input, gw.NATSConn, accountID)
 	}),
+	"RebootInstances": ec2Handler(func(input *ec2.RebootInstancesInput, gw *GatewayConfig, accountID string) (any, error) {
+		return gateway_ec2_instance.RebootInstances(input, gw.NATSConn, accountID)
+	}),
 	"TerminateInstances": ec2Handler(func(input *ec2.TerminateInstancesInput, gw *GatewayConfig, accountID string) (any, error) {
 		return gateway_ec2_instance.TerminateInstances(input, gw.NATSConn, accountID)
 	}),
@@ -0,0 +1,79 @@
+package gateway_ec2_instance
+
+import (
+	"encoding/json"
+	"errors"
+	"fmt"
+	"log/slog"
+	"time"
+
+	"github.com/aws/aws-sdk-go/service/ec2"
+	"github.com/mulgadc/hive/hive/awserrors"
+	"github.com/mulgadc/hive/hive/types"
+	"github.com/mulgadc/hive/hive/utils"
+	"github.com/nats-io/nats.go"
+)
+
+// RebootInstances sends reboot commands to specified instances via NATS.
+// Unlike stop+start, reboot keeps the instance running and sends a QMP system_reset.
+// Returns an empty response on success (AWS returns no state-change data).
+func RebootInstances(input *ec2.RebootInstancesInput, natsConn *nats.Conn, accountID string) (*ec2.RebootInstancesOutput, error) {
+	if len(input.InstanceIds) == 0 {
+		return nil, fmt.Errorf("no instance IDs provided")
+	}
+
+	slog.Info("RebootInstances: Processing request", "instance_count", len(input.InstanceIds))
+
+	for _, instanceIDPtr := range input.InstanceIds {
+		if instanceIDPtr == nil {
+			continue
+		}
+		instanceID := *instanceIDPtr
+
+		command := types.EC2InstanceCommand{
+			ID: instanceID,
+			Attributes: types.EC2CommandAttributes{
+				RebootInstance: true,
+			},
+		}
+
+		jsonData, err := json.Marshal(command)
+		if err != nil {
+			slog.Error("RebootInstances: Failed to marshal command", "instance_id", instanceID, "err", err)
+			continue
+		}
+
+		subject := fmt.Sprintf("ec2.cmd.%s", instanceID)
+		reqMsg := nats.NewMsg(subject)
+		reqMsg.Data = jsonData
+		reqMsg.Header.Set(utils.AccountIDHeader, accountID)
+		msg, err := natsConn.RequestMsg(reqMsg, 5*time.Second)
+		if err != nil {
+			slog.Error("RebootInstances: Failed to send command", "instance_id", instanceID, "err", err)
+
+			// NATS timeout means no daemon has a subscription for this instance.
+			// Check if the instance exists in the stopped-instances KV bucket —
+			// if so, return IncorrectInstanceState instead of NotFound.
+			describeInput := &ec2.DescribeInstancesInput{
+				InstanceIds: []*string{&instanceID},
+			}
+			describeData, _ := json.Marshal(describeInput)
+			if reservations := queryInstanceBucket(natsConn, "ec2.DescribeStoppedInstances", describeData, accountID); len(reservations) > 0 {
+				return nil, errors.New(awserrors.ErrorIncorrectInstanceState)
+			}
+
+			return nil, errors.New(awserrors.ErrorInvalidInstanceIDNotFound)
+		}
+
+		// Check if the daemon returned an error response (e.g. ownership check failure)
+		if responseError, parseErr := utils.ValidateErrorPayload(msg.Data); parseErr != nil {
+			slog.Error("RebootInstances: Daemon returned error", "instance_id", instanceID, "code", *responseError.Code)
+			return nil, errors.New(*responseError.Code)
+		}
+
+		slog.Info("RebootInstances: Command sent successfully", "instance_id", instanceID)
+	}
+
+	slog.Info("RebootInstances: Completed", "total_instances", len(input.InstanceIds))
+	return &ec2.RebootInstancesOutput{}, nil
+}
@@ -0,0 +1,130 @@
+package gateway_ec2_instance
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/service/ec2"
+	"github.com/mulgadc/hive/hive/awserrors"
+	"github.com/mulgadc/hive/hive/types"
+	"github.com/mulgadc/hive/hive/utils"
+	"github.com/nats-io/nats.go"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestRebootInstances_Success(t *testing.T) {
+	_, nc := startTestNATSServer(t)
+
+	instanceID := "i-0123456789abcdef0"
+
+	nc.Subscribe("ec2.cmd."+instanceID, func(msg *nats.Msg) {
+		var cmd types.EC2InstanceCommand
+		err := json.Unmarshal(msg.Data, &cmd)
+		require.NoError(t, err)
+
+		assert.Equal(t, instanceID, cmd.ID)
+		assert.True(t, cmd.Attributes.RebootInstance)
+
+		msg.Respond([]byte(`{}`))
+	})
+
+	input := &ec2.RebootInstancesInput{
+		InstanceIds: []*string{aws.String(instanceID)},
+	}
+
+	output, err := RebootInstances(input, nc, "123456789012")
+
+	require.NoError(t, err)
+	require.NotNil(t, output)
+}
+
+func TestRebootInstances_MultipleInstances(t *testing.T) {
+	_, nc := startTestNATSServer(t)
+
+	ids := []string{"i-001", "i-002"}
+
+	for _, id := range ids {
+		nc.Subscribe("ec2.cmd."+id, func(msg *nats.Msg) {
+			msg.Respond([]byte(`{}`))
+		})
+	}
+
+	input := &ec2.RebootInstancesInput{
+		InstanceIds: []*string{aws.String(ids[0]), aws.String(ids[1])},
+	}
+
+	output, err := RebootInstances(input, nc, "123456789012")
+
+	require.NoError(t, err)
+	require.NotNil(t, output)
+}
+
+func TestRebootInstances_EmptyInstanceIds(t *testing.T) {
+	_, nc := startTestNATSServer(t)
+
+	input := &ec2.RebootInstancesInput{
+		InstanceIds: []*string{},
+	}
+
+	_, err := RebootInstances(input, nc, "123456789012")
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "no instance IDs provided")
+}
+
+func TestRebootInstances_NATSRequestFails(t *testing.T) {
+	_, nc := startTestNATSServer(t)
+
+	input := &ec2.RebootInstancesInput{
+		InstanceIds: []*string{aws.String("i-nosubscriber")},
+	}
+
+	_, err := RebootInstances(input, nc, "123456789012")
+	require.Error(t, err)
+	assert.Equal(t, awserrors.ErrorInvalidInstanceIDNotFound, err.Error())
+}
+
+func TestRebootInstances_StoppedInstance(t *testing.T) {
+	_, nc := startTestNATSServer(t)
+
+	instanceID := "i-stopped"
+
+	// No ec2.cmd subscription — simulates a stopped instance with no daemon listener.
+	// Subscribe to the stopped-instances describe topic to return the instance.
+	nc.Subscribe("ec2.DescribeStoppedInstances", func(msg *nats.Msg) {
+		resp := ec2.DescribeInstancesOutput{
+			Reservations: []*ec2.Reservation{
+				{Instances: []*ec2.Instance{{InstanceId: aws.String(instanceID)}}},
+			},
+		}
+		data, _ := json.Marshal(resp)
+		msg.Respond(data)
+	})
+
+	input := &ec2.RebootInstancesInput{
+		InstanceIds: []*string{aws.String(instanceID)},
+	}
+
+	_, err := RebootInstances(input, nc, "123456789012")
+	require.Error(t, err)
+	assert.Equal(t, awserrors.ErrorIncorrectInstanceState, err.Error())
+}
+
+func TestRebootInstances_DaemonError(t *testing.T) {
+	_, nc := startTestNATSServer(t)
+
+	instanceID := "i-error"
+
+	nc.Subscribe("ec2.cmd."+instanceID, func(msg *nats.Msg) {
+		msg.Respond(utils.GenerateErrorPayload(awserrors.ErrorIncorrectInstanceState))
+	})
+
+	input := &ec2.RebootInstancesInput{
+		InstanceIds: []*string{aws.String(instanceID)},
+	}
+
+	_, err := RebootInstances(input, nc, "123456789012")
+	require.Error(t, err)
+	assert.Equal(t, awserrors.ErrorIncorrectInstanceState, err.Error())
+}
@@ -891,7 +891,7 @@ func TestSlogRequestLogger_CapturesStatusCode(t *testing.T) {
 func TestEC2ActionMapCompleteness(t *testing.T) {
 	expectedActions := []string{
 		"DescribeInstances", "RunInstances", "StartInstances", "StopInstances",
-		"TerminateInstances", "DescribeInstanceTypes", "GetConsoleOutput",
+		"TerminateInstances", "RebootInstances", "DescribeInstanceTypes", "GetConsoleOutput",
 		"ModifyInstanceAttribute",
 		"CreateKeyPair", "DeleteKeyPair", "DescribeKeyPairs", "ImportKeyPair",
 		"DescribeImages", "CreateImage",
@@ -17,6 +17,7 @@ type EC2CommandAttributes struct {
 	StartInstance     bool `json:"start_instance"`
 	AttachVolume      bool `json:"attach_volume"`
 	DetachVolume      bool `json:"detach_volume"`
+	RebootInstance    bool `json:"reboot_instance"`
 }

 // AttachVolumeData carries parameters for an attach-volume command.
@@ -965,6 +965,11 @@ if [ "$STATE" != "stopped" ]; then
    exit 1
 fi

+# Reboot stopped instance (should fail)
+echo "Attempting reboot of stopped instance (should fail)..."
+expect_error "IncorrectInstanceState" aws ec2 reboot-instances --instance-ids "$INSTANCE_ID"
+echo "Reboot-stopped correctly rejected"
+
 # Phase 7a: Attach volume to stopped instance (should fail)
 echo "Phase 7a: Attach Volume to Stopped Instance (Error Path)"
 echo "Creating a volume to test attach-to-stopped..."
@@ -1077,6 +1082,55 @@ echo "Memory verified"

 echo "ModifyInstanceAttribute test passed (type change + vCPU + memory verified via SSH)"

+# Phase 7c-pre: Reboot Running Instance
+echo "Phase 7c-pre: Reboot Instance"
+
+# Capture pre-reboot metadata
+PRE_REBOOT_IP=$(aws ec2 describe-instances --instance-ids "$INSTANCE_ID" \
+    --query 'Reservations[0].Instances[0].PrivateIpAddress' --output text)
+
+echo "Rebooting instance..."
+aws ec2 reboot-instances --instance-ids "$INSTANCE_ID"
+
+# Verify state stays running (poll a few times to confirm no transient state change)
+COUNT=0
+while [ $COUNT -lt 5 ]; do
+    STATE=$(aws ec2 describe-instances --instance-ids "$INSTANCE_ID" \
+        --query 'Reservations[0].Instances[0].State.Name' --output text)
+    if [ "$STATE" != "running" ]; then
+        echo "Instance unexpectedly left running state: $STATE"
+        exit 1
+    fi
+    sleep 2
+    COUNT=$((COUNT + 1))
+done
+echo "Instance state remained running during reboot"
+
+# Wait for SSH to come back (guest OS restarts)
+SSH_PORT=$(get_ssh_port "$INSTANCE_ID")
+SSH_HOST=$(get_ssh_host "$INSTANCE_ID")
+wait_for_ssh "$SSH_HOST" "$SSH_PORT" "test-key-1.pem" 30
+
+# Verify guest actually rebooted (uptime < 120 seconds)
+UPTIME_SECS=$(ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null \
+    -o LogLevel=ERROR -o ConnectTimeout=5 -o BatchMode=yes \
+    -p "$SSH_PORT" -i "test-key-1.pem" \
+    ec2-user@"$SSH_HOST" 'cat /proc/uptime | cut -d. -f1' | tr -d '[:space:]')
+echo "Guest uptime after reboot: ${UPTIME_SECS}s"
+if [ "$UPTIME_SECS" -gt 120 ]; then
+    echo "Guest uptime too high — reboot may not have occurred"
+    exit 1
+fi
+
+# Verify metadata unchanged
+POST_REBOOT_IP=$(aws ec2 describe-instances --instance-ids "$INSTANCE_ID" \
+    --query 'Reservations[0].Instances[0].PrivateIpAddress' --output text)
+if [ "$PRE_REBOOT_IP" != "$POST_REBOOT_IP" ]; then
+    echo "IP changed after reboot: $PRE_REBOOT_IP -> $POST_REBOOT_IP"
+    exit 1
+fi
+echo "Reboot instance test passed"
+
 # Phase 7c: RunInstances with count > 1
 echo "Phase 7c: RunInstances with MinCount/MaxCount > 1"
 echo "Launching 2 instances in a single run-instances call..."
@@ -1234,6 +1288,10 @@ echo "8q: ModifyInstanceAttribute on running instance..."
 expect_error "InvalidInstanceID.NotFound" aws ec2 modify-instance-attribute \
    --instance-id "$INSTANCE_ID" --instance-type "{\"Value\": \"$INSTANCE_TYPE\"}"

+# 8r: Reboot non-existent instance
+echo "8r: Reboot non-existent instance..."
+expect_error "InvalidInstanceID.NotFound" aws ec2 reboot-instances --instance-ids "i-nonexistent"
+
 echo "Negative test suite passed"

 # =============================================================================
@@ -1880,6 +1938,10 @@ expect_error "InvalidInstanceID.NotFound" \
    aws ec2 terminate-instances --instance-ids "$ALPHA_INST" --profile hive-team-beta
 echo "  Beta cannot terminate Alpha's instance"

+expect_error "InvalidInstanceID.NotFound" \
+    aws ec2 reboot-instances --instance-ids "$BETA_INST" --profile hive-team-alpha
+echo "  Alpha cannot reboot Beta's instance"
+
 # Stop Alpha's instance for cross-account start/modify/console tests
 aws ec2 stop-instances --instance-ids "$ALPHA_INST" --profile hive-team-alpha > /dev/null
 COUNT=0