usb-server/internal/client/stream_test.go

package client

import (
	"bytes"
	"errors"
	"io"
	"sync"
	"testing"
	"time"
)

func TestStreamBufferRoundTrip(t *testing.T) {
	s := newStreamBuffer()

	want := []byte("usbip frame")
	if _, err := s.Write(want); err != nil {
		t.Fatalf("Write: %v", err)
	}

	got := make([]byte, len(want))
	if _, err := io.ReadFull(s, got); err != nil {
		t.Fatalf("ReadFull: %v", err)
	}
	if !bytes.Equal(got, want) {
		t.Errorf("read %q, want %q", got, want)
	}
}

// The whole point of replacing io.Pipe: a write must return immediately even
// when nobody is reading, because it happens on the WebSocket read loop.
func TestStreamBufferWriteNeverBlocks(t *testing.T) {
	s := newStreamBuffer()

	done := make(chan struct{})
	go func() {
		defer close(done)
		for i := 0; i < 100; i++ {
			if _, err := s.Write(make([]byte, 1024)); err != nil {
				t.Errorf("Write %d: %v", i, err)
				return
			}
		}
	}()

	select {
	case <-done:
	case <-time.After(2 * time.Second):
		t.Fatal("writes blocked with no reader — this is what froze the client")
	}

	if got := s.Buffered(); got != 100*1024 {
		t.Errorf("buffered %d bytes, want %d", got, 100*1024)
	}
}

func TestStreamBufferReadBlocksUntilData(t *testing.T) {
	s := newStreamBuffer()

	read := make(chan []byte, 1)
	go func() {
		buf := make([]byte, 4)
		n, err := s.Read(buf)
		if err != nil {
			t.Errorf("Read: %v", err)
			read <- nil
			return
		}
		read <- buf[:n]
	}()

	// Give the reader time to park in Read before any data exists.
	time.Sleep(50 * time.Millisecond)
	select {
	case <-read:
		t.Fatal("Read returned before data was written")
	default:
	}

	s.Write([]byte("ping"))

	select {
	case got := <-read:
		if string(got) != "ping" {
			t.Errorf("read %q, want %q", got, "ping")
		}
	case <-time.After(time.Second):
		t.Fatal("Read did not wake up after Write")
	}
}

func TestStreamBufferCloseGivesEOFAfterDraining(t *testing.T) {
	s := newStreamBuffer()
	s.Write([]byte("tail"))
	s.Close()

	// Buffered data must still be readable after Close.
	got := make([]byte, 4)
	if _, err := io.ReadFull(s, got); err != nil {
		t.Fatalf("reading buffered data after Close: %v", err)
	}
	if string(got) != "tail" {
		t.Errorf("read %q, want %q", got, "tail")
	}

	if _, err := s.Read(make([]byte, 4)); err != io.EOF {
		t.Errorf("Read after drain = %v, want io.EOF", err)
	}
}

func TestStreamBufferCloseWakesBlockedReader(t *testing.T) {
	s := newStreamBuffer()

	errCh := make(chan error, 1)
	go func() {
		_, err := s.Read(make([]byte, 4))
		errCh <- err
	}()

	time.Sleep(50 * time.Millisecond)
	s.Close()

	select {
	case err := <-errCh:
		if err != io.EOF {
			t.Errorf("blocked Read woke with %v, want io.EOF", err)
		}
	case <-time.After(time.Second):
		t.Fatal("Close did not wake the blocked reader")
	}
}

func TestStreamBufferOverflowFailsInsteadOfGrowing(t *testing.T) {
	s := newStreamBufferLimit(1024)

	if _, err := s.Write(make([]byte, 1000)); err != nil {
		t.Fatalf("first write: %v", err)
	}
	if _, err := s.Write(make([]byte, 100)); !errors.Is(err, ErrStreamOverflow) {
		t.Fatalf("overflowing write = %v, want ErrStreamOverflow", err)
	}

	// Further writes keep failing rather than silently resuming.
	if _, err := s.Write([]byte("x")); !errors.Is(err, ErrStreamOverflow) {
		t.Errorf("write after overflow = %v, want ErrStreamOverflow", err)
	}

	// Buffered data is still drainable, then the error surfaces.
	if _, err := io.ReadFull(s, make([]byte, 1000)); err != nil {
		t.Fatalf("draining after overflow: %v", err)
	}
	if _, err := s.Read(make([]byte, 4)); !errors.Is(err, ErrStreamOverflow) {
		t.Errorf("Read after drain = %v, want ErrStreamOverflow", err)
	}
}

func TestStreamBufferWriteAfterClose(t *testing.T) {
	s := newStreamBuffer()
	s.Close()

	if _, err := s.Write([]byte("late")); err != io.ErrClosedPipe {
		t.Errorf("Write after Close = %v, want io.ErrClosedPipe", err)
	}
}

// Concurrent writers and one reader, the shape the share path actually has.
func TestStreamBufferConcurrent(t *testing.T) {
	s := newStreamBuffer()

	const writers = 8
	const perWriter = 200
	const chunk = 64

	var wg sync.WaitGroup
	wg.Add(writers)
	for i := 0; i < writers; i++ {
		go func() {
			defer wg.Done()
			for j := 0; j < perWriter; j++ {
				if _, err := s.Write(make([]byte, chunk)); err != nil {
					t.Errorf("Write: %v", err)
					return
				}
			}
		}()
	}

	total := writers * perWriter * chunk
	readDone := make(chan int, 1)
	go func() {
		got := 0
		buf := make([]byte, 128)
		for got < total {
			n, err := s.Read(buf)
			if err != nil {
				break
			}
			got += n
		}
		readDone <- got
	}()

	wg.Wait()

	select {
	case got := <-readDone:
		if got != total {
			t.Errorf("read %d bytes, want %d", got, total)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("concurrent read/write did not finish")
	}
}