More flexible SSH Agent handling

runner.go

@@ -217,13 +217,15 @@ func (c *CorkTypeContainer) startSSHCommand(stageName string) error {
  Failed: failed,
  SSHKeyPath: c.SSHKeyPath,
  }
+ log.Debugf("ssh command options: %+v", sshCommandOptions)
  command, err := NewDockerSSHCommand(sshCommandOptions)
  if err != nil {
  return err
  }
 
  log.Debugf("Running SSH with env %v", c.Env)
  command.Start(c.Env)
+ defer command.CleanUp()
  clientErrChan := make(chan error)
 
  c.runClient(stageName, clientErrChan)

ssh.go

@@ -1,7 +1,10 @@
 package main
 
 import (
+ "encoding/pem"
+ "errors"
  "fmt"
+ "io/ioutil"
  "net"
  "os"
  "os/user"
@@ -12,21 +15,24 @@ import (
 
  "io"
 
+ "crypto/x509"
+
  "golang.org/x/crypto/ssh"
  "golang.org/x/crypto/ssh/agent"
 )
 
 var maxRetries = 15
 
 type DockerSSHCommand struct {
- Host string
- Port int
- Command string
- Failed chan bool
- SSHKeyPath string
- Stdin io.Reader
- Stdout io.Writer
- Stderr io.Writer
+ Host string
+ Port int
+ Command string
+ Failed chan bool
+ SSHKeyPath string
+ Stdin io.Reader
+ Stdout io.Writer
+ Stderr io.Writer
+ SSHAgentManager SSHAgentManager
 }
 
 type DockerSSHCommandOptions struct {
@@ -37,11 +43,112 @@ type DockerSSHCommandOptions struct {
  SSHKeyPath string
 }
 
+type SSHAgentManager interface {
+ Shutdown() error
+ GetSSHAgent() agent.Agent
+ Forward(client *ssh.Client) error
+}
+
+type SystemSSHAgentManager struct {
+ Connection net.Conn
+ Agent agent.Agent
+ SSHAuthSockPath string
+}
+
+func NewSystemSSHAgentManager(expectedSSHKeyPath string, sshAuthSockPath string) (SSHAgentManager, error) {
+ sshAgent, err := net.Dial("unix", sshAuthSockPath)
+ if err != nil {
+ return nil, err
+ }
+
+ agentClient := agent.NewClient(sshAgent)
+ agentKeys, err := agentClient.List()
+ if err != nil {
+ return nil, err
+ }
+
+ found := false
+ for _, key := range agentKeys {
+ if key.Comment == expectedSSHKeyPath {
+ found = true
+ break
+ }
+ }
+
+ if !found {
+ return nil, fmt.Errorf("CannotRunSSHCommand: Could not find appropriate key")
+ }
+ return &SystemSSHAgentManager{
+ Agent: agentClient,
+ Connection: sshAgent,
+ SSHAuthSockPath: sshAuthSockPath,
+ }, nil
+}
+
+func (s *SystemSSHAgentManager) Shutdown() error {
+ return s.Connection.Close()
+}
+
+func (s *SystemSSHAgentManager) GetSSHAgent() agent.Agent {
+ return s.Agent
+}
+
+func (s *SystemSSHAgentManager) Forward(client *ssh.Client) error {
+ return agent.ForwardToRemote(client, s.SSHAuthSockPath)
+}
+
+type InMemorySSHAgentManager struct {
+ Agent agent.Agent
+}
+
+func NewInMemorySSHAgentManager(pathToPrivateKey string) (SSHAgentManager, error) {
+ privateKeyPemBytes, err := ioutil.ReadFile(pathToPrivateKey)
+ if err != nil {
+ return nil, err
+ }
+
+ block, _ := pem.Decode(privateKeyPemBytes)
+ if block == nil {
+ return nil, errors.New("failed to parse PEM block containing the key")
+ }
+
+ privateKey, err := x509.ParsePKCS1PrivateKey(block.Bytes)
+ if err != nil {
+ return nil, err
+ }
+
+ addedKey := agent.AddedKey{
+ PrivateKey: privateKey,
+ Comment: pathToPrivateKey,
+ }
+
+ agentKeyring := agent.NewKeyring()
+ agentKeyring.Add(addedKey)
+
+ return &InMemorySSHAgentManager{
+ Agent: agentKeyring,
+ }, nil
+}
+
+func (i *InMemorySSHAgentManager) Shutdown() error {
+ return nil
+}
+
+func (i *InMemorySSHAgentManager) GetSSHAgent() agent.Agent {
+ return i.Agent
+}
+
+func (i *InMemorySSHAgentManager) Forward(client *ssh.Client) error {
+ return agent.ForwardToAgent(client, i.Agent)
+}
+
 const DefaultSSHPath = ".ssh/id_rsa"
 
 func NewDockerSSHCommand(options DockerSSHCommandOptions) (*DockerSSHCommand, error) {
+ log.Debugf("Creating NewDockerSSHCommand")
  // FIXME: We should have an SSH Search Path
  if options.SSHKeyPath == "" {
+ log.Debugf("Loading current user")
  usr, err := user.Current()
  if err != nil {
  return nil, err
@@ -59,7 +166,7 @@ func NewDockerSSHCommand(options DockerSSHCommandOptions) (*DockerSSHCommand, er
  SSHKeyPath: options.SSHKeyPath,
  }
 
- err := sshCommand.PerformSelfTest()
+ err := sshCommand.ChooseSSHAgentManager()
  if err != nil {
  return nil, err
  }
@@ -69,6 +176,8 @@ func NewDockerSSHCommand(options DockerSSHCommandOptions) (*DockerSSHCommand, er
 func (d *DockerSSHCommand) Start(envVars []string) {
  // Connect to the ssh
  go func() {
+ log.SetLevel(log.DebugLevel)
+ log.Debug("Starting the docker ssh command")
  err := d.runCommand(envVars)
  if err != nil {
  log.Errorf("Error running ssh command: %v", err)
@@ -114,7 +223,7 @@ func (d *DockerSSHCommand) newSession() (*ssh.Session, error) {
 
  log.Debug("Setting up forwarding for the ssh connection")
  // Setup agent forwarding
- err = agent.ForwardToRemote(connection, os.Getenv("SSH_AUTH_SOCK"))
+ err = d.SSHAgentManager.Forward(connection)
  if err != nil {
  return nil, fmt.Errorf("Failed to forward agent: %s", err)
  }
@@ -153,19 +262,19 @@ func (d *DockerSSHCommand) runCommand(envVars []string) error {
  defer session.Close()
 
  /*
-  for _, envVar := range envVars {
-  splitEnvVar := strings.Split(envVar, "=")
-  if len(splitEnvVar) < 2 {
-  return fmt.Errorf("Invalid env var `%s`", envVar)
-  }
-  envKey := splitEnvVar[0]
-  envValue := strings.Join(splitEnvVar[1:], "=")
-  log.Debugf("Setting `%s=%s` on ssh env", envKey, envValue)
-  err = session.Setenv(envKey, envValue)
-  if err != nil {
-  return err
-  }
-  }
+  for _, envVar := range envVars {
+  splitEnvVar := strings.Split(envVar, "=")
+  if len(splitEnvVar) < 2 {
+  return fmt.Errorf("Invalid env var `%s`", envVar)
+  }
+  envKey := splitEnvVar[0]
+  envValue := strings.Join(splitEnvVar[1:], "=")
+  log.Debugf("Setting `%s=%s` on ssh env", envKey, envValue)
+  err = session.Setenv(envKey, envValue)
+  if err != nil {
+  return err
+  }
+  }
  */
 
  if d.Stdin != nil {
@@ -195,39 +304,34 @@ func (d *DockerSSHCommand) runCommand(envVars []string) error {
  return session.Run(d.Command)
 }
 
-// PerformSelfTest determines if the DockerSSHCommand can run
-func (d *DockerSSHCommand) PerformSelfTest() error {
- sshAgent, err := net.Dial("unix", os.Getenv("SSH_AUTH_SOCK"))
- if err != nil {
- return err
- }
- defer sshAgent.Close()
+func (d *DockerSSHCommand) CleanUp() error {
+ return d.SSHAgentManager.Shutdown()
+}
 
- agentClient := agent.NewClient(sshAgent)
- agentKeys, err := agentClient.List()
- if err != nil {
- return err
- }
+func (d *DockerSSHCommand) ChooseSSHAgentManager() error {
+ log.Debugf("Choosing agent based on capabilities")
 
- found := false
- for _, key := range agentKeys {
- if key.Comment == d.SSHKeyPath {
- found = true
- break
+ // First check if there is an existing agent
+ // If not check if there's an ssh key on the system at ~/.ssh/id_rsa
+ // fail if that key is encrypted.
+
+ sshAuthSockPath := os.Getenv("SSH_AUTH_SOCK")
+ if sshAuthSockPath != "" {
+ agentManager, err := NewSystemSSHAgentManager(d.SSHKeyPath, sshAuthSockPath)
+ if err == nil {
+ d.SSHAgentManager = agentManager
+ return nil
  }
  }
-
- if !found {
- return fmt.Errorf("CannotRunSSHCommand: Could not find appropriate key")
+ agentManager, err := NewInMemorySSHAgentManager(d.SSHKeyPath)
+ if err != nil {
+ return fmt.Errorf("Error choosing ssh agent. If your ssh key is passphrase protected you will need to start the ssh-agent: %+v", err)
  }
-
+ d.SSHAgentManager = agentManager
  return nil
 }
 
 func (d *DockerSSHCommand) getAuthMethod() (ssh.AuthMethod, error) {
- sshAgent, err := net.Dial("unix", os.Getenv("SSH_AUTH_SOCK"))
- if err != nil {
- return nil, err
- }
- return ssh.PublicKeysCallback(agent.NewClient(sshAgent).Signers), err
+ log.Debug("Attempting to connect to ssh agent")
+ return ssh.PublicKeysCallback(d.SSHAgentManager.GetSSHAgent().Signers), nil
 }