From 6c79c218f0abcbc9a637ca9b3544ea909b9294c4 Mon Sep 17 00:00:00 2001 From: Qining Lu Date: Thu, 6 Sep 2018 22:17:22 -0400 Subject: [PATCH] Close ssh tunnel when either end closes the connection So that when users tracing through SSH with a specific number of frames to capture, the connection will be closed after that number of frames are captured. --- core/os/device/remotessh/commands.go | 27 +++++++++++++++++++++++++-- 1 file changed, 25 insertions(+), 2 deletions(-) diff --git a/core/os/device/remotessh/commands.go b/core/os/device/remotessh/commands.go index 6467808502..e718aceddc 100644 --- a/core/os/device/remotessh/commands.go +++ b/core/os/device/remotessh/commands.go @@ -208,8 +208,31 @@ func (b binding) doTunnel(ctx context.Context, local net.Conn, remotePort int) e wg := sync.WaitGroup{} - copy := func(writer io.Writer, reader io.Reader) { - _, err := io.Copy(writer, reader) + copy := func(writer net.Conn, reader net.Conn) { + // Use the same buffer size used in io.Copy + buf := make([]byte, 32*1024) + var err error + for { + nr, er := reader.Read(buf) + if nr > 0 { + nw, ew := writer.Write(buf[0:nr]) + if ew != nil { + err = ew + break + } + if nr != nw { + err = fmt.Errorf("short write") + break + } + } + if er != nil { + if er != io.EOF { + err = er + } + break + } + } + writer.Close() if err != nil { log.E(ctx, "Copy Error %s", err) }