From 409a35189b23ca1490265c4eeb13ea370d3144ba Mon Sep 17 00:00:00 2001
From: Jianmin <jianmin@iscas.ac.cn>
Date: Mon, 23 Dec 2024 17:47:26 +0800
Subject: [PATCH] [backport] ssh: implement strict KEX protocol changes

---
 vendor/golang.org/x/crypto/ssh/handshake.go | 59 +++++++++++++++++++--
 vendor/golang.org/x/crypto/ssh/transport.go | 32 +++++++++--
 2 files changed, 81 insertions(+), 10 deletions(-)

diff --git a/vendor/golang.org/x/crypto/ssh/handshake.go b/vendor/golang.org/x/crypto/ssh/handshake.go
index 653dc4d..e7d4545 100644
--- a/vendor/golang.org/x/crypto/ssh/handshake.go
+++ b/vendor/golang.org/x/crypto/ssh/handshake.go
@@ -34,6 +34,16 @@ type keyingTransport interface {
 	// direction will be effected if a msgNewKeys message is sent
 	// or received.
 	prepareKeyChange(*algorithms, *kexResult) error
+
+	// setStrictMode sets the strict KEX mode, notably triggering
+	// sequence number resets on sending or receiving msgNewKeys.
+	// If the sequence number is already > 1 when setStrictMode
+	// is called, an error is returned.
+	setStrictMode() error
+
+	// setInitialKEXDone indicates to the transport that the initial key exchange
+	// was completed
+	setInitialKEXDone()
 }
 
 // handshakeTransport implements rekeying on top of a keyingTransport
@@ -94,6 +104,10 @@ type handshakeTransport struct {
 
 	// The session ID or nil if first kex did not complete yet.
 	sessionID []byte
+
+	// strictMode indicates if the other side of the handshake indicated
+	// that we should be following the strict KEX protocol restrictions.
+	strictMode bool
 }
 
 type pendingKex struct {
@@ -201,7 +215,10 @@ func (t *handshakeTransport) readLoop() {
 			close(t.incoming)
 			break
 		}
-		if p[0] == msgIgnore || p[0] == msgDebug {
+		// If this is the first kex, and strict KEX mode is enabled,
+		// we don't ignore any messages, as they may be used to manipulate
+		// the packet sequence numbers.
+		if !(t.sessionID == nil && t.strictMode) && (p[0] == msgIgnore || p[0] == msgDebug) {
 			continue
 		}
 		t.incoming <- p
@@ -432,6 +449,11 @@ func (t *handshakeTransport) readOnePacket(first bool) ([]byte, error) {
 	return successPacket, nil
 }
 
+const (
+	kexStrictClient = "kex-strict-c-v00@openssh.com"
+	kexStrictServer = "kex-strict-s-v00@openssh.com"
+)
+
 // sendKexInit sends a key change message.
 func (t *handshakeTransport) sendKexInit() error {
 	t.mu.Lock()
@@ -445,7 +467,6 @@ func (t *handshakeTransport) sendKexInit() error {
 	}
 
 	msg := &kexInitMsg{
-		KexAlgos:                t.config.KeyExchanges,
 		CiphersClientServer:     t.config.Ciphers,
 		CiphersServerClient:     t.config.Ciphers,
 		MACsClientServer:        t.config.MACs,
@@ -455,6 +476,13 @@ func (t *handshakeTransport) sendKexInit() error {
 	}
 	io.ReadFull(rand.Reader, msg.Cookie[:])
 
+	// We mutate the KexAlgos slice, in order to add the kex-strict extension algorithm,
+	// and possibly to add the ext-info extension algorithm. Since the slice may be the
+	// user owned KeyExchanges, we create our own slice in order to avoid using user
+	// owned memory by mistake.
+	msg.KexAlgos = make([]string, 0, len(t.config.KeyExchanges)+2) // room for kex-strict and ext-info
+	msg.KexAlgos = append(msg.KexAlgos, t.config.KeyExchanges...)
+
 	isServer := len(t.hostKeys) > 0
 	if isServer {
 		for _, k := range t.hostKeys {
@@ -474,17 +502,24 @@ func (t *handshakeTransport) sendKexInit() error {
 				msg.ServerHostKeyAlgos = append(msg.ServerHostKeyAlgos, keyFormat)
 			}
 		}
+
+		if t.sessionID == nil {
+			msg.KexAlgos = append(msg.KexAlgos, kexStrictServer)
+		}
 	} else {
 		msg.ServerHostKeyAlgos = t.hostKeyAlgorithms
 
 		// As a client we opt in to receiving SSH_MSG_EXT_INFO so we know what
 		// algorithms the server supports for public key authentication. See RFC
 		// 8308, Section 2.1.
+		//
+		// We also send the strict KEX mode extension algorithm, in order to opt
+		// into the strict KEX mode.
 		if firstKeyExchange := t.sessionID == nil; firstKeyExchange {
-			msg.KexAlgos = make([]string, 0, len(t.config.KeyExchanges)+1)
-			msg.KexAlgos = append(msg.KexAlgos, t.config.KeyExchanges...)
 			msg.KexAlgos = append(msg.KexAlgos, "ext-info-c")
+			msg.KexAlgos = append(msg.KexAlgos, kexStrictClient)
 		}
+
 	}
 
 	packet := Marshal(msg)
@@ -581,6 +616,13 @@ func (t *handshakeTransport) enterKeyExchange(otherInitPacket []byte) error {
 		return err
 	}
 
+	if t.sessionID == nil && ((isClient && contains(serverInit.KexAlgos, kexStrictServer)) || (!isClient && contains(clientInit.KexAlgos, kexStrictClient))) {
+		t.strictMode = true
+		if err := t.conn.setStrictMode(); err != nil {
+			return err
+		}
+	}
+
 	// We don't send FirstKexFollows, but we handle receiving it.
 	//
 	// RFC 4253 section 7 defines the kex and the agreement method for
@@ -615,7 +657,8 @@ func (t *handshakeTransport) enterKeyExchange(otherInitPacket []byte) error {
 		return err
 	}
 
-	if t.sessionID == nil {
+	firstKeyExchange := t.sessionID == nil
+	if firstKeyExchange {
 		t.sessionID = result.H
 	}
 	result.SessionID = t.sessionID
@@ -632,6 +675,12 @@ func (t *handshakeTransport) enterKeyExchange(otherInitPacket []byte) error {
 		return unexpectedMessageError(msgNewKeys, packet[0])
 	}
 
+	if firstKeyExchange {
+		// Indicates to the transport that the first key exchange is completed
+		// after receiving SSH_MSG_NEWKEYS.
+		t.conn.setInitialKEXDone()
+	}
+
 	return nil
 }
 
diff --git a/vendor/golang.org/x/crypto/ssh/transport.go b/vendor/golang.org/x/crypto/ssh/transport.go
index acf5a21..4df45fc 100644
--- a/vendor/golang.org/x/crypto/ssh/transport.go
+++ b/vendor/golang.org/x/crypto/ssh/transport.go
@@ -48,6 +48,9 @@ type transport struct {
 	rand      io.Reader
 	isClient  bool
 	io.Closer
+
+	strictMode     bool
+	initialKEXDone bool
 }
 
 // packetCipher represents a combination of SSH encryption/MAC
@@ -73,6 +76,18 @@ type connectionState struct {
 	pendingKeyChange chan packetCipher
 }
 
+func (t *transport) setStrictMode() error {
+	if t.reader.seqNum != 1 {
+		return errors.New("ssh: sequence number != 1 when strict KEX mode requested")
+	}
+	t.strictMode = true
+	return nil
+}
+
+func (t *transport) setInitialKEXDone() {
+	t.initialKEXDone = true
+}
+
 // prepareKeyChange sets up key material for a keychange. The key changes in
 // both directions are triggered by reading and writing a msgNewKey packet
 // respectively.
@@ -111,11 +126,12 @@ func (t *transport) printPacket(p []byte, write bool) {
 // Read and decrypt next packet.
 func (t *transport) readPacket() (p []byte, err error) {
 	for {
-		p, err = t.reader.readPacket(t.bufReader)
+		p, err = t.reader.readPacket(t.bufReader, t.strictMode)
 		if err != nil {
 			break
 		}
-		if len(p) == 0 || (p[0] != msgIgnore && p[0] != msgDebug) {
+		// in strict mode we pass through DEBUG and IGNORE packets only during the initial KEX
+		if len(p) == 0 || (t.strictMode && !t.initialKEXDone) || (p[0] != msgIgnore && p[0] != msgDebug) {
 			break
 		}
 	}
@@ -126,7 +142,7 @@ func (t *transport) readPacket() (p []byte, err error) {
 	return p, err
 }
 
-func (s *connectionState) readPacket(r *bufio.Reader) ([]byte, error) {
+func (s *connectionState) readPacket(r *bufio.Reader, strictMode bool) ([]byte, error) {
 	packet, err := s.packetCipher.readCipherPacket(s.seqNum, r)
 	s.seqNum++
 	if err == nil && len(packet) == 0 {
@@ -139,6 +155,9 @@ func (s *connectionState) readPacket(r *bufio.Reader) ([]byte, error) {
 			select {
 			case cipher := <-s.pendingKeyChange:
 				s.packetCipher = cipher
+				if strictMode {
+					s.seqNum = 0
+				}
 			default:
 				return nil, errors.New("ssh: got bogus newkeys message")
 			}
@@ -169,10 +188,10 @@ func (t *transport) writePacket(packet []byte) error {
 	if debugTransport {
 		t.printPacket(packet, true)
 	}
-	return t.writer.writePacket(t.bufWriter, t.rand, packet)
+	return t.writer.writePacket(t.bufWriter, t.rand, packet, t.strictMode)
 }
 
-func (s *connectionState) writePacket(w *bufio.Writer, rand io.Reader, packet []byte) error {
+func (s *connectionState) writePacket(w *bufio.Writer, rand io.Reader, packet []byte, strictMode bool) error {
 	changeKeys := len(packet) > 0 && packet[0] == msgNewKeys
 
 	err := s.packetCipher.writeCipherPacket(s.seqNum, w, rand, packet)
@@ -187,6 +206,9 @@ func (s *connectionState) writePacket(w *bufio.Writer, rand io.Reader, packet []
 		select {
 		case cipher := <-s.pendingKeyChange:
 			s.packetCipher = cipher
+			if strictMode {
+				s.seqNum = 0
+			}
 		default:
 			panic("ssh: no key material for msgNewKeys")
 		}
-- 
2.39.5 (Apple Git-154)

