From a23a97efd60d656fba12e50619e9491644c7ae23 Mon Sep 17 00:00:00 2001
From: Angira Kekteeva <kira@nspcc.ru>
Date: Sun, 26 Dec 2021 23:38:16 +0300
Subject: [PATCH] [#290] Add timeout to authmate

Signed-off-by: Angira Kekteeva <kira@nspcc.ru>
---
 cmd/authmate/main.go | 21 ++++++++++++++++++---
 1 file changed, 18 insertions(+), 3 deletions(-)

diff --git a/cmd/authmate/main.go b/cmd/authmate/main.go
index ebf7b71a4..791a3795e 100644
--- a/cmd/authmate/main.go
+++ b/cmd/authmate/main.go
@@ -49,6 +49,7 @@ var (
 	lifetimeFlag           time.Duration
 	containerPolicies      string
 	awcCliCredFile         string
+	timeoutFlag            time.Duration
 )
 
 const (
@@ -124,6 +125,13 @@ func appFlags() []cli.Flag {
 			Usage:       "Enable debug logger level",
 			Destination: &logDebugEnabledFlag,
 		},
+		&cli.DurationFlag{
+			Name: "timeout",
+			Usage: "timeout of processing of the command, for example 2m " +
+				"(note: max time unit is an hour so to set a day you should use 24h)",
+			Destination: &timeoutFlag,
+			Value:       1 * time.Minute,
+		},
 	}
 }
 
@@ -276,10 +284,13 @@ It will be ceil rounded to the nearest amount of epoch.`,
 				AwsCliCredentialsFile: awcCliCredFile,
 			}
 
-			if err = agent.IssueSecret(ctx, os.Stdout, issueSecretOptions); err != nil {
-				return cli.Exit(fmt.Sprintf("failed to issue secret: %s", err), 6)
-			}
+			var tcancel context.CancelFunc
+			ctx, tcancel = context.WithTimeout(ctx, timeoutFlag)
+			defer tcancel()
 
+			if err = agent.IssueSecret(ctx, os.Stdout, issueSecretOptions); err != nil {
+				return cli.Exit(fmt.Sprintf("failed to issue secret: %s", err), 7)
+			}
 			return nil
 		},
 	}
@@ -391,6 +402,10 @@ func obtainSecret() *cli.Command {
 				GatePrivateKey: gateCreds,
 			}
 
+			var tcancel context.CancelFunc
+			ctx, tcancel = context.WithTimeout(ctx, timeoutFlag)
+			defer tcancel()
+
 			if err = agent.ObtainSecret(ctx, os.Stdout, obtainSecretOptions); err != nil {
 				return cli.Exit(fmt.Sprintf("failed to obtain secret: %s", err), 5)
 			}