golang.org/x/net@v0.25.1-0.20240516223405-c87a5b62e243/quic/endpoint_test.go (about)

     1  // Copyright 2023 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  //go:build go1.21
     6  
     7  package quic
     8  
     9  import (
    10  	"bytes"
    11  	"context"
    12  	"crypto/tls"
    13  	"io"
    14  	"log/slog"
    15  	"net/netip"
    16  	"testing"
    17  	"time"
    18  
    19  	"golang.org/x/net/quic/qlog"
    20  )
    21  
    22  func TestConnect(t *testing.T) {
    23  	newLocalConnPair(t, &Config{}, &Config{})
    24  }
    25  
    26  func TestStreamTransfer(t *testing.T) {
    27  	ctx := context.Background()
    28  	cli, srv := newLocalConnPair(t, &Config{}, &Config{})
    29  	data := makeTestData(1 << 20)
    30  
    31  	srvdone := make(chan struct{})
    32  	go func() {
    33  		defer close(srvdone)
    34  		s, err := srv.AcceptStream(ctx)
    35  		if err != nil {
    36  			t.Errorf("AcceptStream: %v", err)
    37  			return
    38  		}
    39  		b, err := io.ReadAll(s)
    40  		if err != nil {
    41  			t.Errorf("io.ReadAll(s): %v", err)
    42  			return
    43  		}
    44  		if !bytes.Equal(b, data) {
    45  			t.Errorf("read data mismatch (got %v bytes, want %v", len(b), len(data))
    46  		}
    47  		if err := s.Close(); err != nil {
    48  			t.Errorf("s.Close() = %v", err)
    49  		}
    50  	}()
    51  
    52  	s, err := cli.NewSendOnlyStream(ctx)
    53  	if err != nil {
    54  		t.Fatalf("NewStream: %v", err)
    55  	}
    56  	n, err := io.Copy(s, bytes.NewBuffer(data))
    57  	if n != int64(len(data)) || err != nil {
    58  		t.Fatalf("io.Copy(s, data) = %v, %v; want %v, nil", n, err, len(data))
    59  	}
    60  	if err := s.Close(); err != nil {
    61  		t.Fatalf("s.Close() = %v", err)
    62  	}
    63  }
    64  
    65  func newLocalConnPair(t testing.TB, conf1, conf2 *Config) (clientConn, serverConn *Conn) {
    66  	t.Helper()
    67  	ctx := context.Background()
    68  	e1 := newLocalEndpoint(t, serverSide, conf1)
    69  	e2 := newLocalEndpoint(t, clientSide, conf2)
    70  	conf2 = makeTestConfig(conf2, clientSide)
    71  	c2, err := e2.Dial(ctx, "udp", e1.LocalAddr().String(), conf2)
    72  	if err != nil {
    73  		t.Fatal(err)
    74  	}
    75  	c1, err := e1.Accept(ctx)
    76  	if err != nil {
    77  		t.Fatal(err)
    78  	}
    79  	return c2, c1
    80  }
    81  
    82  func newLocalEndpoint(t testing.TB, side connSide, conf *Config) *Endpoint {
    83  	t.Helper()
    84  	conf = makeTestConfig(conf, side)
    85  	e, err := Listen("udp", "127.0.0.1:0", conf)
    86  	if err != nil {
    87  		t.Fatal(err)
    88  	}
    89  	t.Cleanup(func() {
    90  		e.Close(canceledContext())
    91  	})
    92  	return e
    93  }
    94  
    95  func makeTestConfig(conf *Config, side connSide) *Config {
    96  	if conf == nil {
    97  		return nil
    98  	}
    99  	newConf := *conf
   100  	conf = &newConf
   101  	if conf.TLSConfig == nil {
   102  		conf.TLSConfig = newTestTLSConfig(side)
   103  	}
   104  	if conf.QLogLogger == nil {
   105  		conf.QLogLogger = slog.New(qlog.NewJSONHandler(qlog.HandlerOptions{
   106  			Level: QLogLevelFrame,
   107  			Dir:   *qlogdir,
   108  		}))
   109  	}
   110  	return conf
   111  }
   112  
   113  type testEndpoint struct {
   114  	t                     *testing.T
   115  	e                     *Endpoint
   116  	now                   time.Time
   117  	recvc                 chan *datagram
   118  	idlec                 chan struct{}
   119  	conns                 map[*Conn]*testConn
   120  	acceptQueue           []*testConn
   121  	configTransportParams []func(*transportParameters)
   122  	configTestConn        []func(*testConn)
   123  	sentDatagrams         [][]byte
   124  	peerTLSConn           *tls.QUICConn
   125  	lastInitialDstConnID  []byte // for parsing Retry packets
   126  }
   127  
   128  func newTestEndpoint(t *testing.T, config *Config) *testEndpoint {
   129  	te := &testEndpoint{
   130  		t:     t,
   131  		now:   time.Date(2020, 1, 1, 0, 0, 0, 0, time.UTC),
   132  		recvc: make(chan *datagram),
   133  		idlec: make(chan struct{}),
   134  		conns: make(map[*Conn]*testConn),
   135  	}
   136  	var err error
   137  	te.e, err = newEndpoint((*testEndpointUDPConn)(te), config, (*testEndpointHooks)(te))
   138  	if err != nil {
   139  		t.Fatal(err)
   140  	}
   141  	t.Cleanup(te.cleanup)
   142  	return te
   143  }
   144  
   145  func (te *testEndpoint) cleanup() {
   146  	te.e.Close(canceledContext())
   147  }
   148  
   149  func (te *testEndpoint) wait() {
   150  	select {
   151  	case te.idlec <- struct{}{}:
   152  	case <-te.e.closec:
   153  	}
   154  	for _, tc := range te.conns {
   155  		tc.wait()
   156  	}
   157  }
   158  
   159  // accept returns a server connection from the endpoint.
   160  // Unlike Endpoint.Accept, connections are available as soon as they are created.
   161  func (te *testEndpoint) accept() *testConn {
   162  	if len(te.acceptQueue) == 0 {
   163  		te.t.Fatalf("accept: expected available conn, but found none")
   164  	}
   165  	tc := te.acceptQueue[0]
   166  	te.acceptQueue = te.acceptQueue[1:]
   167  	return tc
   168  }
   169  
   170  func (te *testEndpoint) write(d *datagram) {
   171  	te.recvc <- d
   172  	te.wait()
   173  }
   174  
   175  var testClientAddr = netip.MustParseAddrPort("10.0.0.1:8000")
   176  
   177  func (te *testEndpoint) writeDatagram(d *testDatagram) {
   178  	te.t.Helper()
   179  	logDatagram(te.t, "<- endpoint under test receives", d)
   180  	var buf []byte
   181  	for _, p := range d.packets {
   182  		tc := te.connForDestination(p.dstConnID)
   183  		if p.ptype != packetTypeRetry && tc != nil {
   184  			space := spaceForPacketType(p.ptype)
   185  			if p.num >= tc.peerNextPacketNum[space] {
   186  				tc.peerNextPacketNum[space] = p.num + 1
   187  			}
   188  		}
   189  		if p.ptype == packetTypeInitial {
   190  			te.lastInitialDstConnID = p.dstConnID
   191  		}
   192  		pad := 0
   193  		if p.ptype == packetType1RTT {
   194  			pad = d.paddedSize - len(buf)
   195  		}
   196  		buf = append(buf, encodeTestPacket(te.t, tc, p, pad)...)
   197  	}
   198  	for len(buf) < d.paddedSize {
   199  		buf = append(buf, 0)
   200  	}
   201  	te.write(&datagram{
   202  		b:        buf,
   203  		peerAddr: d.addr,
   204  	})
   205  }
   206  
   207  func (te *testEndpoint) connForDestination(dstConnID []byte) *testConn {
   208  	for _, tc := range te.conns {
   209  		for _, loc := range tc.conn.connIDState.local {
   210  			if bytes.Equal(loc.cid, dstConnID) {
   211  				return tc
   212  			}
   213  		}
   214  	}
   215  	return nil
   216  }
   217  
   218  func (te *testEndpoint) connForSource(srcConnID []byte) *testConn {
   219  	for _, tc := range te.conns {
   220  		for _, loc := range tc.conn.connIDState.remote {
   221  			if bytes.Equal(loc.cid, srcConnID) {
   222  				return tc
   223  			}
   224  		}
   225  	}
   226  	return nil
   227  }
   228  
   229  func (te *testEndpoint) read() []byte {
   230  	te.t.Helper()
   231  	te.wait()
   232  	if len(te.sentDatagrams) == 0 {
   233  		return nil
   234  	}
   235  	d := te.sentDatagrams[0]
   236  	te.sentDatagrams = te.sentDatagrams[1:]
   237  	return d
   238  }
   239  
   240  func (te *testEndpoint) readDatagram() *testDatagram {
   241  	te.t.Helper()
   242  	buf := te.read()
   243  	if buf == nil {
   244  		return nil
   245  	}
   246  	p, _ := parseGenericLongHeaderPacket(buf)
   247  	tc := te.connForSource(p.dstConnID)
   248  	d := parseTestDatagram(te.t, te, tc, buf)
   249  	logDatagram(te.t, "-> endpoint under test sends", d)
   250  	return d
   251  }
   252  
   253  // wantDatagram indicates that we expect the Endpoint to send a datagram.
   254  func (te *testEndpoint) wantDatagram(expectation string, want *testDatagram) {
   255  	te.t.Helper()
   256  	got := te.readDatagram()
   257  	if !datagramEqual(got, want) {
   258  		te.t.Fatalf("%v:\ngot datagram:  %v\nwant datagram: %v", expectation, got, want)
   259  	}
   260  }
   261  
   262  // wantIdle indicates that we expect the Endpoint to not send any more datagrams.
   263  func (te *testEndpoint) wantIdle(expectation string) {
   264  	if got := te.readDatagram(); got != nil {
   265  		te.t.Fatalf("expect: %v\nunexpectedly got: %v", expectation, got)
   266  	}
   267  }
   268  
   269  // advance causes time to pass.
   270  func (te *testEndpoint) advance(d time.Duration) {
   271  	te.t.Helper()
   272  	te.advanceTo(te.now.Add(d))
   273  }
   274  
   275  // advanceTo sets the current time.
   276  func (te *testEndpoint) advanceTo(now time.Time) {
   277  	te.t.Helper()
   278  	if te.now.After(now) {
   279  		te.t.Fatalf("time moved backwards: %v -> %v", te.now, now)
   280  	}
   281  	te.now = now
   282  	for _, tc := range te.conns {
   283  		if !tc.timer.After(te.now) {
   284  			tc.conn.sendMsg(timerEvent{})
   285  			tc.wait()
   286  		}
   287  	}
   288  }
   289  
   290  // testEndpointHooks implements endpointTestHooks.
   291  type testEndpointHooks testEndpoint
   292  
   293  func (te *testEndpointHooks) timeNow() time.Time {
   294  	return te.now
   295  }
   296  
   297  func (te *testEndpointHooks) newConn(c *Conn) {
   298  	tc := newTestConnForConn(te.t, (*testEndpoint)(te), c)
   299  	te.conns[c] = tc
   300  }
   301  
   302  // testEndpointUDPConn implements UDPConn.
   303  type testEndpointUDPConn testEndpoint
   304  
   305  func (te *testEndpointUDPConn) Close() error {
   306  	close(te.recvc)
   307  	return nil
   308  }
   309  
   310  func (te *testEndpointUDPConn) LocalAddr() netip.AddrPort {
   311  	return netip.MustParseAddrPort("127.0.0.1:443")
   312  }
   313  
   314  func (te *testEndpointUDPConn) Read(f func(*datagram)) {
   315  	for {
   316  		select {
   317  		case d, ok := <-te.recvc:
   318  			if !ok {
   319  				return
   320  			}
   321  			f(d)
   322  		case <-te.idlec:
   323  		}
   324  	}
   325  }
   326  
   327  func (te *testEndpointUDPConn) Write(dgram datagram) error {
   328  	te.sentDatagrams = append(te.sentDatagrams, append([]byte(nil), dgram.b...))
   329  	return nil
   330  }