diff --git a/ssh/handshake.go b/ssh/handshake.go index 2b84c35716..07a1843e0a 100644 --- a/ssh/handshake.go +++ b/ssh/handshake.go @@ -58,11 +58,13 @@ type handshakeTransport struct { incoming chan []byte readError error - mu sync.Mutex - writeError error - sentInitPacket []byte - sentInitMsg *kexInitMsg - pendingPackets [][]byte // Used when a key exchange is in progress. + mu sync.Mutex + writeError error + sentInitPacket []byte + sentInitMsg *kexInitMsg + pendingPackets [][]byte // Used when a key exchange is in progress. + writePacketsLeft uint32 + writeBytesLeft int64 // If the read loop wants to schedule a kex, it pings this // channel, and the write loop will send out a kex @@ -71,7 +73,8 @@ type handshakeTransport struct { // If the other side requests or confirms a kex, its kexInit // packet is sent here for the write loop to find it. - startKex chan *pendingKex + startKex chan *pendingKex + kexLoopDone chan struct{} // closed (with writeError non-nil) when kexLoop exits // data for host key checking hostKeyCallback HostKeyCallback @@ -86,12 +89,10 @@ type handshakeTransport struct { // Algorithms agreed in the last key exchange. algorithms *algorithms + // Counters exclusively owned by readLoop. readPacketsLeft uint32 readBytesLeft int64 - writePacketsLeft uint32 - writeBytesLeft int64 - // The session ID or nil if first kex did not complete yet. sessionID []byte } @@ -108,7 +109,8 @@ func newHandshakeTransport(conn keyingTransport, config *Config, clientVersion, clientVersion: clientVersion, incoming: make(chan []byte, chanSize), requestKex: make(chan struct{}, 1), - startKex: make(chan *pendingKex, 1), + startKex: make(chan *pendingKex), + kexLoopDone: make(chan struct{}), config: config, } @@ -340,16 +342,17 @@ write: t.mu.Unlock() } + // Unblock reader. + t.conn.Close() + // drain startKex channel. We don't service t.requestKex // because nobody does blocking sends there. - go func() { - for init := range t.startKex { - init.done <- t.writeError - } - }() + for request := range t.startKex { + request.done <- t.getWriteError() + } - // Unblock reader. - t.conn.Close() + // Mark that the loop is done so that Close can return. + close(t.kexLoopDone) } // The protocol uses uint32 for packet counters, so we can't let them @@ -545,7 +548,16 @@ func (t *handshakeTransport) writePacket(p []byte) error { } func (t *handshakeTransport) Close() error { - return t.conn.Close() + // Close the connection. This should cause the readLoop goroutine to wake up + // and close t.startKex, which will shut down kexLoop if running. + err := t.conn.Close() + + // Wait for the kexLoop goroutine to complete. + // At that point we know that the readLoop goroutine is complete too, + // because kexLoop itself waits for readLoop to close the startKex channel. + <-t.kexLoopDone + + return err } func (t *handshakeTransport) enterKeyExchange(otherInitPacket []byte) error {