sia run-after-scripts for cmd line options must be in blocking mode (#…

…2529) Signed-off-by: Henry Avetisyan <[email protected]> Co-authored-by: Henry Avetisyan <[email protected]>
AthenZ · Feb 27, 2024 · aebd6fd · aebd6fd
1 parent 11e5262
commit aebd6fd
Show file tree

Hide file tree

Showing 5 changed files with 46 additions and 20 deletions.
diff --git a/libs/go/sia/agent/agent.go b/libs/go/sia/agent/agent.go
@@ -629,15 +629,15 @@ func RunAgent(siaCmd, ztsUrl string, opts *options.Options) {
 			log.Fatalf("unable to fetch %d out of %d requested role certificates\n", failures, count)
 		}
 		if count != 0 {
-			util.ExecuteScriptWithoutBlock(opts.RunAfterParts)
+			util.ExecuteScript(opts.RunAfterParts)
 		}
 	case "token":
 		if tokenOpts != nil {
 			err := fetchAccessToken(tokenOpts)
 			if err != nil && !skipErrors {
 				log.Fatalf("Unable to fetch access tokens, err: %v\n", err)
 			}
-			util.ExecuteScriptWithoutBlock(opts.RunAfterTokensParts)
+			util.ExecuteScript(opts.RunAfterTokensParts)
 		} else {
 			log.Print("unable to fetch access tokens, invalid or missing configuration")
 		}
@@ -646,14 +646,14 @@ func RunAgent(siaCmd, ztsUrl string, opts *options.Options) {
 		if err != nil {
 			log.Fatalf("Unable to register identity, err: %v\n", err)
 		}
-		util.ExecuteScriptWithoutBlock(opts.RunAfterParts)
+		util.ExecuteScript(opts.RunAfterParts)
 		log.Printf("identity registered for services: %s\n", svcs)
 	case "rotate", "refresh":
 		err = RefreshInstance(ztsUrl, opts)
 		if err != nil {
 			log.Fatalf("Refresh identity failed, err: %v\n", err)
 		}
-		util.ExecuteScriptWithoutBlock(opts.RunAfterParts)
+		util.ExecuteScript(opts.RunAfterParts)
 		log.Printf("Identity successfully refreshed for services: %s\n", svcs)
 	case "init":
 		err := RegisterInstance(ztsUrl, opts, false)
@@ -665,13 +665,13 @@ func RunAgent(siaCmd, ztsUrl string, opts *options.Options) {
 		if failures != 0 && !skipErrors {
 			log.Fatalf("unable to fetch %d out of %d requested role certificates\n", failures, count)
 		}
-		util.ExecuteScriptWithoutBlock(opts.RunAfterParts)
+		util.ExecuteScript(opts.RunAfterParts)
 		if tokenOpts != nil {
 			err := fetchAccessToken(tokenOpts)
 			if err != nil && !skipErrors {
 				log.Fatalf("Unable to fetch access tokens, err: %v\n", err)
 			}
-			util.ExecuteScriptWithoutBlock(opts.RunAfterTokensParts)
+			util.ExecuteScript(opts.RunAfterTokensParts)
 		}
 	default:
 		// we're going to iterate through our configured services.

diff --git a/libs/go/sia/aws/agent/agent.go b/libs/go/sia/aws/agent/agent.go
@@ -625,15 +625,15 @@ func RunAgent(siaCmd, ztsUrl string, opts *options.Options) {
 			log.Fatalf("unable to fetch %d out of %d requested role certificates\n", failures, count)
 		}
 		if count != 0 {
-			util.ExecuteScriptWithoutBlock(opts.RunAfterParts)
+			util.ExecuteScript(opts.RunAfterParts)
 		}
 	case "token":
 		if tokenOpts != nil {
 			err := fetchAccessToken(tokenOpts)
 			if err != nil && !skipErrors {
 				log.Fatalf("Unable to fetch access tokens, err: %v\n", err)
 			}
-			util.ExecuteScriptWithoutBlock(opts.RunAfterTokensParts)
+			util.ExecuteScript(opts.RunAfterTokensParts)
 		} else {
 			log.Print("unable to fetch access tokens, invalid or missing configuration")
 		}
@@ -642,14 +642,14 @@ func RunAgent(siaCmd, ztsUrl string, opts *options.Options) {
 		if err != nil {
 			log.Fatalf("Unable to register identity, err: %v\n", err)
 		}
-		util.ExecuteScriptWithoutBlock(opts.RunAfterParts)
+		util.ExecuteScript(opts.RunAfterParts)
 		log.Printf("identity registered for services: %s\n", svcs)
 	case "rotate", "refresh":
 		err = RefreshInstance(data, ztsUrl, opts)
 		if err != nil {
 			log.Fatalf("Refresh identity failed, err: %v\n", err)
 		}
-		util.ExecuteScriptWithoutBlock(opts.RunAfterParts)
+		util.ExecuteScript(opts.RunAfterParts)
 		log.Printf("Identity successfully refreshed for services: %s\n", svcs)
 	case "init":
 		err := RegisterInstance(data, ztsUrl, opts, false)
@@ -661,13 +661,13 @@ func RunAgent(siaCmd, ztsUrl string, opts *options.Options) {
 		if failures != 0 && !skipErrors {
 			log.Fatalf("unable to fetch %d out of %d requested role certificates\n", failures, count)
 		}
-		util.ExecuteScriptWithoutBlock(opts.RunAfterParts)
+		util.ExecuteScript(opts.RunAfterParts)
 		if tokenOpts != nil {
 			err := fetchAccessToken(tokenOpts)
 			if err != nil && !skipErrors {
 				log.Fatalf("Unable to fetch access tokens, err: %v\n", err)
 			}
-			util.ExecuteScriptWithoutBlock(opts.RunAfterTokensParts)
+			util.ExecuteScript(opts.RunAfterTokensParts)
 		}
 	default:
 		// we're going to iterate through our configured services.

diff --git a/libs/go/sia/util/data/test_after_script.sh b/libs/go/sia/util/data/test_after_script.sh
@@ -0,0 +1,3 @@
+#!/bin/bash
+
+touch /tmp/test-after-script
diff --git a/libs/go/sia/util/util.go b/libs/go/sia/util/util.go
@@ -1280,19 +1280,26 @@ func ParseScriptArguments(script string) []string {
 	return parts
 }
 
-// ExecuteScriptWithoutBlock executes a script along with the provided
-// arguments in a go subroutine without blocking the agent
-func ExecuteScriptWithoutBlock(script []string) {
+// ExecuteScript executes a script along with the provided
+// arguments while blocking the agent
+func ExecuteScript(script []string) error {
 	// execute run after script (if provided)
 	if len(script) == 0 {
-		return
+		return nil
+	}
+	log.Printf("executing run after hook for: %v", script)
+	err := exec.Command(script[0], script[1:]...).Run()
+	if err != nil {
+		log.Printf("unable to execute: %q, err: %v", script, err)
 	}
+	return err
+}
 
+// ExecuteScriptWithoutBlock executes a script along with the provided
+// arguments in a go subroutine without blocking the agent
+func ExecuteScriptWithoutBlock(script []string) {
 	go func() {
-		log.Printf("executing run after hook for: %v", script)
-		if err := exec.Command(script[0], script[1:]...).Run(); err != nil {
-			log.Printf("unable to execute: %q, err: %v", script, err)
-		}
+		ExecuteScript(script)
 	}()
 }
 

diff --git a/libs/go/sia/util/util_test.go b/libs/go/sia/util/util_test.go
@@ -1681,3 +1681,19 @@ func TestParseSiaCmd(test *testing.T) {
 		})
 	}
 }
+
+func TestExecuteScript(test *testing.T) {
+
+	// non-existent script
+	err := ExecuteScript([]string{"unknown-script"})
+	assert.NotNil(test, err)
+
+	// remove our test file if it exists
+	os.Remove("/tmp/test-after-script")
+	// valid script
+	err = ExecuteScript([]string{"data/test_after_script.sh"})
+	assert.Nil(test, err)
+	// verify our test file was created
+	_, err = os.Stat("/tmp/test-after-script")
+	assert.Nil(test, err)
+}