systemd-notify-all option to notify systemd after role certificates

libs/go/sia/agent/agent.go

@@ -702,13 +702,7 @@ func RunAgent(siaCmd, ztsUrl string, opts *options.Options) {
 			}
 		}
 
-		if cmd == "systemd-notify" {
-			err = util.NotifySystemdReady()
-			if err != nil {
-				log.Printf("failed to notify systemd: %v", err)
-			}
-		}
-
+		util.NotifySystemdReadyForCommand(cmd, "systemd-notify")
 		log.Printf("Identity established for services: %s\n", svcs)
 
 		stop := make(chan bool, 1)
@@ -756,6 +750,7 @@ func RunAgent(siaCmd, ztsUrl string, opts *options.Options) {
 				}
 				GetRoleCertificates(ztsUrl, opts)
 				util.ExecuteScriptWithoutBlock(opts.RunAfterParts, opts.RunAfterFailExit)
+				util.NotifySystemdReadyForCommand(cmd, "systemd-notify-all")
 
 				if opts.SDSUdsPath != "" {
 					certUpdates <- true

libs/go/sia/aws/agent/agent.go

@@ -698,13 +698,7 @@ func RunAgent(siaCmd, ztsUrl string, opts *options.Options) {
 			}
 		}
 
-		if cmd == "systemd-notify" {
-			err = util.NotifySystemdReady()
-			if err != nil {
-				log.Printf("failed to notify systemd: %v", err)
-			}
-		}
-
+		util.NotifySystemdReadyForCommand(cmd, "systemd-notify")
 		log.Printf("Identity established for services: %s\n", svcs)
 
 		stop := make(chan bool, 1)
@@ -757,6 +751,7 @@ func RunAgent(siaCmd, ztsUrl string, opts *options.Options) {
 				}
 				GetRoleCertificates(ztsUrl, opts)
 				util.ExecuteScriptWithoutBlock(opts.RunAfterParts, opts.RunAfterFailExit)
+				util.NotifySystemdReadyForCommand(cmd, "systemd-notify-all")
 
 				if opts.SDSUdsPath != "" {
 					certUpdates <- true

libs/go/sia/util/util.go

@@ -1330,7 +1330,22 @@ func ParseSiaCmd(siaCmd string) (string, bool) {
 	}
 }
 
-// NotifySystemDReady sends a notification to systemd that the service is ready
+// NotifySystemdReadyForCommand sends a notification to systemd that the
+// service is ready if the clientCmd argument matches to the notifyCmd
+func NotifySystemdReadyForCommand(clientCmd, notifyCmd string) error {
+	// if our client command is not the requested value
+	// then we're going to skip the notification
+	if clientCmd != notifyCmd {
+		return nil
+	}
+	err := NotifySystemdReady()
+	if err != nil {
+		log.Printf("failed to notify systemd: %v", err)
+	}
+	return err
+}
+
+// NotifySystemdReady sends a notification to systemd that the service is ready
 func NotifySystemdReady() error {
 	notifySocket := os.Getenv("NOTIFY_SOCKET")
 	if notifySocket == "" {

libs/go/sia/util/util_test.go

@@ -1793,4 +1793,18 @@ func TestNotifySystemdReady(test *testing.T) {
 	defer os.Remove(notifySocket)
 	err = NotifySystemdReady()
 	assert.Nil(test, err)
+	os.Unsetenv("NOTIFY_SOCKET")
+}
+
+func TestNotifySystemdReadyForCommand(test *testing.T) {
+
+	// when commands don't match there is always nil error
+	err := NotifySystemdReadyForCommand("systemd-notify", "init")
+	assert.Nil(test, err)
+
+	// when commands match, the code is executed, so we should
+	// get an error that notify socket is not set
+	err = NotifySystemdReadyForCommand("systemd-notify", "systemd-notify")
+	assert.NotNil(test, err)
+	assert.Equal(test, err.Error(), "notify socket is not set")
 }