Handle errors on failing to accept connections

* permanent errors do not result in a retry
* Add -N flag for non-interactive ssh

[#90425212]

cf-plugin/cmd/fakes/fake_secure_client.go

@@ -13,14 +13,14 @@ type FakeSecureClient struct {
 	NewSessionStub        func() (cmd.SecureSession, error)
 	newSessionMutex       sync.RWMutex
 	newSessionArgsForCall []struct{}
-	newSessionReturns     struct {
+	newSessionReturns struct {
 		result1 cmd.SecureSession
 		result2 error
 	}
 	ConnStub        func() ssh.Conn
 	connMutex       sync.RWMutex
 	connArgsForCall []struct{}
-	connReturns     struct {
+	connReturns struct {
 		result1 ssh.Conn
 	}
 	DialStub        func(network, address string) (net.Conn, error)
@@ -33,10 +33,16 @@ type FakeSecureClient struct {
 		result1 net.Conn
 		result2 error
 	}
+	WaitStub        func() error
+	waitMutex       sync.RWMutex
+	waitArgsForCall []struct{}
+	waitReturns struct {
+		result1 error
+	}
 	CloseStub        func() error
 	closeMutex       sync.RWMutex
 	closeArgsForCall []struct{}
-	closeReturns     struct {
+	closeReturns struct {
 		result1 error
 	}
 }
@@ -124,6 +130,30 @@ func (fake *FakeSecureClient) DialReturns(result1 net.Conn, result2 error) {
 	}{result1, result2}
 }
 
+func (fake *FakeSecureClient) Wait() error {
+	fake.waitMutex.Lock()
+	fake.waitArgsForCall = append(fake.waitArgsForCall, struct{}{})
+	fake.waitMutex.Unlock()
+	if fake.WaitStub != nil {
+		return fake.WaitStub()
+	} else {
+		return fake.waitReturns.result1
+	}
+}
+
+func (fake *FakeSecureClient) WaitCallCount() int {
+	fake.waitMutex.RLock()
+	defer fake.waitMutex.RUnlock()
+	return len(fake.waitArgsForCall)
+}
+
+func (fake *FakeSecureClient) WaitReturns(result1 error) {
+	fake.WaitStub = nil
+	fake.waitReturns = struct {
+		result1 error
+	}{result1}
+}
+
 func (fake *FakeSecureClient) Close() error {
 	fake.closeMutex.Lock()
 	fake.closeArgsForCall = append(fake.closeArgsForCall, struct{}{})

cf-plugin/cmd/fakes/fake_secure_session.go

@@ -35,21 +35,21 @@ type FakeSecureSession struct {
 	StdinPipeStub        func() (io.WriteCloser, error)
 	stdinPipeMutex       sync.RWMutex
 	stdinPipeArgsForCall []struct{}
-	stdinPipeReturns     struct {
+	stdinPipeReturns struct {
 		result1 io.WriteCloser
 		result2 error
 	}
 	StdoutPipeStub        func() (io.Reader, error)
 	stdoutPipeMutex       sync.RWMutex
 	stdoutPipeArgsForCall []struct{}
-	stdoutPipeReturns     struct {
+	stdoutPipeReturns struct {
 		result1 io.Reader
 		result2 error
 	}
 	StderrPipeStub        func() (io.Reader, error)
 	stderrPipeMutex       sync.RWMutex
 	stderrPipeArgsForCall []struct{}
-	stderrPipeReturns     struct {
+	stderrPipeReturns struct {
 		result1 io.Reader
 		result2 error
 	}
@@ -64,19 +64,19 @@ type FakeSecureSession struct {
 	ShellStub        func() error
 	shellMutex       sync.RWMutex
 	shellArgsForCall []struct{}
-	shellReturns     struct {
+	shellReturns struct {
 		result1 error
 	}
 	WaitStub        func() error
 	waitMutex       sync.RWMutex
 	waitArgsForCall []struct{}
-	waitReturns     struct {
+	waitReturns struct {
 		result1 error
 	}
 	CloseStub        func() error
 	closeMutex       sync.RWMutex
 	closeArgsForCall []struct{}
-	closeReturns     struct {
+	closeReturns struct {
 		result1 error
 	}
 }

cf-plugin/cmd/ssh.go

@@ -29,6 +29,7 @@ type SecureShell interface {
 	Connect(opts *options.SSHOptions) error
 	InteractiveSession() error
 	LocalPortForward() error
+	Wait() error
 	Close() error
 }
 
@@ -42,6 +43,7 @@ type SecureClient interface {
 	NewSession() (SecureSession, error)
 	Conn() ssh.Conn
 	Dial(network, address string) (net.Conn, error)
+	Wait() error
 	Close() error
 }
 
@@ -159,10 +161,18 @@ func (c *secureShell) LocalPortForward() error {
 }
 
 func (c *secureShell) localForwardAcceptLoop(listener net.Listener, addr string) {
+	defer listener.Close()
+
 	for {
 		conn, err := listener.Accept()
 		if err == nil {
 			go c.handleForwardConnection(conn, addr)
+		} else {
+			if netErr, ok := err.(net.Error); ok && netErr.Temporary() {
+				time.Sleep(100 * time.Millisecond)
+				continue
+			}
+			return
 		}
 	}
 }
@@ -286,6 +296,10 @@ func (c *secureShell) InteractiveSession() error {
 	return session.Wait()
 }
 
+func (c *secureShell) Wait() error {
+	return c.secureClient.Wait()
+}
+
 func (c *secureShell) validateTarget(app app.App, opts *options.SSHOptions) error {
 	if app.State != "STARTED" {
 		return fmt.Errorf("Application %q is not in the STARTED state", opts.AppName)
@@ -422,6 +436,7 @@ type secureClient struct{ client *ssh.Client }
 
 func (sc *secureClient) Close() error   { return sc.client.Close() }
 func (sc *secureClient) Conn() ssh.Conn { return sc.client.Conn }
+func (sc *secureClient) Wait() error    { return sc.client.Wait() }
 func (sc *secureClient) Dial(n, addr string) (net.Conn, error) {
 	return sc.client.Dial(n, addr)
 }

cf-plugin/cmd/ssh_test.go

@@ -22,6 +22,7 @@ import (
 	"github.com/cloudfoundry-incubator/diego-ssh/cf-plugin/terminal/terminal_helper_fakes"
 	"github.com/cloudfoundry-incubator/diego-ssh/server"
 	fake_server "github.com/cloudfoundry-incubator/diego-ssh/server/fakes"
+	"github.com/cloudfoundry-incubator/diego-ssh/test_helpers"
 	"github.com/cloudfoundry-incubator/diego-ssh/test_helpers/fake_io"
 	"github.com/cloudfoundry-incubator/diego-ssh/test_helpers/fake_net"
 	"github.com/cloudfoundry-incubator/diego-ssh/test_helpers/fake_ssh"
@@ -873,6 +874,7 @@ var _ = Describe("Diego SSH Plugin", func() {
 			echoListener.AddrStub = realListener.Addr
 
 			fakeLocalListener = &fake_net.FakeListener{}
+			fakeLocalListener.AcceptReturns(nil, errors.New("Not Accepting Connections"))
 
 			echoServer = server.NewServer(logger.Session("echo"), "", echoHandler)
 			echoServer.SetListener(echoListener)
@@ -981,7 +983,7 @@ var _ = Describe("Diego SSH Plugin", func() {
 				realLocalListener2, err = net.Listen("tcp", "127.0.0.1:0")
 				Expect(err).NotTo(HaveOccurred())
 
-				localAddress2 = realLocalListener.Addr().String()
+				localAddress2 = realLocalListener2.Addr().String()
 
 				fakeListenerFactory.ListenStub = func(network, addr string) (net.Listener, error) {
 					if addr == localAddress {
@@ -1030,20 +1032,18 @@ var _ = Describe("Diego SSH Plugin", func() {
 
 			Context("when the secure client is closed", func() {
 				BeforeEach(func() {
-					fakeConn := &fake_net.FakeConn{}
-					fakeConn.ReadReturns(0, io.EOF)
-
-					fakeLocalListener.AcceptReturns(fakeConn, nil)
 					fakeListenerFactory.ListenReturns(fakeLocalListener, nil)
+					fakeLocalListener.AcceptReturns(nil, errors.New("not accepting connections"))
 				})
 
 				It("closes the listeners ", func() {
 					Eventually(fakeListenerFactory.ListenCallCount).Should(Equal(2))
-					Eventually(fakeLocalListener.AcceptCallCount).Should(BeNumerically(">=", 2))
+					Eventually(fakeLocalListener.AcceptCallCount).Should(Equal(2))
 
+					originalCloseCount := fakeLocalListener.CloseCallCount()
 					err := secureShell.Close()
 					Expect(err).NotTo(HaveOccurred())
-					Expect(fakeLocalListener.CloseCallCount()).Should(Equal(2))
+					Expect(fakeLocalListener.CloseCallCount()).Should(Equal(originalCloseCount + 2))
 				})
 			})
 		})
@@ -1059,43 +1059,73 @@ var _ = Describe("Diego SSH Plugin", func() {
 		})
 
 		Context("when the client it closed", func() {
-			var fakeConn *fake_net.FakeConn
-
 			BeforeEach(func() {
-				fakeConn = &fake_net.FakeConn{}
-				fakeConn.ReadReturns(0, io.EOF)
-
-				fakeLocalListener.AcceptReturns(fakeConn, nil)
 				fakeListenerFactory.ListenReturns(fakeLocalListener, nil)
+				fakeLocalListener.AcceptReturns(nil, errors.New("not accepting and connections"))
 			})
 
 			It("closes the listener when the client is closed", func() {
 				Eventually(fakeListenerFactory.ListenCallCount).Should(Equal(1))
+				Eventually(fakeLocalListener.AcceptCallCount).Should(Equal(1))
 
+				originalCloseCount := fakeLocalListener.CloseCallCount()
 				err := secureShell.Close()
 				Expect(err).NotTo(HaveOccurred())
-				Expect(fakeLocalListener.CloseCallCount()).Should(Equal(1))
+				Expect(fakeLocalListener.CloseCallCount()).Should(Equal(originalCloseCount + 1))
 			})
 		})
 
 		Context("when accept fails", func() {
 			var fakeConn *fake_net.FakeConn
-
 			BeforeEach(func() {
 				fakeConn = &fake_net.FakeConn{}
 				fakeConn.ReadReturns(0, io.EOF)
 
-				fakeLocalListener.AcceptStub = func() (net.Conn, error) {
-					if fakeLocalListener.AcceptCallCount() == 1 {
-						return nil, errors.New("boom")
-					}
-					return fakeConn, nil
-				}
 				fakeListenerFactory.ListenReturns(fakeLocalListener, nil)
 			})
 
-			PIt("continues to listen for new connections", func() {
-				Eventually(fakeLocalListener.AcceptCallCount).Should(BeNumerically(">", 1))
+			Context("with a permanent error", func() {
+				BeforeEach(func() {
+					fakeLocalListener.AcceptReturns(nil, errors.New("boom"))
+				})
+
+				It("stops trying to accept connections", func() {
+					Eventually(fakeLocalListener.AcceptCallCount).Should(Equal(1))
+					Consistently(fakeLocalListener.AcceptCallCount).Should(Equal(1))
+					Expect(fakeLocalListener.CloseCallCount()).To(Equal(1))
+				})
+			})
+
+			Context("with a temporary error", func() {
+				var timeCh chan time.Time
+
+				BeforeEach(func() {
+					timeCh = make(chan time.Time, 3)
+
+					fakeLocalListener.AcceptStub = func() (net.Conn, error) {
+						timeCh := timeCh
+						if fakeLocalListener.AcceptCallCount() > 3 {
+							close(timeCh)
+							return nil, test_helpers.NewTestNetError(false, false)
+						} else {
+							timeCh <- time.Now()
+							return nil, test_helpers.NewTestNetError(false, true)
+						}
+					}
+				})
+
+				It("retries connecting after a short delay", func() {
+					Eventually(fakeLocalListener.AcceptCallCount).Should(Equal(3))
+					Expect(timeCh).To(HaveLen(3))
+
+					times := make([]time.Time, 0)
+					for t := range timeCh {
+						times = append(times, t)
+					}
+
+					Expect(times[1]).To(BeTemporally("~", times[0].Add(100*time.Millisecond), 20*time.Millisecond))
+					Expect(times[2]).To(BeTemporally("~", times[1].Add(100*time.Millisecond), 20*time.Millisecond))
+				})
 			})
 		})
 
@@ -1113,6 +1143,35 @@ var _ = Describe("Diego SSH Plugin", func() {
 		})
 	})
 
+	Describe("Wait", func() {
+		var opts *options.SSHOptions
+
+		BeforeEach(func() {
+			opts = &options.SSHOptions{
+				AppName: "app-1",
+			}
+
+			fakeAppFactory.GetReturns(app.App{
+				State: "STARTED",
+				Diego: true,
+			}, nil)
+			fakeInfoFactory.GetReturns(info.Info{}, nil)
+			fakeCredFactory.GetReturns(credential.Credential{}, nil)
+		})
+
+		JustBeforeEach(func() {
+			connectErr := secureShell.Connect(opts)
+			Expect(connectErr).NotTo(HaveOccurred())
+		})
+
+		It("calls close on the secureClient", func() {
+			err := secureShell.Wait()
+			Expect(err).NotTo(HaveOccurred())
+
+			Expect(fakeSecureClient.WaitCallCount()).To(Equal(1))
+		})
+	})
+
 	Describe("Close", func() {
 		var opts *options.SSHOptions
 

cf-plugin/options/ssh.go

@@ -25,16 +25,18 @@ type ForwardSpec struct {
 }
 
 type SSHOptions struct {
-	AppName            string
-	Command            []string
-	Instance           uint
-	SkipHostValidation bool
-	TerminalRequest    TTYRequest
-	ForwardSpecs       []ForwardSpec
+	AppName             string
+	Command             []string
+	Instance            uint
+	SkipHostValidation  bool
+	SkipRemoteExecution bool
+	TerminalRequest     TTYRequest
+	ForwardSpecs        []ForwardSpec
 
 	getoptSet                       *getopt.Set
 	instanceOption                  getopt.Option
 	skipHostValidationOption        getopt.Option
+	skipRemoteExecutionOption       getopt.Option
 	disableTerminalAllocationOption getopt.Option
 	forceTerminalAllocationOption   getopt.Option
 	localForwardingOption           getopt.Option
@@ -57,6 +59,12 @@ func NewSSHOptions() *SSHOptions {
 		"skip host key validation",
 	).SetFlag()
 
+	sshOptions.skipRemoteExecutionOption = opts.BoolVar(
+		&sshOptions.SkipRemoteExecution,
+		'N',
+		"do not execute a remote command",
+	).SetFlag()
+
 	var force, disable bool
 	sshOptions.forceTerminalAllocationOption = opts.BoolVar(&force, 't', "force pseudo-tty allocation").SetFlag()
 	sshOptions.disableTerminalAllocationOption = opts.BoolVar(&disable, 'T', "disable pseudo-tty allocation").SetFlag()