github.com/ava-labs/avalanchego@v1.11.11/network/dialer_test.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package network
     5  
     6  import (
     7  	"context"
     8  	"errors"
     9  	"net"
    10  	"net/netip"
    11  
    12  	"github.com/ava-labs/avalanchego/network/dialer"
    13  )
    14  
    15  var (
    16  	errRefused = errors.New("connection refused")
    17  
    18  	_ dialer.Dialer = (*testDialer)(nil)
    19  )
    20  
    21  type testDialer struct {
    22  	// maps [ip.String] to a listener
    23  	listeners map[netip.AddrPort]*testListener
    24  }
    25  
    26  func newTestDialer() *testDialer {
    27  	return &testDialer{
    28  		listeners: make(map[netip.AddrPort]*testListener),
    29  	}
    30  }
    31  
    32  func (d *testDialer) NewListener() (netip.AddrPort, *testListener) {
    33  	// Uses a private IP to easily enable testing AllowPrivateIPs
    34  	addrPort := netip.AddrPortFrom(
    35  		netip.AddrFrom4([4]byte{10, 0, 0, 0}),
    36  		uint16(len(d.listeners)+1),
    37  	)
    38  	listener := newTestListener(addrPort)
    39  	d.AddListener(addrPort, listener)
    40  	return addrPort, listener
    41  }
    42  
    43  func (d *testDialer) AddListener(ip netip.AddrPort, listener *testListener) {
    44  	d.listeners[ip] = listener
    45  }
    46  
    47  func (d *testDialer) Dial(ctx context.Context, ip netip.AddrPort) (net.Conn, error) {
    48  	listener, ok := d.listeners[ip]
    49  	if !ok {
    50  		return nil, errRefused
    51  	}
    52  	serverConn, clientConn := net.Pipe()
    53  	server := &testConn{
    54  		Conn: serverConn,
    55  		localAddr: &net.TCPAddr{
    56  			IP:   net.IPv6loopback,
    57  			Port: 1,
    58  		},
    59  		remoteAddr: &net.TCPAddr{
    60  			IP:   net.IPv6loopback,
    61  			Port: 2,
    62  		},
    63  	}
    64  	client := &testConn{
    65  		Conn: clientConn,
    66  		localAddr: &net.TCPAddr{
    67  			IP:   net.IPv6loopback,
    68  			Port: 3,
    69  		},
    70  		remoteAddr: &net.TCPAddr{
    71  			IP:   net.IPv6loopback,
    72  			Port: 4,
    73  		},
    74  	}
    75  	select {
    76  	case listener.inbound <- server:
    77  		return client, nil
    78  	case <-ctx.Done():
    79  		return nil, ctx.Err()
    80  	case <-listener.closed:
    81  		return nil, errRefused
    82  	}
    83  }