github.com/mholt/caddy-l4@v0.0.0-20241104153248-ec8fae209322/modules/l4proxyprotocol/handler_test.go (about)

     1  package l4proxyprotocol
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"net"
     7  	"sync"
     8  	"testing"
     9  
    10  	"github.com/caddyserver/caddy/v2"
    11  	"go.uber.org/zap"
    12  
    13  	"github.com/mholt/caddy-l4/layer4"
    14  )
    15  
    16  func assertString(t *testing.T, expected string, value string) {
    17  	t.Helper()
    18  	if value != expected {
    19  		t.Fatalf("Expected '%s' but got '%s'\n", expected, value)
    20  	}
    21  }
    22  
    23  func TestProxyProtocolHandleV1(t *testing.T) {
    24  	wg := &sync.WaitGroup{}
    25  	in, out := net.Pipe()
    26  	defer closePipe(wg, in, out)
    27  
    28  	cx := layer4.WrapConnection(in, []byte{}, zap.NewNop())
    29  	go func() {
    30  		wg.Add(1)
    31  		defer wg.Done()
    32  		defer func() { _ = out.Close() }()
    33  		_, err := out.Write(ProxyV1Example)
    34  		assertNoError(t, err)
    35  	}()
    36  
    37  	ctx, cancel := caddy.NewContext(caddy.Context{Context: context.Background()})
    38  	defer cancel()
    39  
    40  	handler := Handler{}
    41  	err := handler.Provision(ctx)
    42  	assertNoError(t, err)
    43  
    44  	var nextCx *layer4.Connection
    45  	err = handler.Handle(cx, layer4.HandlerFunc(func(c *layer4.Connection) error {
    46  		nextCx = c
    47  		return nil
    48  	}))
    49  	assertNoError(t, err)
    50  
    51  	if nextCx == nil {
    52  		t.Fatalf("handler did not call next")
    53  	}
    54  
    55  	assertString(t, "192.168.0.1:56324", nextCx.RemoteAddr().String())
    56  	assertString(t, "192.168.0.11:443", nextCx.LocalAddr().String())
    57  
    58  	_, _ = io.Copy(io.Discard, in)
    59  }
    60  
    61  func TestProxyProtocolHandleV2(t *testing.T) {
    62  	wg := &sync.WaitGroup{}
    63  	in, out := net.Pipe()
    64  	defer closePipe(wg, in, out)
    65  
    66  	cx := layer4.WrapConnection(in, []byte{}, zap.NewNop())
    67  	go func() {
    68  		wg.Add(1)
    69  		defer wg.Done()
    70  		defer func() { _ = out.Close() }()
    71  		_, err := out.Write(ProxyV2Example)
    72  		assertNoError(t, err)
    73  	}()
    74  
    75  	ctx, cancel := caddy.NewContext(caddy.Context{Context: context.Background()})
    76  	defer cancel()
    77  
    78  	handler := Handler{}
    79  	err := handler.Provision(ctx)
    80  	assertNoError(t, err)
    81  
    82  	var nextCx *layer4.Connection
    83  	err = handler.Handle(cx, layer4.HandlerFunc(func(c *layer4.Connection) error {
    84  		nextCx = c
    85  		return nil
    86  	}))
    87  	assertNoError(t, err)
    88  
    89  	if nextCx == nil {
    90  		t.Fatalf("handler did not call next")
    91  	}
    92  
    93  	assertString(t, "127.0.0.1:47111", nextCx.RemoteAddr().String())
    94  	assertString(t, "127.0.0.1:443", nextCx.LocalAddr().String())
    95  
    96  	_, _ = io.Copy(io.Discard, in)
    97  }
    98  
    99  func TestProxyProtocolHandleGarbage(t *testing.T) {
   100  	wg := &sync.WaitGroup{}
   101  	in, out := net.Pipe()
   102  	defer closePipe(wg, in, out)
   103  
   104  	cx := layer4.WrapConnection(in, []byte{}, zap.NewNop())
   105  	go func() {
   106  		wg.Add(1)
   107  		defer wg.Done()
   108  		defer func() { _ = out.Close() }()
   109  		_, err := out.Write([]byte("some garbage"))
   110  		assertNoError(t, err)
   111  	}()
   112  
   113  	ctx, cancel := caddy.NewContext(caddy.Context{Context: context.Background()})
   114  	defer cancel()
   115  
   116  	handler := Handler{}
   117  	err := handler.Provision(ctx)
   118  	assertNoError(t, err)
   119  
   120  	var nextCx *layer4.Connection
   121  	err = handler.Handle(cx, layer4.HandlerFunc(func(c *layer4.Connection) error {
   122  		nextCx = c
   123  		return nil
   124  	}))
   125  	if err == nil || err.Error() != "parsing the PROXY header: invalid signature" {
   126  		t.Fatalf("handler did not return an error or the wrong error -> %s", err)
   127  	}
   128  
   129  	if nextCx != nil {
   130  		t.Fatalf("handler did call next")
   131  	}
   132  
   133  	_, _ = io.Copy(io.Discard, in)
   134  }