From 79ee4b81082f0df689b03be1ecb4503b8437aa0a Mon Sep 17 00:00:00 2001
From: Kazuyoshi Kato <kaz@fly.io>
Date: Tue, 12 Nov 2024 13:46:24 -0800
Subject: [PATCH] fix: prevent a panic if the machine is missing

Not really sure the condition to trigger the panic, but returning
an error would make investigating the issue slightly easier.
---
 .../deploy/machines_deploymachinesapp.go      |  6 +++-
 .../deploy/machines_deploymachinesapp_test.go | 36 +++++++++++++++++++
 internal/command/deploy/mock_client_test.go   |  8 ++++-
 internal/command/deploy/plan.go               |  1 +
 4 files changed, 49 insertions(+), 2 deletions(-)
 create mode 100644 internal/command/deploy/machines_deploymachinesapp_test.go

diff --git a/internal/command/deploy/machines_deploymachinesapp.go b/internal/command/deploy/machines_deploymachinesapp.go
index 51261f77b5..71e5f1f09f 100644
--- a/internal/command/deploy/machines_deploymachinesapp.go
+++ b/internal/command/deploy/machines_deploymachinesapp.go
@@ -591,6 +591,7 @@ func (md *machineDeployment) updateExistingMachines(ctx context.Context, updateE
 	return err
 }
 
+// updateExistingMachinesWRecovery updates existing machines.
 // The code duplication is on purpose here. The plan is to completely move over to updateExistingMachinesWRecovery
 func (md *machineDeployment) updateExistingMachinesWRecovery(ctx context.Context, updateEntries []*machineUpdateEntry) (err error) {
 	ctx, span := tracing.GetTracer().Start(ctx, "update_existing_machines_w_recovery", trace.WithAttributes(
@@ -652,9 +653,12 @@ func (md *machineDeployment) updateExistingMachinesWRecovery(ctx context.Context
 		canaryAppState.Machines = []*fly.Machine{oldAppState.Machines[0]}
 
 		newCanaryAppState := newAppState
-		canaryMach, _ := lo.Find(newAppState.Machines, func(m *fly.Machine) bool {
+		canaryMach, exists := lo.Find(newAppState.Machines, func(m *fly.Machine) bool {
 			return m.ID == oldAppState.Machines[0].ID
 		})
+		if !exists {
+			return fmt.Errorf("failed to find machine %s under app %s", oldAppState.Machines[0].ID, md.app.Name)
+		}
 		newCanaryAppState.Machines = []*fly.Machine{canaryMach}
 
 		if err := md.updateMachinesWRecovery(ctx, &canaryAppState, &newCanaryAppState, nil, updateMachineSettings{
diff --git a/internal/command/deploy/machines_deploymachinesapp_test.go b/internal/command/deploy/machines_deploymachinesapp_test.go
new file mode 100644
index 0000000000..0799e3dc93
--- /dev/null
+++ b/internal/command/deploy/machines_deploymachinesapp_test.go
@@ -0,0 +1,36 @@
+package deploy
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/superfly/fly-go"
+	"github.com/superfly/flyctl/internal/machine"
+	"github.com/superfly/flyctl/iostreams"
+)
+
+func TestUpdateExistingMachinesWRecovery(t *testing.T) {
+	ios, _, _, _ := iostreams.Test()
+	client := &mockFlapsClient{}
+	client.machines = []*fly.Machine{{ID: "test-machine-id"}}
+	md := &machineDeployment{
+		app:         &fly.AppCompact{},
+		io:          ios,
+		colorize:    ios.ColorScheme(),
+		flapsClient: client,
+		strategy:    "canary",
+	}
+
+	ctx := context.Background()
+	err := md.updateExistingMachinesWRecovery(ctx, nil)
+	assert.NoError(t, err)
+
+	err = md.updateExistingMachinesWRecovery(ctx, []*machineUpdateEntry{
+		{
+			leasableMachine: machine.NewLeasableMachine(client, ios, &fly.Machine{}, false),
+			launchInput:     &fly.LaunchMachineInput{},
+		},
+	})
+	assert.Error(t, err, "failed to find machine test-machine-id")
+}
diff --git a/internal/command/deploy/mock_client_test.go b/internal/command/deploy/mock_client_test.go
index 6e78fe07f9..243c52dc22 100644
--- a/internal/command/deploy/mock_client_test.go
+++ b/internal/command/deploy/mock_client_test.go
@@ -21,6 +21,9 @@ type mockFlapsClient struct {
 	breakWait        bool
 	breakUncordon    bool
 	breakSetMetadata bool
+	breakList        bool
+
+	machines []*fly.Machine
 }
 
 func (m *mockFlapsClient) AcquireLease(ctx context.Context, machineID string, ttl *int) (*fly.MachineLease, error) {
@@ -123,7 +126,10 @@ func (m *mockFlapsClient) Launch(ctx context.Context, builder fly.LaunchMachineI
 }
 
 func (m *mockFlapsClient) List(ctx context.Context, state string) ([]*fly.Machine, error) {
-	return nil, fmt.Errorf("failed to list machines")
+	if m.breakList {
+		return nil, fmt.Errorf("failed to list machines")
+	}
+	return m.machines, nil
 }
 
 func (m *mockFlapsClient) ListActive(ctx context.Context) ([]*fly.Machine, error) {
diff --git a/internal/command/deploy/plan.go b/internal/command/deploy/plan.go
index 62a18f13d7..aa64c63882 100644
--- a/internal/command/deploy/plan.go
+++ b/internal/command/deploy/plan.go
@@ -60,6 +60,7 @@ type machinePairing struct {
 	newMachine *fly.Machine
 }
 
+// appState returns the app's state from Flaps.
 func (md *machineDeployment) appState(ctx context.Context, existingAppState *AppState) (*AppState, error) {
 	ctx, span := tracing.GetTracer().Start(ctx, "app_state")
 	defer span.End()