From 67cc60fbb6d73d877889cec082e787eea93bc862 Mon Sep 17 00:00:00 2001
From: Vishal Nayak <vishalnayak@users.noreply.github.com>
Date: Tue, 6 Feb 2018 13:06:17 -0500
Subject: [PATCH] command/ssh: create and reuse the api client (#3909)

* pass around the api client

* reuse the client object in the base command
---
 command/base.go |  3 ++-
 command/ssh.go  | 28 ++++++++++------------------
 2 files changed, 12 insertions(+), 19 deletions(-)

diff --git a/command/base.go b/command/base.go
index 7dcca7671b5a..be6925b3e0d8 100644
--- a/command/base.go
+++ b/command/base.go
@@ -45,7 +45,6 @@ type BaseCommand struct {
 
 	tokenHelper token.TokenHelper
 
-	// For testing
 	client *api.Client
 }
 
@@ -110,6 +109,8 @@ func (c *BaseCommand) Client() (*api.Client, error) {
 		client.SetToken(token)
 	}
 
+	c.client = client
+
 	return client, nil
 }
 
diff --git a/command/ssh.go b/command/ssh.go
index c35fbb9cfd99..8d2922cecda3 100644
--- a/command/ssh.go
+++ b/command/ssh.go
@@ -244,6 +244,13 @@ func (c *SSHCommand) Run(args []string) int {
 		sshArgs = args[1:]
 	}
 
+	// Set the client in the command
+	_, err = c.Client()
+	if err != nil {
+		c.UI.Error(err.Error())
+		return 1
+	}
+
 	// Credentials are generated only against a registered role. If user
 	// does not specify a role with the SSH command, then lookup API is used
 	// to fetch all the roles with which this IP is associated. If there is
@@ -331,13 +338,7 @@ func (c *SSHCommand) handleTypeCA(username, ip string, sshArgs []string) int {
 		return 1
 	}
 
-	client, err := c.Client()
-	if err != nil {
-		c.UI.Error(err.Error())
-		return 1
-	}
-
-	sshClient := client.SSHWithMountPoint(c.flagMountPoint)
+	sshClient := c.client.SSHWithMountPoint(c.flagMountPoint)
 
 	// Attempt to sign the public key
 	secret, err := sshClient.SignKey(c.flagRole, map[string]interface{}{
@@ -611,12 +612,7 @@ func (c *SSHCommand) handleTypeDynamic(username, ip string, sshArgs []string) in
 // generateCredential generates a credential for the given role and returns the
 // decoded secret data.
 func (c *SSHCommand) generateCredential(username, ip string) (*api.Secret, *SSHCredentialResp, error) {
-	client, err := c.Client()
-	if err != nil {
-		return nil, nil, err
-	}
-
-	sshClient := client.SSHWithMountPoint(c.flagMountPoint)
+	sshClient := c.client.SSHWithMountPoint(c.flagMountPoint)
 
 	// Attempt to generate the credential.
 	secret, err := sshClient.Credential(c.flagRole, map[string]interface{}{
@@ -683,11 +679,7 @@ func (c *SSHCommand) defaultRole(mountPoint, ip string) (string, error) {
 	data := map[string]interface{}{
 		"ip": ip,
 	}
-	client, err := c.Client()
-	if err != nil {
-		return "", err
-	}
-	secret, err := client.Logical().Write(mountPoint+"/lookup", data)
+	secret, err := c.client.Logical().Write(mountPoint+"/lookup", data)
 	if err != nil {
 		return "", fmt.Errorf("Error finding roles for IP %q: %q", ip, err)