ssh/channel_test.go - crypto - Git at Google

 // Copyright 2025 The Go Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.

 package ssh

 import (
 	"encoding/binary"
 	"io"
 	"math/bits"
 	"testing"
 	"time"
 	"unsafe"
 )

 func TestMinPayloadSize(t *testing.T) {
 	// 4 GiB (2^32). Declared as a var (not a const) so that int(bigPayload)
 	// is a runtime conversion: a constant conversion would fail to compile
 	// on 32-bit platforms with "constant 4294967296 overflows int". On
 	// 32-bit the value truncates to 0 at runtime, but the is64Bit cases
 	// that reference it are skipped by the runtime check below.
 	var bigPayload int64 = 1 << 32

 	tests := []struct {
 		name       string
 		maxPayload uint32
 		dataLen    int
 		want       uint32
 		is64Bit    bool // Flag to run only on 64-bit architectures
 	}{
 		{
 			name:       "Normal Case - Data fits in payload",
 			maxPayload: 32768,
 			dataLen:    1000,
 			want:       1000,
 		},
 		{
 			name:       "Normal Case - Data larger than payload",
 			maxPayload: 32768,
 			dataLen:    50000,
 			want:       32768,
 		},
 		{
 			name:       "Boundary Case - Data zero",
 			maxPayload: 32768,
 			dataLen:    0,
 			want:       0,
 		},
 		{
 			name:       "Overflow Case - Data is exactly 4GB (1<<32)",
 			maxPayload: 32768,
 			dataLen:    int(bigPayload),
 			want:       32768,
 			is64Bit:    true,
 		},
 		{
 			name:       "Overflow Case - Data is 4GB + small amount",
 			maxPayload: 32768,
 			dataLen:    int(bigPayload + 100),
 			want:       32768,
 			is64Bit:    true,
 		},
 	}

 	is64Bit := bits.UintSize == 64

 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			if tt.is64Bit && !is64Bit {
 				t.Skip("Skipping test requiring 64-bit int")
 			}
 			got := minPayloadSize(tt.maxPayload, tt.dataLen)
 			if got != tt.want {
 				t.Errorf("minPayloadSize(%d, %d) = %d; want %d", tt.maxPayload, tt.dataLen, got, tt.want)
 			}
 		})
 	}
 }

 // TestWriteExtendedNoInfiniteLoopOnLargeWrite is an end-to-end regression
 // test for the integer-overflow bug in WriteExtended. Before the fix, a
 // write whose len(data) was a multiple of 2^32 caused minPayloadSize to
 // return 0; WriteExtended then spun forever, reserving 0 bytes per
 // iteration and never advancing the data slice.
 //
 // We exercise the real WriteExtended path with a slice whose declared
 // length is exactly 2^32. Allocating 4 GiB is unnecessary: each iteration
 // only reads up to maxRemotePayload bytes from the head of the slice, and
 // the loop blocks in remoteWin.reserve() once the channel window is
 // exhausted — before the slice base advances past the underlying buffer.
 //
 // With the fix, the loop blocks in reserve(); we detect that via
 // waitWriterBlocked(), then close the window to let WriteExtended return.
 // With the bug, the loop never blocks and the test times out.
 //
 //go:nocheckptr
 func TestWriteExtendedNoInfiniteLoopOnLargeWrite(t *testing.T) {
 	if bits.UintSize < 64 {
 		t.Skip("test requires 64-bit int to construct a slice with len >= 2^32")
 	}

 	reader, writer, mux := channelPair(t)
 	defer reader.Close()
 	defer writer.Close()
 	defer mux.Close()

 	// Sized to hold the full pre-update remote window so that no iteration
 	// reads past the backing buffer before reserve() blocks.
 	backing := make([]byte, channelWindowSize)
 	var bigLen int64 = 1 << 32
 	bigSlice := unsafe.Slice(&backing[0], int(bigLen))

 	done := make(chan int, 1)
 	go func() {
 		n, _ := writer.Write(bigSlice)
 		done <- n
 	}()

 	blocked := make(chan struct{})
 	go func() {
 		writer.remoteWin.waitWriterBlocked()
 		close(blocked)
 	}()

 	select {
 	case <-blocked:
 		// Good — the loop made progress and is now blocked in reserve().
 		// Close the window to let WriteExtended return.
 		writer.remoteWin.close()
 	case <-time.After(2 * time.Second):
 		t.Fatal("WriteExtended did not block in reserve within 2s — minPayloadSize likely returned 0 (integer overflow regression)")
 	}

 	select {
 	case n := <-done:
 		if n == 0 {
 			t.Fatalf("WriteExtended returned n=0; expected progress")
 		}
 	case <-time.After(2 * time.Second):
 		t.Fatal("WriteExtended did not return after closing the window")
 	}
 }

 func TestDiscardedExtendedDataReturnsWindowCredit(t *testing.T) {
 	client, _, serverChans := forwardingPair(t)

 	clientCh, reqs, err := client.OpenChannel("test", nil)
 	if err != nil {
 		t.Fatalf("OpenChannel: %v", err)
 	}
 	go DiscardRequests(reqs)
 	defer clientCh.Close()

 	serverCh := (<-serverChans).(*channel)
 	defer serverCh.Close()

 	// Write more than the default channel receive window (2 MiB) of
 	// extended data with a type code other than stderr. The peer
 	// discards such data: if it did not return the window credit the
 	// write would block forever once the window is exhausted.
 	const payload = 4 * 1024 * 1024
 	done := make(chan error, 1)
 	go func() {
 		buf := make([]byte, 32*1024)
 		remaining := payload
 		for remaining > 0 {
 			n := min(len(buf), remaining)
 			nw, err := serverCh.WriteExtended(buf[:n], 2)
 			if err != nil {
 				done <- err
 				return
 			}
 			remaining -= nw
 		}
 		done <- nil
 	}()

 	select {
 	case err := <-done:
 		if err != nil {
 			t.Fatalf("write of discarded extended data: %v", err)
 		}
 	case <-time.After(10 * time.Second):
 		t.Fatal("write of discarded extended data blocked: window credit is not returned")
 	}

 	// The main stream must still be usable.
 	want := []byte("main stream after extended data")
 	if _, err := serverCh.Write(want); err != nil {
 		t.Fatalf("write to main stream: %v", err)
 	}
 	got := make([]byte, len(want))
 	if _, err := io.ReadFull(clientCh, got); err != nil {
 		t.Fatalf("read from main stream: %v", err)
 	}
 	if string(got) != string(want) {
 		t.Fatalf("read %q from main stream, want %q", got, want)
 	}

 	// Reading the main stream payload above orders the test after all
 	// preceding packets: everything received has now been either read or
 	// discarded, so the full window credit must have been accounted.
 	// Anything less indicates a partial-credit leak that the liveness
 	// check above would not detect.
 	cc := clientCh.(*channel)
 	cc.windowMu.Lock()
 	window := cc.myWindow + cc.myConsumed
 	cc.windowMu.Unlock()
 	if window != channelWindowSize {
 		t.Errorf("myWindow+myConsumed = %d, want %d", window, channelWindowSize)
 	}
 }

 func TestDiscardedExtendedDataAfterClose(t *testing.T) {
 	client, _, serverChans := forwardingPair(t)

 	clientCh, reqs, err := client.OpenChannel("test", nil)
 	if err != nil {
 		t.Fatalf("OpenChannel: %v", err)
 	}
 	go DiscardRequests(reqs)
 	serverCh := <-serverChans
 	defer serverCh.Close()

 	ch := clientCh.(*channel)
 	// Close the channel locally: sentClose is now set, but the channel
 	// stays in the mux chanList until the peer's close arrives, so
 	// in-flight data packets are still dispatched to handleData.
 	if err := ch.Close(); err != nil {
 		t.Fatalf("Close: %v", err)
 	}

 	// Deliver extended data packets with an unknown type code, enough to
 	// cross the adjustWindow send threshold of 3*maxIncomingPayload.
 	// The window adjust message cannot be sent after the close and
 	// adjustWindow fails with io.EOF: handleData must swallow that
 	// error, since any error it returns terminates the mux read loop
 	// and tears down every channel on the connection.
 	const packetLen = 32 * 1024
 	packet := make([]byte, 13+packetLen)
 	packet[0] = msgChannelExtendedData
 	binary.BigEndian.PutUint32(packet[1:], ch.localId)
 	binary.BigEndian.PutUint32(packet[5:], 2)
 	binary.BigEndian.PutUint32(packet[9:], packetLen)
 	for range 5 {
 		if err := ch.handleData(packet); err != nil {
 			t.Fatalf("handleData after local close: %v", err)
 		}
 	}
 }
	// Copyright 2025 The Go Authors. All rights reserved.
	// Use of this source code is governed by a BSD-style
	// license that can be found in the LICENSE file.

	package ssh

	import (
	"encoding/binary"
	"io"
	"math/bits"
	"testing"
	"time"
	"unsafe"
	)

	func TestMinPayloadSize(t *testing.T) {
	// 4 GiB (2^32). Declared as a var (not a const) so that int(bigPayload)
	// is a runtime conversion: a constant conversion would fail to compile
	// on 32-bit platforms with "constant 4294967296 overflows int". On
	// 32-bit the value truncates to 0 at runtime, but the is64Bit cases
	// that reference it are skipped by the runtime check below.
	var bigPayload int64 = 1 << 32

	tests := []struct {
	name string
	maxPayload uint32
	dataLen int
	want uint32
	is64Bit bool // Flag to run only on 64-bit architectures
	}{
	{
	name: "Normal Case - Data fits in payload",
	maxPayload: 32768,
	dataLen: 1000,
	want: 1000,
	},
	{
	name: "Normal Case - Data larger than payload",
	maxPayload: 32768,
	dataLen: 50000,
	want: 32768,
	},
	{
	name: "Boundary Case - Data zero",
	maxPayload: 32768,
	dataLen: 0,
	want: 0,
	},
	{
	name: "Overflow Case - Data is exactly 4GB (1<<32)",
	maxPayload: 32768,
	dataLen: int(bigPayload),
	want: 32768,
	is64Bit: true,
	},
	{
	name: "Overflow Case - Data is 4GB + small amount",
	maxPayload: 32768,
	dataLen: int(bigPayload + 100),
	want: 32768,
	is64Bit: true,
	},
	}

	is64Bit := bits.UintSize == 64

	for _, tt := range tests {
	t.Run(tt.name, func(t *testing.T) {
	if tt.is64Bit && !is64Bit {
	t.Skip("Skipping test requiring 64-bit int")
	}
	got := minPayloadSize(tt.maxPayload, tt.dataLen)
	if got != tt.want {
	t.Errorf("minPayloadSize(%d, %d) = %d; want %d", tt.maxPayload, tt.dataLen, got, tt.want)
	}
	})
	}
	}

	// TestWriteExtendedNoInfiniteLoopOnLargeWrite is an end-to-end regression
	// test for the integer-overflow bug in WriteExtended. Before the fix, a
	// write whose len(data) was a multiple of 2^32 caused minPayloadSize to
	// return 0; WriteExtended then spun forever, reserving 0 bytes per
	// iteration and never advancing the data slice.
	//
	// We exercise the real WriteExtended path with a slice whose declared
	// length is exactly 2^32. Allocating 4 GiB is unnecessary: each iteration
	// only reads up to maxRemotePayload bytes from the head of the slice, and
	// the loop blocks in remoteWin.reserve() once the channel window is
	// exhausted — before the slice base advances past the underlying buffer.
	//
	// With the fix, the loop blocks in reserve(); we detect that via
	// waitWriterBlocked(), then close the window to let WriteExtended return.
	// With the bug, the loop never blocks and the test times out.
	//
	//go:nocheckptr
	func TestWriteExtendedNoInfiniteLoopOnLargeWrite(t *testing.T) {
	if bits.UintSize < 64 {
	t.Skip("test requires 64-bit int to construct a slice with len >= 2^32")
	}

	reader, writer, mux := channelPair(t)
	defer reader.Close()
	defer writer.Close()
	defer mux.Close()

	// Sized to hold the full pre-update remote window so that no iteration
	// reads past the backing buffer before reserve() blocks.
	backing := make([]byte, channelWindowSize)
	var bigLen int64 = 1 << 32
	bigSlice := unsafe.Slice(&backing[0], int(bigLen))

	done := make(chan int, 1)
	go func() {
	n, _ := writer.Write(bigSlice)
	done <- n
	}()

	blocked := make(chan struct{})
	go func() {
	writer.remoteWin.waitWriterBlocked()
	close(blocked)
	}()

	select {
	case <-blocked:
	// Good — the loop made progress and is now blocked in reserve().
	// Close the window to let WriteExtended return.
	writer.remoteWin.close()
	case <-time.After(2 * time.Second):
	t.Fatal("WriteExtended did not block in reserve within 2s — minPayloadSize likely returned 0 (integer overflow regression)")
	}

	select {
	case n := <-done:
	if n == 0 {
	t.Fatalf("WriteExtended returned n=0; expected progress")
	}
	case <-time.After(2 * time.Second):
	t.Fatal("WriteExtended did not return after closing the window")
	}
	}

	func TestDiscardedExtendedDataReturnsWindowCredit(t *testing.T) {
	client, _, serverChans := forwardingPair(t)

	clientCh, reqs, err := client.OpenChannel("test", nil)
	if err != nil {
	t.Fatalf("OpenChannel: %v", err)
	}
	go DiscardRequests(reqs)
	defer clientCh.Close()

	serverCh := (<-serverChans).(*channel)
	defer serverCh.Close()

	// Write more than the default channel receive window (2 MiB) of
	// extended data with a type code other than stderr. The peer
	// discards such data: if it did not return the window credit the
	// write would block forever once the window is exhausted.
	const payload = 4 * 1024 * 1024
	done := make(chan error, 1)
	go func() {
	buf := make([]byte, 32*1024)
	remaining := payload
	for remaining > 0 {
	n := min(len(buf), remaining)
	nw, err := serverCh.WriteExtended(buf[:n], 2)
	if err != nil {
	done <- err
	return
	}
	remaining -= nw
	}
	done <- nil
	}()

	select {
	case err := <-done:
	if err != nil {
	t.Fatalf("write of discarded extended data: %v", err)
	}
	case <-time.After(10 * time.Second):
	t.Fatal("write of discarded extended data blocked: window credit is not returned")
	}

	// The main stream must still be usable.
	want := []byte("main stream after extended data")
	if _, err := serverCh.Write(want); err != nil {
	t.Fatalf("write to main stream: %v", err)
	}
	got := make([]byte, len(want))
	if _, err := io.ReadFull(clientCh, got); err != nil {
	t.Fatalf("read from main stream: %v", err)
	}
	if string(got) != string(want) {
	t.Fatalf("read %q from main stream, want %q", got, want)
	}

	// Reading the main stream payload above orders the test after all
	// preceding packets: everything received has now been either read or
	// discarded, so the full window credit must have been accounted.
	// Anything less indicates a partial-credit leak that the liveness
	// check above would not detect.
	cc := clientCh.(*channel)
	cc.windowMu.Lock()
	window := cc.myWindow + cc.myConsumed
	cc.windowMu.Unlock()
	if window != channelWindowSize {
	t.Errorf("myWindow+myConsumed = %d, want %d", window, channelWindowSize)
	}
	}

	func TestDiscardedExtendedDataAfterClose(t *testing.T) {
	client, _, serverChans := forwardingPair(t)

	clientCh, reqs, err := client.OpenChannel("test", nil)
	if err != nil {
	t.Fatalf("OpenChannel: %v", err)
	}
	go DiscardRequests(reqs)
	serverCh := <-serverChans
	defer serverCh.Close()

	ch := clientCh.(*channel)
	// Close the channel locally: sentClose is now set, but the channel
	// stays in the mux chanList until the peer's close arrives, so
	// in-flight data packets are still dispatched to handleData.
	if err := ch.Close(); err != nil {
	t.Fatalf("Close: %v", err)
	}

	// Deliver extended data packets with an unknown type code, enough to
	// cross the adjustWindow send threshold of 3*maxIncomingPayload.
	// The window adjust message cannot be sent after the close and
	// adjustWindow fails with io.EOF: handleData must swallow that
	// error, since any error it returns terminates the mux read loop
	// and tears down every channel on the connection.
	const packetLen = 32 * 1024
	packet := make([]byte, 13+packetLen)
	packet[0] = msgChannelExtendedData
	binary.BigEndian.PutUint32(packet[1:], ch.localId)
	binary.BigEndian.PutUint32(packet[5:], 2)
	binary.BigEndian.PutUint32(packet[9:], packetLen)
	for range 5 {
	if err := ch.handleData(packet); err != nil {
	t.Fatalf("handleData after local close: %v", err)
	}
	}
	}