ssh: improve DH GEX group selection using PreferredBits

Previously, the server selected the Diffie-Hellman group based solely on
the MaxBits value provided by the client. This resulted in suboptimal
modulus selection, often ignoring the client's PreferredBits or selecting
a larger-than-necessary group.

This change implements a "best fit" selection algorithm similar to
OpenSSH's choose_dh logic. It attempts to find the smallest available
group larger than or equal to the client's PreferredBits, falling back to
the largest available group within the accepted range if no group above
the preference is available.

Additionally, this commit caches the parsed Oakley groups using
sync.OnceValue, avoiding repeated big.Int parsing on every handshake
while keeping the cost out of package initialization.

This issue was found during a security audit by NCC Group Cryptography
Services, sponsored by Teleport, and was assessed and is being fixed as
a non-security bug.

Change-Id: Idfa81bbcf354a7fb7b541cb4bbeb6e4a0181398a
Reviewed-on: https://go-review.googlesource.com/c/crypto/+/782424
Auto-Submit: Nicola Murino <nicola.murino@gmail.com>
Reviewed-by: Filippo Valsorda <filippo@golang.org>
Reviewed-by: Mark Freeman <markfreeman@google.com>
LUCI-TryBot-Result: golang-scoped@luci-project-accounts.iam.gserviceaccount.com <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: David Chase <drchase@google.com>

Author: drakkan

ssh/kex.go

@@ -16,6 +16,7 @@ import (
 	"io"
 	"math/big"
 	"slices"
+	"sync"
 
 	"golang.org/x/crypto/curve25519"
 )
@@ -718,15 +719,9 @@ func (gex *dhGEXSHA) Server(c packetConn, randSource io.Reader, magics *handshak
 			kexDHGexRequest.MaxBits, kexDHGexRequest.PreferredBits)
 	}
 
-	var p *big.Int
-	// We hardcode sending Oakley Group 14 (2048 bits), Oakley Group 15 (3072
-	// bits) or Oakley Group 16 (4096 bits), based on the requested max size.
-	if kexDHGexRequest.MaxBits < 3072 {
-		p, _ = new(big.Int).SetString(oakleyGroup14, 16)
-	} else if kexDHGexRequest.MaxBits < 4096 {
-		p, _ = new(big.Int).SetString(oakleyGroup15, 16)
-	} else {
-		p, _ = new(big.Int).SetString(oakleyGroup16, 16)
+	p, err := chooseDH(kexDHGexRequest)
+	if err != nil {
+		return nil, err
 	}
 
 	g := big.NewInt(2)
@@ -805,3 +800,65 @@ func (gex *dhGEXSHA) Server(c packetConn, randSource io.Reader, magics *handshak
 		Hash:      gex.hashFunc,
 	}, err
 }
+
+type dhKEXGroup struct {
+	size int
+	p    *big.Int
+}
+
+// supportedDHKEXGroups returns the DH groups the server is willing to offer
+// for diffie-hellman-group-exchange-* key exchanges. The list is built lazily
+// on first use to keep the hex-to-big.Int parse out of package initialization.
+var supportedDHKEXGroups = sync.OnceValue(func() []dhKEXGroup {
+	specs := []struct {
+		size int
+		hex  string
+	}{
+		{2048, oakleyGroup14},
+		{3072, oakleyGroup15},
+		{4096, oakleyGroup16},
+	}
+	out := make([]dhKEXGroup, 0, len(specs))
+	for _, s := range specs {
+		p, _ := new(big.Int).SetString(s.hex, 16)
+		out = append(out, dhKEXGroup{size: s.size, p: p})
+	}
+	return out
+})
+
+// chooseDH picks a DH group for the given client request, mirroring the
+// algorithm used by OpenSSH's choose_dh in dh.c: prefer the smallest known
+// group larger than  or equal to the client's PreferredBits, and otherwise pick
+// the largest group within the accepted [MinBits, MaxBits] range.
+func chooseDH(req kexDHGexRequestMsg) (*big.Int, error) {
+	var best *big.Int
+	bestSize := 0
+	wantBits := int(req.PreferredBits)
+
+	for _, group := range supportedDHKEXGroups() {
+		if uint32(group.size) < req.MinBits || uint32(group.size) > req.MaxBits {
+			continue
+		}
+
+		if bestSize == 0 {
+			best = group.p
+			bestSize = group.size
+			continue
+		}
+
+		closerFromAbove := group.size >= wantBits && group.size < bestSize
+		closerFromBelow := group.size > bestSize && bestSize < wantBits
+
+		if closerFromAbove || closerFromBelow {
+			best = group.p
+			bestSize = group.size
+		}
+	}
+
+	if bestSize == 0 {
+		return nil, fmt.Errorf("ssh: no suitable DH group found for request min: %d, preferred: %d, max: %d",
+			req.MinBits, req.PreferredBits, req.MaxBits)
+	}
+
+	return best, nil
+}

ssh/kex_test.go

@@ -9,6 +9,7 @@ package ssh
 import (
 	"crypto/rand"
 	"fmt"
+	"math/big"
 	"reflect"
 	"sync"
 	"testing"
@@ -65,6 +66,88 @@ func TestKexes(t *testing.T) {
 	}
 }
 
+func TestChooseDH(t *testing.T) {
+	oakley := map[int]string{
+		2048: oakleyGroup14,
+		3072: oakleyGroup15,
+		4096: oakleyGroup16,
+	}
+	expected := func(size int) *big.Int {
+		hex, ok := oakley[size]
+		if !ok {
+			t.Fatalf("test setup: no Oakley group for size %d", size)
+		}
+		p, _ := new(big.Int).SetString(hex, 16)
+		return p
+	}
+
+	tests := []struct {
+		name    string
+		request kexDHGexRequestMsg
+		want    int // expected bit size; 0 means error expected
+		wantErr bool
+	}{
+		{
+			name:    "Standard 2048 request",
+			request: kexDHGexRequestMsg{MinBits: 1024, PreferredBits: 2048, MaxBits: 8192},
+			want:    2048,
+		},
+		{
+			name:    "Standard 3072 request",
+			request: kexDHGexRequestMsg{MinBits: 1024, PreferredBits: 3072, MaxBits: 8192},
+			want:    3072,
+		},
+		{
+			name:    "Standard 4096 request",
+			request: kexDHGexRequestMsg{MinBits: 1024, PreferredBits: 4096, MaxBits: 8192},
+			want:    4096,
+		},
+		{
+			name:    "Preferred 2500 -> Expect 3072 (round up)",
+			request: kexDHGexRequestMsg{MinBits: 1024, PreferredBits: 2500, MaxBits: 8192},
+			want:    3072,
+		},
+		{
+			name:    "Preferred 3500 -> Expect 4096 (round up)",
+			request: kexDHGexRequestMsg{MinBits: 1024, PreferredBits: 3500, MaxBits: 8192},
+			want:    4096,
+		},
+		{
+			name:    "Preferred too high (8000) -> Expect 4096 (cap at max available)",
+			request: kexDHGexRequestMsg{MinBits: 1024, PreferredBits: 8000, MaxBits: 8192},
+			want:    4096,
+		},
+		{
+			name:    "No group in range",
+			request: kexDHGexRequestMsg{MinBits: 2500, PreferredBits: 2500, MaxBits: 2900},
+			wantErr: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := chooseDH(tt.request)
+
+			if (err != nil) != tt.wantErr {
+				t.Errorf("chooseDH() error = %v, wantErr %t", err, tt.wantErr)
+				return
+			}
+			if tt.wantErr {
+				return
+			}
+			if got == nil {
+				t.Fatalf("chooseDH() returned nil big.Int but expected success")
+			}
+			if got.BitLen() != tt.want {
+				t.Errorf("chooseDH() got size = %d, want %d", got.BitLen(), tt.want)
+			}
+			if want := expected(tt.want); got.Cmp(want) != 0 {
+				t.Errorf("chooseDH() returned the wrong group for size %d", tt.want)
+			}
+		})
+	}
+}
+
 func BenchmarkKexes(b *testing.B) {
 	type kexResultErr struct {
 		result *kexResult