github.com/SagerNet/gvisor@v0.0.0-20210707092255-7731c139d75c/pkg/tcpip/link/nested/nested_test.go (about)

     1  // Copyright 2020 The gVisor Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package nested_test
    16  
    17  import (
    18  	"testing"
    19  
    20  	"github.com/SagerNet/gvisor/pkg/tcpip"
    21  	"github.com/SagerNet/gvisor/pkg/tcpip/header"
    22  	"github.com/SagerNet/gvisor/pkg/tcpip/link/nested"
    23  	"github.com/SagerNet/gvisor/pkg/tcpip/stack"
    24  )
    25  
    26  type parentEndpoint struct {
    27  	nested.Endpoint
    28  }
    29  
    30  var _ stack.LinkEndpoint = (*parentEndpoint)(nil)
    31  var _ stack.NetworkDispatcher = (*parentEndpoint)(nil)
    32  
    33  type childEndpoint struct {
    34  	stack.LinkEndpoint
    35  	dispatcher stack.NetworkDispatcher
    36  }
    37  
    38  var _ stack.LinkEndpoint = (*childEndpoint)(nil)
    39  
    40  func (c *childEndpoint) Attach(dispatcher stack.NetworkDispatcher) {
    41  	c.dispatcher = dispatcher
    42  }
    43  
    44  func (c *childEndpoint) IsAttached() bool {
    45  	return c.dispatcher != nil
    46  }
    47  
    48  type counterDispatcher struct {
    49  	count int
    50  }
    51  
    52  var _ stack.NetworkDispatcher = (*counterDispatcher)(nil)
    53  
    54  func (d *counterDispatcher) DeliverNetworkPacket(tcpip.LinkAddress, tcpip.LinkAddress, tcpip.NetworkProtocolNumber, *stack.PacketBuffer) {
    55  	d.count++
    56  }
    57  
    58  func (d *counterDispatcher) DeliverOutboundPacket(tcpip.LinkAddress, tcpip.LinkAddress, tcpip.NetworkProtocolNumber, *stack.PacketBuffer) {
    59  	panic("unimplemented")
    60  }
    61  
    62  func TestNestedLinkEndpoint(t *testing.T) {
    63  	const emptyAddress = tcpip.LinkAddress("")
    64  
    65  	var (
    66  		childEP  childEndpoint
    67  		nestedEP parentEndpoint
    68  		disp     counterDispatcher
    69  	)
    70  	nestedEP.Endpoint.Init(&childEP, &nestedEP)
    71  
    72  	if childEP.IsAttached() {
    73  		t.Error("On init, childEP.IsAttached() = true, want = false")
    74  	}
    75  	if nestedEP.IsAttached() {
    76  		t.Error("On init, nestedEP.IsAttached() = true, want = false")
    77  	}
    78  
    79  	nestedEP.Attach(&disp)
    80  	if disp.count != 0 {
    81  		t.Fatalf("After attach, got disp.count = %d, want = 0", disp.count)
    82  	}
    83  	if !childEP.IsAttached() {
    84  		t.Error("After attach, childEP.IsAttached() = false, want = true")
    85  	}
    86  	if !nestedEP.IsAttached() {
    87  		t.Error("After attach, nestedEP.IsAttached() = false, want = true")
    88  	}
    89  
    90  	nestedEP.DeliverNetworkPacket(emptyAddress, emptyAddress, header.IPv4ProtocolNumber, stack.NewPacketBuffer(stack.PacketBufferOptions{}))
    91  	if disp.count != 1 {
    92  		t.Errorf("After first packet with dispatcher attached, got disp.count = %d, want = 1", disp.count)
    93  	}
    94  
    95  	nestedEP.Attach(nil)
    96  	if childEP.IsAttached() {
    97  		t.Error("After detach, childEP.IsAttached() = true, want = false")
    98  	}
    99  	if nestedEP.IsAttached() {
   100  		t.Error("After detach, nestedEP.IsAttached() = true, want = false")
   101  	}
   102  
   103  	disp.count = 0
   104  	nestedEP.DeliverNetworkPacket(emptyAddress, emptyAddress, header.IPv4ProtocolNumber, stack.NewPacketBuffer(stack.PacketBufferOptions{}))
   105  	if disp.count != 0 {
   106  		t.Errorf("After second packet with dispatcher detached, got disp.count = %d, want = 0", disp.count)
   107  	}
   108  
   109  }