Skip to content
Permalink
Browse files

Merge pull request #55 from antoniomika/fix_no_tty

Fixed connection close during a notty session
  • Loading branch information
antoniomika committed Nov 29, 2019
2 parents 9226761 + 6761b05 commit ff31f2bbf146788aa23629dbf72e2125c16584c6
Showing with 70 additions and 26 deletions.
  1. +10 −2 channels.go
  2. +5 −0 handle.go
  3. +14 −0 main.go
  4. +41 −24 requests.go
@@ -42,6 +42,8 @@ func handleSession(newChannel ssh.NewChannel, sshConn *SSHConnection, state *Sta
data := make([]byte, 4096)
dataRead, err := connection.Read(data)
if err != nil && err == io.EOF {
break
} else if err != nil {
select {
case <-sshConn.Close:
break
@@ -122,8 +124,14 @@ func handleAlias(newChannel ssh.NewChannel, sshConn *SSHConnection, state *State

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

copyBoth(conn, connection, false)
sshConn.CleanUp(state)
copyBoth(conn, connection)

select {
case <-sshConn.Close:
break
default:
sshConn.CleanUp(state)
}
}

func writeToSession(connection ssh.Channel, c string) {
@@ -22,6 +22,11 @@ func handleRequest(newRequest *ssh.Request, sshConn *SSHConnection, state *State
case "tcpip-forward":
go checkSession(newRequest, sshConn, state)
handleRemoteForward(newRequest, sshConn, state)
case "keepalive@openssh.com":
err := newRequest.Reply(true, nil)
if err != nil {
log.Println("Error replying to socket request:", err)
}
default:
err := newRequest.Reply(false, nil)
if err != nil {
14 main.go
@@ -267,6 +267,20 @@ func main() {

state.SSHConnections.Store(sshConn.RemoteAddr(), holderConn)

go func() {
err := sshConn.Wait()
if err != nil && *debug {
log.Println("Closing SSH connection:", err)
}

select {
case <-holderConn.Close:
break
default:
holderConn.CleanUp(state)
}
}()

go handleRequests(reqs, holderConn, state)
go handleChannels(chans, holderConn, state)

@@ -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 ff31f2b

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