some cleanup

[#155484737]

cmd/sshd/main_test.go

@@ -716,7 +716,10 @@ var _ = Describe("SSH daemon", func() {
 		})
 
 		Context("when a client requests a remote port forward", func() {
-			var server *ghttp.Server
+			var (
+				server *ghttp.Server
+				ln     net.Listener
+			)
 
 			BeforeEach(func() {
 				server = ghttp.NewServer()
@@ -725,10 +728,13 @@ var _ = Describe("SSH daemon", func() {
 				)
 			})
 
-			It("forwards the remote port from server side to the target", func() {
-				ln, err := client.Listen("tcp", "127.0.0.1:0")
+			JustBeforeEach(func() {
+				var err error
+				ln, err = client.Listen("tcp", "127.0.0.1:0")
 				Expect(err).NotTo(HaveOccurred())
+			})
 
+			It("forwards the remote port from server side to the target", func() {
 				go func() {
 					for {
 						conn, err := ln.Accept()
@@ -767,6 +773,30 @@ var _ = Describe("SSH daemon", func() {
 				Expect(err).NotTo(HaveOccurred())
 				Expect(line).To(ContainSubstring("hello from the other side"))
 			})
+
+			Context("when the connection is closed", func() {
+				JustBeforeEach(func() {
+					Expect(client.Close()).To(Succeed())
+				})
+
+				It("closes the listeners associated with this conn", func() {
+					Eventually(func() error {
+						_, err := http.Get(fmt.Sprintf("http://%s", ln.Addr()))
+						return err
+					}).Should(MatchError(ContainSubstring("connection refused")))
+				})
+			})
+
+			Context("when the listener is closed", func() {
+				JustBeforeEach(func() {
+					Expect(ln.Close()).To(Succeed())
+				})
+
+				It("responds with a client refused error to clients", func() {
+					_, err := http.Get(fmt.Sprintf("http://%s", ln.Addr()))
+					Expect(err).To(MatchError(ContainSubstring("connection refused")))
+				})
+			})
 		})
 	})
 })

handlers/direct_tcpip_channel_handler_test.go

@@ -37,6 +37,8 @@ var _ = Describe("DirectTcpipChannelHandler", func() {
 		echoHandler *fake_server.FakeConnectionHandler
 		echoServer  *server.Server
 		echoAddress string
+
+		handleConnFinished chan struct{}
 	)
 
 	BeforeEach(func() {
@@ -76,14 +78,20 @@ var _ = Describe("DirectTcpipChannelHandler", func() {
 		serverNetConn, clientNetConn := test_helpers.Pipe()
 
 		sshd = daemon.New(logger, serverSSHConfig, nil, newChannelHandlers)
-		go sshd.HandleConnection(serverNetConn)
+
+		handleConnFinished = make(chan struct{})
+		go func() {
+			sshd.HandleConnection(serverNetConn)
+			close(handleConnFinished)
+		}()
 
 		client = test_helpers.NewClient(clientNetConn, nil)
 	})
 
 	AfterEach(func() {
 		client.Close()
 		echoServer.Shutdown()
+		Eventually(handleConnFinished).Should(BeClosed())
 	})
 
 	Context("when a session is opened", func() {

handlers/globalrequest/cancel_tcpip_forward_handler.go

@@ -4,6 +4,7 @@ import (
 	"net"
 	"strconv"
 
+	"code.cloudfoundry.org/diego-ssh/handlers/globalrequest/internal"
 	"code.cloudfoundry.org/diego-ssh/helpers"
 	"code.cloudfoundry.org/lager"
 	"golang.org/x/crypto/ssh"
@@ -21,7 +22,7 @@ func (h *CancelTCPIPForwardHandler) HandleRequest(logger lager.Logger, request *
 	logger.Info("start")
 	defer logger.Info("done")
 
-	var tcpipForwardMessage tcpipForwardMsg
+	var tcpipForwardMessage internal.TCPIPForwardRequest
 
 	err := ssh.Unmarshal(request.Payload, &tcpipForwardMessage)
 	if err != nil {

handlers/globalrequest/cancel_tcpip_forward_handler_test.go

@@ -0,0 +1,108 @@
+package globalrequest_test
+
+import (
+	"fmt"
+	"net"
+
+	"code.cloudfoundry.org/diego-ssh/daemon"
+	"code.cloudfoundry.org/diego-ssh/handlers"
+	"code.cloudfoundry.org/diego-ssh/handlers/globalrequest"
+	"code.cloudfoundry.org/diego-ssh/handlers/globalrequest/internal"
+	"code.cloudfoundry.org/diego-ssh/test_helpers"
+	"code.cloudfoundry.org/lager/lagertest"
+	"code.cloudfoundry.org/localip"
+	"golang.org/x/crypto/ssh"
+
+	. "github.com/onsi/ginkgo"
+	. "github.com/onsi/gomega"
+)
+
+var _ = Describe("CancelTcpipForwardHandler", func() {
+	var (
+		remoteAddress string
+		sshClient     *ssh.Client
+		logger        *lagertest.TestLogger
+	)
+
+	BeforeEach(func() {
+		logger = lagertest.NewTestLogger("tcpip-forward-test")
+
+		remotePort, err := localip.LocalPort()
+		Expect(err).NotTo(HaveOccurred())
+		remoteAddress = fmt.Sprintf("127.0.0.1:%d", remotePort)
+
+		globalRequestHandlers := map[string]handlers.GlobalRequestHandler{
+			globalrequest.TCPIPForward:       new(globalrequest.TCPIPForwardHandler),
+			globalrequest.CancelTCPIPForward: new(globalrequest.CancelTCPIPForwardHandler),
+		}
+
+		serverSSHConfig := &ssh.ServerConfig{
+			NoClientAuth: true,
+		}
+		serverSSHConfig.AddHostKey(TestHostKey)
+
+		sshd := daemon.New(logger, serverSSHConfig, globalRequestHandlers, nil)
+
+		serverNetConn, clientNetConn := test_helpers.Pipe()
+		go sshd.HandleConnection(serverNetConn)
+		sshClient = test_helpers.NewClient(clientNetConn, nil)
+	})
+
+	Context("when the request is invalid", func() {
+		It("rejects the request", func() {
+			payload := ssh.Marshal(struct {
+				port uint16
+			}{
+				port: 10,
+			})
+			ok, _, err := sshClient.SendRequest("cancel-tcpip-forward", true, payload)
+			Expect(err).NotTo(HaveOccurred())
+			Expect(ok).NotTo(BeTrue())
+		})
+	})
+
+	Context("when the listener isn't found", func() {
+		It("rejects the request", func() {
+			payload := ssh.Marshal(internal.TCPIPForwardRequest{
+				Address: "127.0.0.1",
+				Port:    9090,
+			})
+			ok, _, err := sshClient.SendRequest("cancel-tcpip-forward", true, payload)
+			Expect(err).NotTo(HaveOccurred())
+			Expect(ok).NotTo(BeTrue())
+		})
+	})
+
+	Context("when the listener exists", func() {
+		var (
+			ok  bool
+			err error
+		)
+
+		BeforeEach(func() {
+			addr, err := net.ResolveTCPAddr("tcp", "127.0.0.1:9090")
+			Expect(err).NotTo(HaveOccurred())
+			_, err = sshClient.ListenTCP(addr)
+			Expect(err).NotTo(HaveOccurred())
+
+			_, err = net.Dial("tcp", "127.0.0.1:9090")
+			Expect(err).NotTo(HaveOccurred())
+
+			payload := ssh.Marshal(internal.TCPIPForwardRequest{
+				Address: "127.0.0.1",
+				Port:    9090,
+			})
+			ok, _, err = sshClient.SendRequest("cancel-tcpip-forward", true, payload)
+		})
+
+		It("successfully process the request", func() {
+			Expect(err).NotTo(HaveOccurred())
+			Expect(ok).To(BeTrue())
+		})
+
+		It("stops listening to the port", func() {
+			_, err := net.Dial("tcp", "127.0.0.1:9090")
+			Expect(err).To(MatchError(ContainSubstring("connection refused")))
+		})
+	})
+})

handlers/globalrequest/internal/package.go

@@ -0,0 +1 @@
+package internal // import "code.cloudfoundry.org/diego-ssh/handlers/globalrequest/internal"

handlers/globalrequest/internal/tcpip_forward_messages.go

@@ -0,0 +1,10 @@
+package internal
+
+type TCPIPForwardRequest struct {
+	Address string
+	Port    uint32
+}
+
+type TCPIPForwardResponse struct {
+	Port uint32
+}

handlers/globalrequest/tcpip_forward_handler.go

@@ -5,6 +5,7 @@ import (
 	"strconv"
 	"sync"
 
+	"code.cloudfoundry.org/diego-ssh/handlers/globalrequest/internal"
 	"code.cloudfoundry.org/diego-ssh/helpers"
 	"code.cloudfoundry.org/lager"
 	"golang.org/x/crypto/ssh"
@@ -14,11 +15,6 @@ const TCPIPForward = "tcpip-forward"
 
 type TCPIPForwardHandler struct{}
 
-type tcpipForwardMsg struct {
-	Address string
-	Port    uint32
-}
-
 func (h *TCPIPForwardHandler) HandleRequest(logger lager.Logger, request *ssh.Request, conn ssh.Conn, lnStore *helpers.ListenerStore) {
 	logger = logger.Session("tcpip-forward", lager.Data{
 		"type":       request.Type,
@@ -27,8 +23,7 @@ func (h *TCPIPForwardHandler) HandleRequest(logger lager.Logger, request *ssh.Re
 	logger.Info("start")
 	defer logger.Info("done")
 
-	var tcpipForwardMessage tcpipForwardMsg
-
+	var tcpipForwardMessage internal.TCPIPForwardRequest
 	err := ssh.Unmarshal(request.Payload, &tcpipForwardMessage)
 	if err != nil {
 		logger.Error("unmarshal-failed", err)
@@ -50,11 +45,10 @@ func (h *TCPIPForwardHandler) HandleRequest(logger lager.Logger, request *ssh.Re
 		return
 	}
 
-	lnStore.AddListener(address, listener)
-
 	var listenerAddr string
 	var listenerPort uint32
 	if addr, ok := listener.Addr().(*net.TCPAddr); ok {
+		address = addr.String()
 		listenerAddr = addr.IP.String()
 		listenerPort = uint32(addr.Port)
 	}
@@ -63,21 +57,20 @@ func (h *TCPIPForwardHandler) HandleRequest(logger lager.Logger, request *ssh.Re
 		"port": listenerPort,
 	})
 
+	lnStore.AddListener(address, listener)
+
 	go h.forwardAcceptLoop(listener, logger, conn, tcpipForwardMessage.Address, listenerPort)
 
-	var tcpipForwardResponse struct {
-		Port uint32
-	}
+	var tcpipForwardResponse internal.TCPIPForwardResponse
 	tcpipForwardResponse.Port = listenerPort
 
 	var replyPayload []byte
+
 	if tcpipForwardMessage.Port == 0 {
+		// See RFC 4254, section 7.1
 		replyPayload = ssh.Marshal(tcpipForwardResponse)
 	}
 
-	// WARN (CEV): I think we only want to reply if WantReply
-	// is true and this appears to always reply!
-	//
 	// Reply() will only send something when WantReply is true
 	_ = request.Reply(true, replyPayload)
 }
@@ -90,9 +83,7 @@ type forwardedTCPPayload struct {
 	OriginPort uint32
 }
 
-func (h *TCPIPForwardHandler) forwardAcceptLoop(listener net.Listener, logger lager.Logger, sshConn ssh.Conn,
-	lnAddr string, lnPort uint32) {
-
+func (h *TCPIPForwardHandler) forwardAcceptLoop(listener net.Listener, logger lager.Logger, sshConn ssh.Conn, lnAddr string, lnPort uint32) {
 	logger = logger.Session("forward-accept-loop")
 	logger.Info("start")
 	defer logger.Info("done")
@@ -130,11 +121,12 @@ func (h *TCPIPForwardHandler) forwardAcceptLoop(listener net.Listener, logger la
 			var wg sync.WaitGroup
 			wg.Add(2)
 
-			go helpers.CopyAndClose(logger.Session("to-target"), &wg, conn, channel,
-				func() { conn.(*net.TCPConn).CloseWrite() })
-
-			go helpers.CopyAndClose(logger.Session("to-channel"), &wg, channel, conn,
-				func() { channel.CloseWrite() })
+			go helpers.CopyAndClose(logger.Session("to-target"), &wg, conn, channel, func() {
+				conn.Close()
+			})
+			go helpers.CopyAndClose(logger.Session("to-channel"), &wg, channel, conn, func() {
+				channel.CloseWrite()
+			})
 
 			wg.Wait()
 		}(conn)