Skip to content
Permalink
Browse files

Fixed leaky goroutines

  • Loading branch information
antoniomika committed Nov 29, 2019
1 parent 5e19adc commit 6761b0593535c4dbd8066c16fea09f88e98911bf
Showing with 42 additions and 25 deletions.
  1. +1 −1 channels.go
  2. +41 −24 requests.go
@@ -124,7 +124,7 @@ func handleAlias(newChannel ssh.NewChannel, sshConn *SSHConnection, state *State

sshConn.Listeners.Store(conn.RemoteAddr(), nil)

copyBoth(conn, connection, false)
copyBoth(conn, connection)

select {
case <-sshConn.Close:
@@ -8,7 +8,6 @@ import (
"net"
"os"
"strconv"
"sync"
"time"

"github.com/logrusorgru/aurora"
@@ -227,47 +226,65 @@ func handleRemoteForward(newRequest *ssh.Request, sshConn *SSHConnection, state
}
}

go copyBoth(cl, newChan, false)
go copyBoth(cl, newChan)
go ssh.DiscardRequests(newReqs)
}
}

func copyBoth(writer net.Conn, reader ssh.Channel, wait bool) {
// IdleTimeoutConn handles the connection with a context deadline
// code adapted from https://qiita.com/kwi/items/b38d6273624ad3f6ae79
type IdleTimeoutConn struct {
Conn net.Conn
}

// Read is needed to implement the reader part
func (i IdleTimeoutConn) Read(buf []byte) (int, error) {
err := i.Conn.SetDeadline(time.Now().Add(5 * time.Second))
if err != nil {
return 0, err
}

return i.Conn.Read(buf)
}

// Write is needed to implement the writer part
func (i IdleTimeoutConn) Write(buf []byte) (int, error) {
err := i.Conn.SetDeadline(time.Now().Add(5 * time.Second))
if err != nil {
return 0, err
}

return i.Conn.Write(buf)
}

func copyBoth(writer net.Conn, reader ssh.Channel) {
closeBoth := func() {
time.Sleep(100 * time.Millisecond)
writer.Close()
reader.Close()
writer.Close()
}

var wg sync.WaitGroup

go func() {
if wait {
wg.Add(1)
defer wg.Done()
}
tcon := IdleTimeoutConn{
Conn: writer,
}

_, err := io.Copy(reader, writer)
copyToReader := func() {
_, err := io.Copy(reader, tcon)
if err != nil && *debug {
log.Println("Error copying to reader:", err)
}
}()

func() {
if wait {
wg.Add(1)
defer wg.Done()
}
closeBoth()
}

_, err := io.Copy(writer, reader)
copyToWriter := func() {
_, err := io.Copy(tcon, reader)
if err != nil && *debug {
log.Println("Error copying to writer:", err)
}
}()

if wait {
wg.Wait()
closeBoth()
}

closeBoth()
go copyToReader()
copyToWriter()
}

0 comments on commit 6761b05

Please sign in to comment.
You can’t perform that action at this time.