Merge pull request #455 from Darkren/fix/ssh_data_race

Fix data race in SSH

internal/therealssh/channel.go

@@ -11,6 +11,7 @@ import (
 	"os/user"
 	"path/filepath"
 	"strings"
+	"sync"
 
 	"github.com/kr/pty"
 	"github.com/skycoin/dmsg/cipher"
@@ -34,12 +35,18 @@ type SSHChannel struct {
 
 	session  *Session
 	listener *net.UnixListener
+
+	dataChMx sync.Mutex
 	dataCh   chan []byte
+
+	doneOnce sync.Once
+	done     chan struct{}
 }
 
 // OpenChannel constructs new SSHChannel with empty Session.
 func OpenChannel(remoteID uint32, remoteAddr *app.Addr, conn net.Conn) *SSHChannel {
-	return &SSHChannel{RemoteID: remoteID, conn: conn, RemoteAddr: remoteAddr, msgCh: make(chan []byte), dataCh: make(chan []byte)}
+	return &SSHChannel{RemoteID: remoteID, conn: conn, RemoteAddr: remoteAddr, msgCh: make(chan []byte),
+		dataCh: make(chan []byte), done: make(chan struct{})}
 }
 
 // OpenClientChannel constructs new client SSHChannel with empty Session.
@@ -243,35 +250,66 @@ func (sshCh *SSHChannel) WindowChange(sz *pty.Winsize) error {
 	return sshCh.session.WindowChange(sz)
 }
 
-// Close safely closes Channel resources.
-func (sshCh *SSHChannel) Close() error {
-	select {
-	case <-sshCh.dataCh:
-	default:
-		close(sshCh.dataCh)
-	}
-	close(sshCh.msgCh)
+func (sshCh *SSHChannel) close() (closed bool, err error) {
+	sshCh.doneOnce.Do(func() {
+		closed = true
 
-	var sErr, lErr error
-	if sshCh.session != nil {
-		sErr = sshCh.session.Close()
-	}
+		close(sshCh.done)
 
-	if sshCh.listener != nil {
-		lErr = sshCh.listener.Close()
-	}
+		select {
+		case <-sshCh.dataCh:
+		default:
+			sshCh.dataChMx.Lock()
+			close(sshCh.dataCh)
+			sshCh.dataChMx.Unlock()
+		}
+		close(sshCh.msgCh)
 
-	if sErr != nil {
-		return sErr
-	}
+		var sErr, lErr error
+		if sshCh.session != nil {
+			sErr = sshCh.session.Close()
+		}
+
+		if sshCh.listener != nil {
+			lErr = sshCh.listener.Close()
+		}
 
-	if lErr != nil {
-		return lErr
+		if sErr != nil {
+			err = sErr
+			return
+		}
+
+		if lErr != nil {
+			err = lErr
+		}
+	})
+
+	return closed, err
+}
+
+// Close safely closes Channel resources.
+func (sshCh *SSHChannel) Close() error {
+	closed, err := sshCh.close()
+	if err != nil {
+		return err
+	}
+	if !closed {
+		return errors.New("channel is already closed")
 	}
 
 	return nil
 }
 
+// IsClosed returns whether the Channel is closed.
+func (sshCh *SSHChannel) IsClosed() bool {
+	select {
+	case <-sshCh.done:
+		return true
+	default:
+		return false
+	}
+}
+
 func debug(format string, v ...interface{}) {
 	if !Debug {
 		return

internal/therealssh/client.go

@@ -132,7 +132,11 @@ func (c *Client) serveConn(conn net.Conn) error {
 		case CmdChannelOpenResponse, CmdChannelResponse:
 			sshCh.msgCh <- data
 		case CmdChannelData:
-			sshCh.dataCh <- data
+			sshCh.dataChMx.Lock()
+			if !sshCh.IsClosed() {
+				sshCh.dataCh <- data
+			}
+			sshCh.dataChMx.Unlock()
 		case CmdChannelServerClose:
 			err = sshCh.Close()
 		default:

internal/therealssh/server.go

@@ -126,7 +126,11 @@ func (s *Server) HandleData(remotePK cipher.PubKey, localID uint32, data []byte)
 		return errors.New("session is not started")
 	}
 
-	channel.dataCh <- data
+	channel.dataChMx.Lock()
+	if !channel.IsClosed() {
+		channel.dataCh <- data
+	}
+	channel.dataChMx.Unlock()
 	return nil
 }