Skip to content

Commit

Permalink
provide sia config option to exit process if run_after script fails (#…
Browse files Browse the repository at this point in the history 
…2580)
  • Loading branch information
@havetisyan
havetisyan authored Apr 3, 2024
1 parent b2d339a commit a16110d
Show file tree
Hide file tree
Showing 6 changed files with 42 additions and 23 deletions.
18 changes: 9 additions & 9 deletions libs/go/sia/agent/agent.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.ExecuteScript(opts.RunAfterParts)
util.ExecuteScript(opts.RunAfterParts, opts.RunAfterFailExit)
}
case "token":
if tokenOpts != nil {
err := fetchAccessToken(tokenOpts)
if err != nil && !skipErrors {
log.Fatalf("Unable to fetch access tokens, err: %v\n", err)
}
util.ExecuteScript(opts.RunAfterTokensParts)
util.ExecuteScript(opts.RunAfterTokensParts, opts.RunAfterFailExit)
} else {
log.Print("unable to fetch access tokens, invalid or missing configuration")
}
Expand All @@ -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.ExecuteScript(opts.RunAfterParts)
util.ExecuteScript(opts.RunAfterParts, opts.RunAfterFailExit)
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.ExecuteScript(opts.RunAfterParts)
util.ExecuteScript(opts.RunAfterParts, opts.RunAfterFailExit)
log.Printf("Identity successfully refreshed for services: %s\n", svcs)
case "init":
err := RegisterInstance(ztsUrl, opts, false)
Expand All @@ -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.ExecuteScript(opts.RunAfterParts)
util.ExecuteScript(opts.RunAfterParts, opts.RunAfterFailExit)
if tokenOpts != nil {
err := fetchAccessToken(tokenOpts)
if err != nil && !skipErrors {
log.Fatalf("Unable to fetch access tokens, err: %v\n", err)
}
util.ExecuteScript(opts.RunAfterTokensParts)
util.ExecuteScript(opts.RunAfterTokensParts, opts.RunAfterFailExit)
}
default:
// we're going to iterate through our configured services.
Expand Down Expand Up @@ -741,13 +741,13 @@ func RunAgent(siaCmd, ztsUrl string, opts *options.Options) {
if err != nil {
errors <- fmt.Errorf("Unable to fetch access tokens after identity refresh, err: %v\n", err)
} else {
util.ExecuteScriptWithoutBlock(opts.RunAfterTokensParts)
util.ExecuteScriptWithoutBlock(opts.RunAfterTokensParts, opts.RunAfterFailExit)
}
} else {
log.Print("token config does not exist - do not refresh tokens")
}
GetRoleCertificates(ztsUrl, opts)
util.ExecuteScriptWithoutBlock(opts.RunAfterParts)
util.ExecuteScriptWithoutBlock(opts.RunAfterParts, opts.RunAfterFailExit)

if opts.SDSUdsPath != "" {
certUpdates <- true
Expand Down Expand Up @@ -798,7 +798,7 @@ func RunAgent(siaCmd, ztsUrl string, opts *options.Options) {
if err != nil {
errors <- fmt.Errorf("refresh access-token task got error: %v\n", err)
} else {
util.ExecuteScriptWithoutBlock(opts.RunAfterTokensParts)
util.ExecuteScriptWithoutBlock(opts.RunAfterTokensParts, opts.RunAfterFailExit)
}
case <-stop:
errors <- nil
Expand Down
18 changes: 9 additions & 9 deletions libs/go/sia/aws/agent/agent.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.ExecuteScript(opts.RunAfterParts)
util.ExecuteScript(opts.RunAfterParts, opts.RunAfterFailExit)
}
case "token":
if tokenOpts != nil {
err := fetchAccessToken(tokenOpts)
if err != nil && !skipErrors {
log.Fatalf("Unable to fetch access tokens, err: %v\n", err)
}
util.ExecuteScript(opts.RunAfterTokensParts)
util.ExecuteScript(opts.RunAfterTokensParts, opts.RunAfterFailExit)
} else {
log.Print("unable to fetch access tokens, invalid or missing configuration")
}
Expand All @@ -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.ExecuteScript(opts.RunAfterParts)
util.ExecuteScript(opts.RunAfterParts, opts.RunAfterFailExit)
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.ExecuteScript(opts.RunAfterParts)
util.ExecuteScript(opts.RunAfterParts, opts.RunAfterFailExit)
log.Printf("Identity successfully refreshed for services: %s\n", svcs)
case "init":
err := RegisterInstance(data, ztsUrl, opts, false)
Expand All @@ -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.ExecuteScript(opts.RunAfterParts)
util.ExecuteScript(opts.RunAfterParts, opts.RunAfterFailExit)
if tokenOpts != nil {
err := fetchAccessToken(tokenOpts)
if err != nil && !skipErrors {
log.Fatalf("Unable to fetch access tokens, err: %v\n", err)
}
util.ExecuteScript(opts.RunAfterTokensParts)
util.ExecuteScript(opts.RunAfterTokensParts, opts.RunAfterFailExit)
}
default:
// we're going to iterate through our configured services.
Expand Down Expand Up @@ -742,13 +742,13 @@ func RunAgent(siaCmd, ztsUrl string, opts *options.Options) {
if err != nil {
errors <- fmt.Errorf("Unable to fetch access token after identity refresh, err: %v\n", err)
} else {
util.ExecuteScriptWithoutBlock(opts.RunAfterTokensParts)
util.ExecuteScriptWithoutBlock(opts.RunAfterTokensParts, opts.RunAfterFailExit)
}
} else {
log.Print("token config does not exist - do not refresh token")
}
GetRoleCertificates(ztsUrl, opts)
util.ExecuteScriptWithoutBlock(opts.RunAfterParts)
util.ExecuteScriptWithoutBlock(opts.RunAfterParts, opts.RunAfterFailExit)

if opts.SDSUdsPath != "" {
certUpdates <- true
Expand Down Expand Up @@ -799,7 +799,7 @@ func RunAgent(siaCmd, ztsUrl string, opts *options.Options) {
if err != nil {
errors <- fmt.Errorf("refresh access-token task got error: %v\n", err)
} else {
util.ExecuteScriptWithoutBlock(opts.RunAfterTokensParts)
util.ExecuteScriptWithoutBlock(opts.RunAfterTokensParts, opts.RunAfterFailExit)
}
case <-stop:
errors <- nil
Expand Down
8 changes: 8 additions & 0 deletions libs/go/sia/aws/options/options.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -113,6 +113,7 @@ type Config struct {
RunAfterTokens string `json:"run_after_tokens,omitempty"` //execute the command mentioned after tokens are created
SpiffeTrustDomain string `json:"spiffe_trust_domain,omitempty"` //spiffe trust domain - if configured used full spiffe uri with namespace
StoreTokenOption *int `json:"store_token_option,omitempty"` //store access token option
RunAfterFailExit bool `json:"run_after_fail_exit,omitempty"` //exit process if run_after script fails
}

type AccessProfileConfig struct {
Expand Down Expand Up @@ -218,6 +219,7 @@ type Options struct {
SpiffeNamespace string //spiffe uri namespace
OmitDomain bool //attestation role only includes service name
StoreTokenOption *int //store access token option
RunAfterFailExit bool //exit process if run_after script fails
}

const (
Expand Down Expand Up @@ -452,6 +454,9 @@ func InitEnvConfig(config *Config) (*Config, *ConfigAccount, error) {
if !config.AccessManagement {
config.AccessManagement = util.ParseEnvBooleanFlag("ATHENZ_SIA_ACCESS_MANAGEMENT")
}
if !config.RunAfterFailExit {
config.RunAfterFailExit = util.ParseEnvBooleanFlag("ATHENZ_SIA_RUN_AFTER_FAIL_EXIT")
}

roleArn := os.Getenv("ATHENZ_SIA_IAM_ROLE_ARN")
if roleArn == "" {
Expand Down Expand Up @@ -549,6 +554,7 @@ func setOptions(config *Config, account *ConfigAccount, profileConfig *AccessPro
runAfterTokens := ""
spiffeTrustDomain := ""
addlSanDNSEntries := make([]string, 0)
runAfterFailExit := false

var storeTokenOption *int
if config != nil {
Expand All @@ -564,6 +570,7 @@ func setOptions(config *Config, account *ConfigAccount, profileConfig *AccessPro
fileDirectUpdate = config.FileDirectUpdate
accessManagement = config.AccessManagement
storeTokenOption = config.StoreTokenOption
runAfterFailExit = config.RunAfterFailExit

if config.RefreshInterval > 0 {
refreshInterval = config.RefreshInterval
Expand Down Expand Up @@ -781,6 +788,7 @@ func setOptions(config *Config, account *ConfigAccount, profileConfig *AccessPro
OmitDomain: account.OmitDomain,
StoreTokenOption: storeTokenOption,
AddlSanDNSEntries: addlSanDNSEntries,
RunAfterFailExit: runAfterFailExit,
}, nil
}

Expand Down
8 changes: 8 additions & 0 deletions libs/go/sia/options/options.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -121,6 +121,7 @@ type Config struct {
RunAfterTokens string `json:"run_after_tokens,omitempty"` //execute the command mentioned after tokens are created
SpiffeTrustDomain string `json:"spiffe_trust_domain,omitempty"` //spiffe trust domain - if configured generate full spiffe uri with namespace
StoreTokenOption *int `json:"store_token_option,omitempty"` //store access token option
RunAfterFailExit bool `json:"run_after_fail_exit,omitempty"` //exit process if run_after script fails
}

type AccessProfileConfig struct {
Expand Down Expand Up @@ -228,6 +229,7 @@ type Options struct {
SpiffeNamespace string //spiffe uri namespace
OmitDomain bool //attestation role only includes service name
StoreTokenOption *int //store access token option
RunAfterFailExit bool //exit process if run_after script fails
}

const (
Expand Down Expand Up @@ -500,6 +502,9 @@ func InitEnvConfig(config *Config, provider provider.Provider) (*Config, *Config
if !config.AccessManagement {
config.AccessManagement = util.ParseEnvBooleanFlag("ATHENZ_SIA_ACCESS_MANAGEMENT")
}
if !config.RunAfterFailExit {
config.RunAfterFailExit = util.ParseEnvBooleanFlag("ATHENZ_SIA_RUN_AFTER_FAIL_EXIT")
}

config.Threshold = util.ParseEnvFloatFlag("ATHENZ_SIA_ACCOUNT_THRESHOLD", DefaultThreshold)
config.SshThreshold = util.ParseEnvFloatFlag("ATHENZ_SIA_ACCOUNT_SSH_THRESHOLD", DefaultThreshold)
Expand Down Expand Up @@ -610,6 +615,7 @@ func setOptions(config *Config, account *ConfigAccount, profileConfig *AccessPro
runAfterTokens := ""
spiffeTrustDomain := ""
addlSanDNSEntries := make([]string, 0)
runAfterFailExit := false

var storeTokenOption *int
if config != nil {
Expand All @@ -625,6 +631,7 @@ func setOptions(config *Config, account *ConfigAccount, profileConfig *AccessPro
fileDirectUpdate = config.FileDirectUpdate
accessManagement = config.AccessManagement
storeTokenOption = config.StoreTokenOption
runAfterFailExit = config.RunAfterFailExit

if config.RefreshInterval > 0 {
refreshInterval = config.RefreshInterval
Expand Down Expand Up @@ -850,6 +857,7 @@ func setOptions(config *Config, account *ConfigAccount, profileConfig *AccessPro
OmitDomain: account.OmitDomain,
StoreTokenOption: storeTokenOption,
AddlSanDNSEntries: addlSanDNSEntries,
RunAfterFailExit: runAfterFailExit,
}, nil
}

Expand Down
9 changes: 6 additions & 3 deletions libs/go/sia/util/util.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1282,7 +1282,7 @@ func ParseScriptArguments(script string) []string {

// ExecuteScript executes a script along with the provided
// arguments while blocking the agent
func ExecuteScript(script []string) error {
func ExecuteScript(script []string, runAfterFailExit bool) error {
// execute run after script (if provided)
if len(script) == 0 {
return nil
Expand All @@ -1291,15 +1291,18 @@ func ExecuteScript(script []string) error {
err := exec.Command(script[0], script[1:]...).Run()
if err != nil {
log.Printf("unable to execute: %q, err: %v", script, err)
if runAfterFailExit {
os.Exit(1)
}
}
return err
}

// ExecuteScriptWithoutBlock executes a script along with the provided
// arguments in a go subroutine without blocking the agent
func ExecuteScriptWithoutBlock(script []string) {
func ExecuteScriptWithoutBlock(script []string, runAfterFailExit bool) {
go func() {
ExecuteScript(script)
ExecuteScript(script, runAfterFailExit)
}()
}

Expand Down
4 changes: 2 additions & 2 deletions libs/go/sia/util/util_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1685,13 +1685,13 @@ func TestParseSiaCmd(test *testing.T) {
func TestExecuteScript(test *testing.T) {

// non-existent script
err := ExecuteScript([]string{"unknown-script"})
err := ExecuteScript([]string{"unknown-script"}, false)
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"})
err = ExecuteScript([]string{"data/test_after_script.sh"}, true)
assert.Nil(test, err)
// verify our test file was created
_, err = os.Stat("/tmp/test-after-script")
Expand Down

0 comments on commit a16110d

Please sign in to comment.