Add ability to input a session token from a cli flag

Emyrk · Emyrk · commit 1373030838ef · 2021-09-30T22:38:36.000-04:00
diff --git a/internal/cmd/agent.go b/internal/cmd/agent.go
@@ -41,10 +41,11 @@ func agentCmd() *cobra.Command {
 
 func startCmd() *cobra.Command {
 	var (
-		token    string
-		coderURL string
-		logFile  string
-		pullCert bool
+		token        string
+		sessionToken string
+		coderURL     string
+		logFile      string
+		pullCert     bool
 	)
 	cmd := &cobra.Command{
 		Use:   "start --coder-url=<coder_url> --token=<token> --log-file=<path>",
@@ -107,9 +108,9 @@ coder agent start --log-file=/tmp/coder-agent.log
 				}
 			}
 
-			// First inject certs
+			// First inject certs if enabled
 			if pullCert {
-				count, err := writeCoderdCerts(ctx)
+				count, err := writeCoderdCerts(ctx, sessionToken)
 				if err != nil {
 					return xerrors.Errorf("trust certs: %w", err)
 				}
@@ -139,16 +140,17 @@ coder agent start --log-file=/tmp/coder-agent.log
 	}
 
 	cmd.Flags().StringVar(&token, "token", "", "coder agent token")
+	cmd.Flags().StringVar(&sessionToken, "session-token", "", "coder session token to auth as user")
 	cmd.Flags().StringVar(&coderURL, "coder-url", "", "coder access url")
 	cmd.Flags().StringVar(&logFile, "log-file", "", "write a copy of logs to file")
 	cmd.Flags().BoolVar(&pullCert, "pull-cert", true, "pulls the tls certificate from coderd to ensure the cert is trusted")
 
 	return cmd
 }
 
-func writeCoderdCerts(ctx context.Context) (int, error) {
+func writeCoderdCerts(ctx context.Context, sessionToken string) (int, error) {
 	// Inject certs to custom dir and concat with : with existing dir.
-	certs, err := trustCertificate(ctx)
+	certs, err := trustCertificate(ctx, sessionToken)
 	if err != nil {
 		return 0, xerrors.Errorf("trust cert: %w", err)
 	}
@@ -187,7 +189,7 @@ func writeCoderdCerts(ctx context.Context) (int, error) {
 // It will then extend the certs to trust to include this directory.
 // This only happens if coderd can answer the challenge to prove
 // it has the shared secret.
-func trustCertificate(ctx context.Context) ([][]byte, error) {
+func trustCertificate(ctx context.Context, sessionToken string) ([][]byte, error) {
 	conf := &tls.Config{InsecureSkipVerify: true}
 	hc := &http.Client{
 		Timeout: time.Second * 3,
@@ -196,7 +198,7 @@ func trustCertificate(ctx context.Context) ([][]byte, error) {
 		},
 	}
 
-	c, err := newClient(ctx, false, withHTTPClient(hc))
+	c, err := newClient(ctx, false, withHTTPClient(hc), withSessionToken(sessionToken))
 	if err != nil {
 		return nil, xerrors.Errorf("new client: %w", err)
 	}
diff --git a/internal/cmd/auth.go b/internal/cmd/auth.go
@@ -25,6 +25,12 @@ const urlEnv = "CODER_URL"
 
 type coderOpt func(opt *coder.ClientOptions)
 
+func withSessionToken(token string) coderOpt {
+	return func(opt *coder.ClientOptions) {
+		opt.Token = token
+	}
+}
+
 func withHTTPClient(hc *http.Client) coderOpt {
 	return func(opt *coder.ClientOptions) {
 		opt.HTTPClient = hc
@@ -38,30 +44,39 @@ func newClient(ctx context.Context, checkVersion bool, optsF ...coderOpt) (coder
 		rawURL       = os.Getenv(urlEnv)
 	)
 
-	if sessionToken == "" || rawURL == "" {
-		sessionToken, err = config.Session.Read()
-		if err != nil {
-			return nil, errNeedLogin
-		}
+	clientOpts := &coder.ClientOptions{}
+	for _, f := range optsF {
+		f(clientOpts)
+	}
 
-		rawURL, err = config.URL.Read()
-		if err != nil {
-			return nil, errNeedLogin
+	// Missing the token
+	if clientOpts.Token == "" {
+		// If the env var is not set, try the config file
+		if sessionToken == "" {
+			sessionToken, err = config.Session.Read()
+			if err != nil {
+				return nil, errNeedLogin
+			}
 		}
+		clientOpts.Token = sessionToken
 	}
 
-	u, err := url.Parse(rawURL)
-	if err != nil {
-		return nil, xerrors.Errorf("url malformed: %w try running \"coder login\" with a valid URL", err)
-	}
+	// Missing the url
+	if clientOpts.BaseURL == nil {
+		// If the env var is not set, try the config file
+		if rawURL == "" {
+			rawURL, err = config.URL.Read()
+			if err != nil {
+				return nil, errNeedLogin
+			}
+		}
 
-	clientOpts := &coder.ClientOptions{
-		BaseURL: u,
-		Token:   sessionToken,
-	}
+		u, err := url.Parse(rawURL)
+		if err != nil {
+			return nil, xerrors.Errorf("url malformed: %w try running \"coder login\" with a valid URL", err)
+		}
 
-	for _, f := range optsF {
-		f(clientOpts)
+		clientOpts.BaseURL = u
 	}
 
 	c, err := coder.NewClient(*clientOpts)
