github.com/SagerNet/gvisor@v0.0.0-20210707092255-7731c139d75c/pkg/tcpip/network/internal/testutil/testutil.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 testutil defines types and functions used to test Network Layer
    16  // functionality such as IP fragmentation.
    17  package testutil
    18  
    19  import (
    20  	"fmt"
    21  	"math/rand"
    22  
    23  	"github.com/SagerNet/gvisor/pkg/tcpip"
    24  	"github.com/SagerNet/gvisor/pkg/tcpip/buffer"
    25  	"github.com/SagerNet/gvisor/pkg/tcpip/header"
    26  	"github.com/SagerNet/gvisor/pkg/tcpip/stack"
    27  )
    28  
    29  // MockLinkEndpoint is an endpoint used for testing, it stores packets written
    30  // to it and can mock errors.
    31  type MockLinkEndpoint struct {
    32  	// WrittenPackets is where packets written to the endpoint are stored.
    33  	WrittenPackets []*stack.PacketBuffer
    34  
    35  	mtu          uint32
    36  	err          tcpip.Error
    37  	allowPackets int
    38  }
    39  
    40  // NewMockLinkEndpoint creates a new MockLinkEndpoint.
    41  //
    42  // err is the error that will be returned once allowPackets packets are written
    43  // to the endpoint.
    44  func NewMockLinkEndpoint(mtu uint32, err tcpip.Error, allowPackets int) *MockLinkEndpoint {
    45  	return &MockLinkEndpoint{
    46  		mtu:          mtu,
    47  		err:          err,
    48  		allowPackets: allowPackets,
    49  	}
    50  }
    51  
    52  // MTU implements LinkEndpoint.MTU.
    53  func (ep *MockLinkEndpoint) MTU() uint32 { return ep.mtu }
    54  
    55  // Capabilities implements LinkEndpoint.Capabilities.
    56  func (*MockLinkEndpoint) Capabilities() stack.LinkEndpointCapabilities { return 0 }
    57  
    58  // MaxHeaderLength implements LinkEndpoint.MaxHeaderLength.
    59  func (*MockLinkEndpoint) MaxHeaderLength() uint16 { return 0 }
    60  
    61  // LinkAddress implements LinkEndpoint.LinkAddress.
    62  func (*MockLinkEndpoint) LinkAddress() tcpip.LinkAddress { return "" }
    63  
    64  // WritePacket implements LinkEndpoint.WritePacket.
    65  func (ep *MockLinkEndpoint) WritePacket(_ stack.RouteInfo, _ tcpip.NetworkProtocolNumber, pkt *stack.PacketBuffer) tcpip.Error {
    66  	if ep.allowPackets == 0 {
    67  		return ep.err
    68  	}
    69  	ep.allowPackets--
    70  	ep.WrittenPackets = append(ep.WrittenPackets, pkt)
    71  	return nil
    72  }
    73  
    74  // WritePackets implements LinkEndpoint.WritePackets.
    75  func (ep *MockLinkEndpoint) WritePackets(r stack.RouteInfo, pkts stack.PacketBufferList, protocol tcpip.NetworkProtocolNumber) (int, tcpip.Error) {
    76  	var n int
    77  
    78  	for pkt := pkts.Front(); pkt != nil; pkt = pkt.Next() {
    79  		if err := ep.WritePacket(r, protocol, pkt); err != nil {
    80  			return n, err
    81  		}
    82  		n++
    83  	}
    84  
    85  	return n, nil
    86  }
    87  
    88  // Attach implements LinkEndpoint.Attach.
    89  func (*MockLinkEndpoint) Attach(stack.NetworkDispatcher) {}
    90  
    91  // IsAttached implements LinkEndpoint.IsAttached.
    92  func (*MockLinkEndpoint) IsAttached() bool { return false }
    93  
    94  // Wait implements LinkEndpoint.Wait.
    95  func (*MockLinkEndpoint) Wait() {}
    96  
    97  // ARPHardwareType implements LinkEndpoint.ARPHardwareType.
    98  func (*MockLinkEndpoint) ARPHardwareType() header.ARPHardwareType { return header.ARPHardwareNone }
    99  
   100  // AddHeader implements LinkEndpoint.AddHeader.
   101  func (*MockLinkEndpoint) AddHeader(_, _ tcpip.LinkAddress, _ tcpip.NetworkProtocolNumber, _ *stack.PacketBuffer) {
   102  }
   103  
   104  // MakeRandPkt generates a randomized packet. transportHeaderLength indicates
   105  // how many random bytes will be copied in the Transport Header.
   106  // extraHeaderReserveLength indicates how much extra space will be reserved for
   107  // the other headers. The payload is made from Views of the sizes listed in
   108  // viewSizes.
   109  func MakeRandPkt(transportHeaderLength int, extraHeaderReserveLength int, viewSizes []int, proto tcpip.NetworkProtocolNumber) *stack.PacketBuffer {
   110  	var views buffer.VectorisedView
   111  
   112  	for _, s := range viewSizes {
   113  		newView := buffer.NewView(s)
   114  		if _, err := rand.Read(newView); err != nil {
   115  			panic(fmt.Sprintf("rand.Read: %s", err))
   116  		}
   117  		views.AppendView(newView)
   118  	}
   119  
   120  	pkt := stack.NewPacketBuffer(stack.PacketBufferOptions{
   121  		ReserveHeaderBytes: transportHeaderLength + extraHeaderReserveLength,
   122  		Data:               views,
   123  	})
   124  	pkt.NetworkProtocolNumber = proto
   125  	if _, err := rand.Read(pkt.TransportHeader().Push(transportHeaderLength)); err != nil {
   126  		panic(fmt.Sprintf("rand.Read: %s", err))
   127  	}
   128  	return pkt
   129  }