github.com/mholt/caddy-l4@v0.0.0-20241104153248-ec8fae209322/layer4/connection_test.go (about)

     1  package layer4
     2  
     3  import (
     4  	"bytes"
     5  	"net"
     6  	"testing"
     7  
     8  	"go.uber.org/zap"
     9  )
    10  
    11  func TestConnection_FreezeAndUnfreeze(t *testing.T) {
    12  	in, out := net.Pipe()
    13  	defer func() { _ = in.Close() }()
    14  	defer func() { _ = out.Close() }()
    15  
    16  	cx := WrapConnection(out, []byte{}, zap.NewNop())
    17  	defer func() { _ = cx.Close() }()
    18  
    19  	matcherData := []byte("foo")
    20  	consumeData := []byte("bar")
    21  
    22  	buf := make([]byte, len(matcherData))
    23  
    24  	go func() {
    25  		_, _ = in.Write(matcherData)
    26  		_, _ = in.Write(consumeData)
    27  	}()
    28  
    29  	// prefetch like server handler would
    30  	err := cx.prefetch()
    31  	if err != nil {
    32  		t.Fatal(err)
    33  	}
    34  
    35  	// 1st matcher
    36  	cx.freeze()
    37  
    38  	n, err := cx.Read(buf)
    39  	if err != nil {
    40  		t.Fatal(err)
    41  	}
    42  	if n != len(matcherData) {
    43  		t.Fatalf("expected to read %d bytes but got %d", len(matcherData), n)
    44  	}
    45  	if bytes.Compare(matcherData, buf) != 0 {
    46  		t.Fatalf("expected %s but received %s", matcherData, buf)
    47  	}
    48  
    49  	cx.unfreeze()
    50  
    51  	// 2nd matcher (reads same data)
    52  
    53  	cx.freeze()
    54  
    55  	n, err = cx.Read(buf)
    56  	if err != nil {
    57  		t.Fatal(err)
    58  	}
    59  	if n != len(matcherData) {
    60  		t.Fatalf("expected to read %d bytes but got %d", len(matcherData), n)
    61  	}
    62  	if bytes.Compare(matcherData, buf) != 0 {
    63  		t.Fatalf("expected %s but received %s", matcherData, buf)
    64  	}
    65  
    66  	cx.unfreeze()
    67  
    68  	// 1st consumer (no freeze call)
    69  
    70  	n, err = cx.Read(buf)
    71  	if err != nil {
    72  		t.Fatal(err)
    73  	}
    74  	if n != len(matcherData) {
    75  		t.Fatalf("expected to read %d bytes but got %d", len(matcherData), n)
    76  	}
    77  	if bytes.Compare(matcherData, buf) != 0 {
    78  		t.Fatalf("expected %s but received %s", matcherData, buf)
    79  	}
    80  
    81  	// 2nd consumer (reads other data)
    82  
    83  	n, err = cx.Read(buf)
    84  	if err != nil {
    85  		t.Fatal(err)
    86  	}
    87  	if n != len(consumeData) {
    88  		t.Fatalf("expected to read %d bytes but got %d", len(consumeData), n)
    89  	}
    90  	if bytes.Compare(consumeData, buf) != 0 {
    91  		t.Fatalf("expected %s but received %s", consumeData, buf)
    92  	}
    93  }