github.com/lmb/consul@v1.4.1/connect/proxy/testing.go (about)

     1  package proxy
     2  
     3  import (
     4  	"fmt"
     5  	"io"
     6  	"log"
     7  	"net"
     8  	"sync/atomic"
     9  	"time"
    10  
    11  	"github.com/hashicorp/consul/connect"
    12  	"github.com/hashicorp/consul/lib/freeport"
    13  	"github.com/mitchellh/go-testing-interface"
    14  	"github.com/stretchr/testify/require"
    15  )
    16  
    17  // TestLocalAddr makes a localhost address on the given port
    18  func TestLocalAddr(port int) string {
    19  	return fmt.Sprintf("localhost:%d", port)
    20  }
    21  
    22  // TestTCPServer is a simple TCP echo server for use during tests.
    23  type TestTCPServer struct {
    24  	l                        net.Listener
    25  	stopped                  int32
    26  	accepted, closed, active int32
    27  }
    28  
    29  // NewTestTCPServer opens as a listening socket on the given address and returns
    30  // a TestTCPServer serving requests to it. The server is already started and can
    31  // be stopped by calling Close().
    32  func NewTestTCPServer(t testing.T) *TestTCPServer {
    33  	port := freeport.GetT(t, 1)
    34  	addr := TestLocalAddr(port[0])
    35  
    36  	l, err := net.Listen("tcp", addr)
    37  	require.NoError(t, err)
    38  
    39  	log.Printf("test tcp server listening on %s", addr)
    40  	s := &TestTCPServer{
    41  		l: l,
    42  	}
    43  	go s.accept()
    44  
    45  	return s
    46  }
    47  
    48  // Close stops the server
    49  func (s *TestTCPServer) Close() {
    50  	atomic.StoreInt32(&s.stopped, 1)
    51  	if s.l != nil {
    52  		s.l.Close()
    53  	}
    54  }
    55  
    56  // Addr returns the address that this server is listening on.
    57  func (s *TestTCPServer) Addr() net.Addr {
    58  	return s.l.Addr()
    59  }
    60  
    61  func (s *TestTCPServer) accept() error {
    62  	for {
    63  		conn, err := s.l.Accept()
    64  		if err != nil {
    65  			if atomic.LoadInt32(&s.stopped) == 1 {
    66  				log.Printf("test tcp echo server %s stopped", s.l.Addr())
    67  				return nil
    68  			}
    69  			log.Printf("test tcp echo server %s failed: %s", s.l.Addr(), err)
    70  			return err
    71  		}
    72  
    73  		atomic.AddInt32(&s.accepted, 1)
    74  		atomic.AddInt32(&s.active, 1)
    75  
    76  		go func(c net.Conn) {
    77  			io.Copy(c, c)
    78  			atomic.AddInt32(&s.closed, 1)
    79  			atomic.AddInt32(&s.active, -1)
    80  		}(conn)
    81  	}
    82  }
    83  
    84  // TestEchoConn attempts to write some bytes to conn and expects to read them
    85  // back within a short timeout (10ms). If prefix is not empty we expect it to be
    86  // poresent at the start of all echoed responses (for example to distinguish
    87  // between multiple echo server instances).
    88  func TestEchoConn(t testing.T, conn net.Conn, prefix string) {
    89  	t.Helper()
    90  
    91  	// Write some bytes and read them back
    92  	n, err := conn.Write([]byte("Hello World"))
    93  	require.Equal(t, 11, n)
    94  	require.Nil(t, err)
    95  
    96  	expectLen := 11 + len(prefix)
    97  
    98  	buf := make([]byte, expectLen)
    99  	// read until our buffer is full - it might be separate packets if prefix is
   100  	// in use.
   101  	got := 0
   102  	for got < expectLen {
   103  		n, err = conn.Read(buf[got:])
   104  		require.Nilf(t, err, "err: %s", err)
   105  		got += n
   106  	}
   107  	require.Equal(t, expectLen, got)
   108  	require.Equal(t, prefix+"Hello World", string(buf[:]))
   109  
   110  	// Addresses test flakiness around returning before Write or Read finish
   111  	// see PR #4498
   112  	time.Sleep(time.Millisecond)
   113  }
   114  
   115  // TestStaticUpstreamResolverFunc returns a function that will return a static
   116  // resolver for testing UpstreamListener.
   117  func TestStaticUpstreamResolverFunc(r connect.Resolver) func(UpstreamConfig) (connect.Resolver, error) {
   118  	return func(UpstreamConfig) (connect.Resolver, error) {
   119  		return r, nil
   120  	}
   121  }