func (handler *DirectTcpipChannelHandler) HandleNewChannel(logger lager.Logger, newChannel ssh.NewChannel) {
type channelOpenDirectTcpipMsg struct {
TargetAddr string
TargetPort uint32
OriginAddr string
OriginPort uint32
}
var directTcpipMessage channelOpenDirectTcpipMsg
err := ssh.Unmarshal(newChannel.ExtraData(), &directTcpipMessage)
if err != nil {
newChannel.Reject(ssh.ConnectionFailed, "Failed to parse open channel message")
return
}
destination := fmt.Sprintf("%s:%d", directTcpipMessage.TargetAddr, directTcpipMessage.TargetPort)
conn, err := handler.dialer.Dial("tcp", destination)
if err != nil {
newChannel.Reject(ssh.ConnectionFailed, err.Error())
return
}
channel, requests, err := newChannel.Accept()
go ssh.DiscardRequests(requests)
wg := &sync.WaitGroup{}
wg.Add(2)
go helpers.CopyAndClose(logger.Session("to-target"), wg, conn, channel)
go helpers.CopyAndClose(logger.Session("to-channel"), wg, channel, conn)
wg.Wait()
}
func (sess *session) run(command *exec.Cmd) error {
logger := sess.logger.Session("run")
command.Stdout = sess.channel
command.Stderr = sess.channel.Stderr()
stdin, err := command.StdinPipe()
if err != nil {
return err
}
go helpers.CopyAndClose(logger.Session("to-stdin"), nil, stdin, sess.channel)
return sess.runner.Start(command)
}
})
})
Describe("CopyAndClose", func() {
var reader io.Reader
var fakeWriteCloser *fake_io.FakeWriteCloser
var wg *sync.WaitGroup
BeforeEach(func() {
reader = strings.NewReader("message")
fakeWriteCloser = &fake_io.FakeWriteCloser{}
wg = nil
})
JustBeforeEach(func() {
helpers.CopyAndClose(logger, wg, fakeWriteCloser, reader)
})
It("copies from source to target", func() {
Expect(fakeWriteCloser.WriteCallCount()).To(Equal(1))
Expect(string(fakeWriteCloser.WriteArgsForCall(0))).To(Equal("message"))
})
It("closes the target when the copy is complete", func() {
Expect(fakeWriteCloser.CloseCallCount()).To(Equal(1))
})
Context("when a wait group is provided", func() {
BeforeEach(func() {
wg = &sync.WaitGroup{}
wg.Add(1)