diff options
Diffstat (limited to 'test/packetimpact/testbench/connections.go')
| -rw-r--r-- | test/packetimpact/testbench/connections.go | 37 |
1 files changed, 21 insertions, 16 deletions
diff --git a/test/packetimpact/testbench/connections.go b/test/packetimpact/testbench/connections.go index 56ac3fa54..463fd0556 100644 --- a/test/packetimpact/testbench/connections.go +++ b/test/packetimpact/testbench/connections.go @@ -17,7 +17,6 @@ package testbench import ( - "flag" "fmt" "math/rand" "net" @@ -32,13 +31,6 @@ import ( "gvisor.dev/gvisor/pkg/tcpip/seqnum" ) -var localIPv4 = flag.String("local_ipv4", "", "local IPv4 address for test packets") -var remoteIPv4 = flag.String("remote_ipv4", "", "remote IPv4 address for test packets") -var localIPv6 = flag.String("local_ipv6", "", "local IPv6 address for test packets") -var remoteIPv6 = flag.String("remote_ipv6", "", "remote IPv6 address for test packets") -var localMAC = flag.String("local_mac", "", "local mac address for test packets") -var remoteMAC = flag.String("remote_mac", "", "remote mac address for test packets") - func portFromSockaddr(sa unix.Sockaddr) (uint16, error) { switch sa := sa.(type) { case *unix.SockaddrInet4: @@ -64,11 +56,11 @@ func pickPort(domain, typ int) (fd int, sa unix.Sockaddr, err error) { switch domain { case unix.AF_INET: var sa4 unix.SockaddrInet4 - copy(sa4.Addr[:], net.ParseIP(*localIPv4).To4()) + copy(sa4.Addr[:], net.ParseIP(LocalIPv4).To4()) sa = &sa4 case unix.AF_INET6: var sa6 unix.SockaddrInet6 - copy(sa6.Addr[:], net.ParseIP(*localIPv6).To16()) + copy(sa6.Addr[:], net.ParseIP(LocalIPv6).To16()) sa = &sa6 default: return -1, nil, fmt.Errorf("invalid domain %d, it should be one of unix.AF_INET or unix.AF_INET6", domain) @@ -120,12 +112,12 @@ var _ layerState = (*etherState)(nil) // newEtherState creates a new etherState. func newEtherState(out, in Ether) (*etherState, error) { - lMAC, err := tcpip.ParseMACAddress(*localMAC) + lMAC, err := tcpip.ParseMACAddress(LocalMAC) if err != nil { return nil, err } - rMAC, err := tcpip.ParseMACAddress(*remoteMAC) + rMAC, err := tcpip.ParseMACAddress(RemoteMAC) if err != nil { return nil, err } @@ -172,8 +164,8 @@ var _ layerState = (*ipv4State)(nil) // newIPv4State creates a new ipv4State. func newIPv4State(out, in IPv4) (*ipv4State, error) { - lIP := tcpip.Address(net.ParseIP(*localIPv4).To4()) - rIP := tcpip.Address(net.ParseIP(*remoteIPv4).To4()) + lIP := tcpip.Address(net.ParseIP(LocalIPv4).To4()) + rIP := tcpip.Address(net.ParseIP(RemoteIPv4).To4()) s := ipv4State{ out: IPv4{SrcAddr: &lIP, DstAddr: &rIP}, in: IPv4{SrcAddr: &rIP, DstAddr: &lIP}, @@ -217,8 +209,8 @@ var _ layerState = (*ipv6State)(nil) // newIPv6State creates a new ipv6State. func newIPv6State(out, in IPv6) (*ipv6State, error) { - lIP := tcpip.Address(net.ParseIP(*localIPv6).To16()) - rIP := tcpip.Address(net.ParseIP(*remoteIPv6).To16()) + lIP := tcpip.Address(net.ParseIP(LocalIPv6).To16()) + rIP := tcpip.Address(net.ParseIP(RemoteIPv6).To16()) s := ipv6State{ out: IPv6{SrcAddr: &lIP, DstAddr: &rIP}, in: IPv6{SrcAddr: &rIP, DstAddr: &lIP}, @@ -841,6 +833,7 @@ func (conn *UDPIPv4) SendIP(additionalLayers ...Layer) { // Expect expects a frame with the UDP layer matching the provided UDP within // the timeout specified. If it doesn't arrive in time, an error is returned. func (conn *UDPIPv4) Expect(udp UDP, timeout time.Duration) (*UDP, error) { + conn.t.Helper() layer, err := (*Connection)(conn).Expect(&udp, timeout) if layer == nil { return nil, err @@ -852,6 +845,18 @@ func (conn *UDPIPv4) Expect(udp UDP, timeout time.Duration) (*UDP, error) { return gotUDP, err } +// ExpectData is a convenient method that expects a Layer and the Layer after +// it. If it doens't arrive in time, it returns nil. +func (conn *UDPIPv4) ExpectData(udp UDP, payload Payload, timeout time.Duration) (Layers, error) { + conn.t.Helper() + expected := make([]Layer, len(conn.layerStates)) + expected[len(expected)-1] = &udp + if payload.length() != 0 { + expected = append(expected, &payload) + } + return (*Connection)(conn).ExpectFrame(expected, timeout) +} + // Close frees associated resources held by the UDPIPv4 connection. func (conn *UDPIPv4) Close() { (*Connection)(conn).Close() |