From 8a7e36cf6b350b2630a4f6216e7e1a2c5690843a Mon Sep 17 00:00:00 2001
From: Matt Calhoun <matt@mattcalhoun.com>
Date: Mon, 16 Dec 2024 13:47:27 -0500
Subject: [PATCH] error checking on retreiving context

---
 cmd/cmd_utils.go   | 7 ++++++-
 cmd/helmfile.go    | 9 ++++++++-
 cmd/terraform.go   | 9 ++++++++-
 cmd/vendor_pull.go | 8 +++++++-
 cmd/version.go     | 6 +++++-
 5 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/cmd/cmd_utils.go b/cmd/cmd_utils.go
index 74364cdde..0c9185108 100644
--- a/cmd/cmd_utils.go
+++ b/cmd/cmd_utils.go
@@ -476,7 +476,12 @@ func CheckForAtmosUpdateAndPrintMessage(cliConfig schema.CliConfiguration) {
 }
 
 func customHelpMessageToUpgradeToAtmosLatestRelease(cmd *cobra.Command, args []string) {
-	atmosConfig := cmd.Context().Value(contextKey("atmos_config")).(*schema.CliConfiguration)
+	value := cmd.Context().Value(contextKey("atmos_config"))
+	if value == nil {
+		u.LogErrorAndExit(schema.CliConfiguration{}, fmt.Errorf("atmos configuration not found in context"))
+	}
+	atmosConfig := value.(*schema.CliConfiguration)
+
 	originalHelpFunc(cmd, args)
 	CheckForAtmosUpdateAndPrintMessage(*atmosConfig)
 }
diff --git a/cmd/helmfile.go b/cmd/helmfile.go
index bb16d0e85..f4d1ac007 100644
--- a/cmd/helmfile.go
+++ b/cmd/helmfile.go
@@ -1,6 +1,8 @@
 package cmd
 
 import (
+	"fmt"
+
 	"github.com/samber/lo"
 	"github.com/spf13/cobra"
 
@@ -34,7 +36,12 @@ var helmfileCmd = &cobra.Command{
 		// Exit on help
 		if info.NeedHelp {
 			// Check for the latest Atmos release on GitHub and print update message
-			atmosConfig := cmd.Context().Value(contextKey("atmos_config")).(*schema.CliConfiguration)
+			value := cmd.Context().Value(contextKey("atmos_config"))
+			if value == nil {
+				u.LogErrorAndExit(schema.CliConfiguration{}, fmt.Errorf("atmos configuration not found in context"))
+			}
+			atmosConfig := value.(*schema.CliConfiguration)
+
 			CheckForAtmosUpdateAndPrintMessage(*atmosConfig)
 			return
 		}
diff --git a/cmd/terraform.go b/cmd/terraform.go
index a645216ab..98f125911 100644
--- a/cmd/terraform.go
+++ b/cmd/terraform.go
@@ -1,6 +1,8 @@
 package cmd
 
 import (
+	"fmt"
+
 	"github.com/samber/lo"
 	"github.com/spf13/cobra"
 
@@ -34,7 +36,12 @@ var terraformCmd = &cobra.Command{
 		// Exit on help
 		if info.NeedHelp {
 			// Check for the latest Atmos release on GitHub and print update message
-			atmosConfig := cmd.Context().Value(contextKey("atmos_config")).(*schema.CliConfiguration)
+			value := cmd.Context().Value(contextKey("atmos_config"))
+			if value == nil {
+				u.LogErrorAndExit(schema.CliConfiguration{}, fmt.Errorf("atmos configuration not found in context"))
+			}
+			atmosConfig := value.(*schema.CliConfiguration)
+
 			CheckForAtmosUpdateAndPrintMessage(*atmosConfig)
 			return
 		}
diff --git a/cmd/vendor_pull.go b/cmd/vendor_pull.go
index 0bdb90cb6..650c70111 100644
--- a/cmd/vendor_pull.go
+++ b/cmd/vendor_pull.go
@@ -1,6 +1,8 @@
 package cmd
 
 import (
+	"fmt"
+
 	"github.com/spf13/cobra"
 
 	e "github.com/cloudposse/atmos/internal/exec"
@@ -15,7 +17,11 @@ var vendorPullCmd = &cobra.Command{
 	Long:               `This command executes 'atmos vendor pull' CLI commands`,
 	FParseErrWhitelist: struct{ UnknownFlags bool }{UnknownFlags: false},
 	Run: func(cmd *cobra.Command, args []string) {
-		atmosConfig := cmd.Context().Value(contextKey("atmos_config")).(*schema.CliConfiguration)
+		value := cmd.Context().Value(contextKey("atmos_config"))
+		if value == nil {
+			u.LogErrorAndExit(schema.CliConfiguration{}, fmt.Errorf("atmos configuration not found in context"))
+		}
+		atmosConfig := value.(*schema.CliConfiguration)
 
 		// WithStackValidation is a functional option that enables/disables stack configuration validation
 		// based on whether the --stack flag is provided
diff --git a/cmd/version.go b/cmd/version.go
index 21b16e844..98067d8b9 100644
--- a/cmd/version.go
+++ b/cmd/version.go
@@ -53,7 +53,11 @@ var versionCmd = &cobra.Command{
 		}
 
 		// Check for the cache and print update message
-		atmosConfig := cmd.Context().Value(contextKey("atmos_config")).(*schema.CliConfiguration)
+		value := cmd.Context().Value(contextKey("atmos_config"))
+		if value == nil {
+			u.LogErrorAndExit(schema.CliConfiguration{}, fmt.Errorf("atmos configuration not found in context"))
+		}
+		atmosConfig := value.(*schema.CliConfiguration)
 		CheckForAtmosUpdateAndPrintMessage(*atmosConfig)
 	},
 }