From 04cfbad47f467bf7e0f67864063659e6513fad4c Mon Sep 17 00:00:00 2001
From: Saket Chaudhary <saketchaudharynit@gmail.com>
Date: Thu, 10 Oct 2024 01:24:57 +0530
Subject: [PATCH] Extension checks

---
 checks/startup_check.go | 19 ++++++++++++-------
 config/config.go        | 24 ++++++++++++++++++++++++
 2 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/checks/startup_check.go b/checks/startup_check.go
index 0e94e66..ccd96a1 100644
--- a/checks/startup_check.go
+++ b/checks/startup_check.go
@@ -18,11 +18,11 @@ type LogSender interface {
 }
 
 /// Register checks here
-var checks = []checkFn{
-	agentVersionCheck,
-	handlerCheck,
-	sanityCheck,
-	vendorCheck,
+var checks = map[string]checkFn{
+    "agent":         agentVersionCheck, 
+    "handler":       handlerCheck,
+    "sanity":        sanityCheck, 
+    "vendor":        vendorCheck, 
 }
 
 func RunChecks(ctx context.Context, conf *config.Configuration, reg *api.RegistrationResponse, logSender LogSender) {
@@ -32,8 +32,13 @@ func RunChecks(ctx context.Context, conf *config.Configuration, reg *api.Registr
 		util.Logln(errLog)
 	}
 
-	for _, check := range checks {
-		runCheck(ctx, conf, reg, runtimeConfig, logSender, check)
+	runAllChecks := len(conf.ExtensionChecks) == 0
+
+	for checkName, check := range checks {
+		if runAllChecks || conf.ExtensionChecks[checkName] {
+			util.Debugf("Running Extension check: %s", checkName)
+			runCheck(ctx, conf, reg, runtimeConfig, logSender, check)
+		}
 	}
 }
 
diff --git a/config/config.go b/config/config.go
index b02eec0..f7f1f66 100644
--- a/config/config.go
+++ b/config/config.go
@@ -21,6 +21,7 @@ var EmptyNRWrapper = "Undefined"
 type Configuration struct {
 	TestingOverride            bool // ignores envioronment specific details when running unit tests
 	ExtensionEnabled           bool
+	ExtensionChecks            map[string]bool
 	LogsEnabled                bool
 	SendFunctionLogs           bool
 	CollectTraceID             bool
@@ -40,6 +41,7 @@ type Configuration struct {
 func ConfigurationFromEnvironment() *Configuration {
 	nrEnabledStr, nrEnabledOverride := os.LookupEnv("NEW_RELIC_ENABLED")
 	nrEnabledRubyStr, nrEnabledRubyOverride := os.LookupEnv("NEW_RELIC_AGENT_ENABLED")
+	nrExtensionChecksStr, nrExtensionChecksOverride := os.LookupEnv("NEW_RELIC_EXTENSION_CHECKS")
 	enabledStr, extensionEnabledOverride := os.LookupEnv("NEW_RELIC_LAMBDA_EXTENSION_ENABLED")
 	licenseKey, lkOverride := os.LookupEnv("NEW_RELIC_LICENSE_KEY")
 	licenseKeySecretId, lkSecretOverride := os.LookupEnv("NEW_RELIC_LICENSE_KEY_SECRET")
@@ -93,6 +95,28 @@ func ConfigurationFromEnvironment() *Configuration {
 		ret.LicenseKeySSMParameterName = licenseKeySSMParameterName
 	}
 
+	if nrExtensionChecksOverride && nrExtensionChecksStr != "" {
+		validChecks := map[string]bool{
+			"agent":   true,
+			"handler": true,
+			"sanity":  true,
+			"vendor":  true,
+		}
+		selectedChecks := make(map[string]bool)
+		checks := strings.Split(nrExtensionChecksStr, ",")
+	
+		for _, check := range checks {
+			trimmedCheck := strings.TrimSpace(check)
+			if trimmedCheck != "" && validChecks[trimmedCheck] {
+				selectedChecks[trimmedCheck] = true
+			}
+		}
+	
+		if len(selectedChecks) > 0 {
+			ret.ExtensionChecks = selectedChecks
+		}
+	}
+
 	if nrOverride {
 		ret.NRHandler = nrHandler
 	} else {