diff --git a/ssh/kex.go b/ssh/kex.go index 5f7fdd8..91b771c 100644 --- a/ssh/kex.go +++ b/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 +} diff --git a/ssh/kex_test.go b/ssh/kex_test.go index cb7f66a..068afc7 100644 --- a/ssh/kex_test.go +++ b/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