github.com/mailgun/holster/v4@v4.20.0/udp/proxy_test.go (about)

     1  package udp_test
     2  
     3  import (
     4  	"fmt"
     5  	"net"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/sirupsen/logrus"
    10  
    11  	"github.com/mailgun/holster/v4/udp"
    12  	"github.com/stretchr/testify/assert"
    13  	"github.com/stretchr/testify/require"
    14  )
    15  
    16  func TestServerClient(t *testing.T) {
    17  	srv, err := udp.NewServer(udp.ServerConfig{
    18  		BindAddress: "127.0.0.1:5001",
    19  		Handler: func(conn net.PacketConn, recv []byte, addr net.Addr) {
    20  			resp := fmt.Sprintf("Hello, %s", string(recv))
    21  			_, err := conn.WriteTo([]byte(resp), addr)
    22  			require.NoError(t, err)
    23  		},
    24  	})
    25  	require.NoError(t, err)
    26  	defer srv.Close()
    27  
    28  	conn, err := udp.NewClient("127.0.0.1:5001")
    29  	require.NoError(t, err)
    30  	defer conn.Close()
    31  
    32  	_, err = conn.Send([]byte("Thrawn"))
    33  	require.NoError(t, err)
    34  
    35  	b := make([]byte, 50)
    36  	n, _, err := conn.Recv(b, time.Second)
    37  	require.NoError(t, err)
    38  
    39  	assert.Equal(t, "Hello, Thrawn", string(b[:n]))
    40  }
    41  
    42  func TestProxy(t *testing.T) {
    43  	logrus.SetLevel(logrus.DebugLevel)
    44  	srv, err := udp.NewServer(udp.ServerConfig{
    45  		BindAddress: "127.0.0.1:5001",
    46  		Handler: func(conn net.PacketConn, recv []byte, addr net.Addr) {
    47  			resp := fmt.Sprintf("Hello, %s", string(recv))
    48  			_, err := conn.WriteTo([]byte(resp), addr)
    49  			require.NoError(t, err)
    50  		},
    51  	})
    52  	require.NoError(t, err)
    53  	defer srv.Close()
    54  
    55  	p := udp.NewProxy(udp.ProxyConfig{
    56  		Listen:          "127.0.0.1:5000",
    57  		Upstream:        "127.0.0.1:5001",
    58  		UpstreamTimeout: time.Second,
    59  	})
    60  
    61  	// Start the proxy for our udp server
    62  	err = p.Start()
    63  	require.NoError(t, err)
    64  	defer p.Stop()
    65  
    66  	conn, err := udp.NewClient("127.0.0.1:5000")
    67  	require.NoError(t, err)
    68  	defer conn.Close()
    69  
    70  	_, err = conn.Send([]byte("Admiral"))
    71  	require.NoError(t, err)
    72  
    73  	b := make([]byte, 50)
    74  	n, _, err := conn.Recv(b, time.Second)
    75  	require.NoError(t, err)
    76  
    77  	assert.Equal(t, "Hello, Admiral", string(b[:n]))
    78  
    79  	// Shutdown the proxy
    80  	p.Stop()
    81  
    82  	// Should not get a response from the upstream server
    83  	_, err = conn.Send([]byte("Not expecting a response"))
    84  	require.NoError(t, err)
    85  	_, _, err = conn.Recv(b, time.Second)
    86  	require.Error(t, err)
    87  	assert.Contains(t, err.Error(), "recvfrom: connection refused")
    88  
    89  	// Start the proxy again
    90  	err = p.Start()
    91  	require.NoError(t, err)
    92  
    93  	// Should get a response
    94  	_, err = conn.Send([]byte("World"))
    95  	require.NoError(t, err)
    96  
    97  	n, _, err = conn.Recv(b, time.Second)
    98  	require.NoError(t, err)
    99  
   100  	assert.Equal(t, "Hello, World", string(b[:n]))
   101  
   102  }