Merge pull request #3375 from kevinGC:ipt-test-early-return

PiperOrigin-RevId: 326693922

pkg/test/testutil/testutil.go

@@ -316,6 +316,11 @@ func Copy(src, dst string) error {
 func Poll(cb func() error, timeout time.Duration) error {
 	ctx, cancel := context.WithTimeout(context.Background(), timeout)
 	defer cancel()
+	return PollContext(ctx, cb)
+}
+
+// PollContext is like Poll, but takes a context instead of a timeout.
+func PollContext(ctx context.Context, cb func() error) error {
 	b := backoff.WithContext(backoff.NewConstantBackOff(100*time.Millisecond), ctx)
 	return backoff.Retry(cb, b)
 }

test/iptables/iptables.go

@@ -16,6 +16,7 @@
 package iptables
 
 import (
+	"context"
 	"fmt"
 	"net"
 	"time"
@@ -29,7 +30,11 @@ const IPExchangePort = 2349
 const TerminalStatement = "Finished!"
 
 // TestTimeout is the timeout used for all tests.
-const TestTimeout = 10 * time.Minute
+const TestTimeout = 10 * time.Second
+
+// NegativeTimeout is the time tests should wait to establish the negative
+// case, i.e. that connections are not made.
+const NegativeTimeout = 2 * time.Second
 
 // A TestCase contains one action to run in the container and one to run
 // locally. The actions run concurrently and each must succeed for the test
@@ -40,10 +45,60 @@ type TestCase interface {
 
 	// ContainerAction runs inside the container. It receives the IP of the
 	// local process.
-	ContainerAction(ip net.IP, ipv6 bool) error
+	ContainerAction(ctx context.Context, ip net.IP, ipv6 bool) error
 
 	// LocalAction runs locally. It receives the IP of the container.
-	LocalAction(ip net.IP, ipv6 bool) error
+	LocalAction(ctx context.Context, ip net.IP, ipv6 bool) error
+
+	// ContainerSufficient indicates whether ContainerAction's return value
+	// alone indicates whether the test succeeded.
+	ContainerSufficient() bool
+
+	// LocalSufficient indicates whether LocalAction's return value alone
+	// indicates whether the test succeeded.
+	LocalSufficient() bool
+}
+
+// baseCase provides defaults for ContainerSufficient and LocalSufficient when
+// both actions are required to finish.
+type baseCase struct{}
+
+// ContainerSufficient implements TestCase.ContainerSufficient.
+func (baseCase) ContainerSufficient() bool {
+	return false
+}
+
+// LocalSufficient implements TestCase.LocalSufficient.
+func (baseCase) LocalSufficient() bool {
+	return false
+}
+
+// localCase provides defaults for ContainerSufficient and LocalSufficient when
+// only the local action is required to finish.
+type localCase struct{}
+
+// ContainerSufficient implements TestCase.ContainerSufficient.
+func (localCase) ContainerSufficient() bool {
+	return false
+}
+
+// LocalSufficient implements TestCase.LocalSufficient.
+func (localCase) LocalSufficient() bool {
+	return true
+}
+
+// containerCase provides defaults for ContainerSufficient and LocalSufficient
+// when only the container action is required to finish.
+type containerCase struct{}
+
+// ContainerSufficient implements TestCase.ContainerSufficient.
+func (containerCase) ContainerSufficient() bool {
+	return true
+}
+
+// LocalSufficient implements TestCase.LocalSufficient.
+func (containerCase) LocalSufficient() bool {
+	return false
 }
 
 // Tests maps test names to TestCase.

test/iptables/iptables_test.go

@@ -16,9 +16,11 @@ package iptables
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"net"
 	"reflect"
+	"sync"
 	"testing"
 
 	"gvisor.dev/gvisor/pkg/test/dockerutil"
@@ -51,9 +53,24 @@ func iptablesTest(t *testing.T, test TestCase, ipv6 bool) {
 		t.Fatalf("no test found with name %q. Has it been registered?", test.Name())
 	}
 
-	ctx := context.Background()
+	// Wait for the local and container goroutines to finish.
+	var wg sync.WaitGroup
+	defer wg.Wait()
+
+	ctx, cancel := context.WithTimeout(context.Background(), TestTimeout)
+	defer cancel()
+
 	d := dockerutil.MakeContainer(ctx, t)
-	defer d.CleanUp(ctx)
+	defer func() {
+		if logs, err := d.Logs(context.Background()); err != nil {
+			t.Logf("Failed to retrieve container logs.")
+		} else {
+			t.Logf("=== Container logs: ===\n%s", logs)
+		}
+		// Use a new context, as cleanup should run even when we
+		// timeout.
+		d.CleanUp(context.Background())
+	}()
 
 	// TODO(gvisor.dev/issue/170): Skipping IPv6 gVisor tests.
 	if ipv6 && dockerutil.Runtime() != "runc" {
@@ -86,15 +103,44 @@ func iptablesTest(t *testing.T, test TestCase, ipv6 bool) {
 	}
 
 	// Run our side of the test.
-	if err := test.LocalAction(ip, ipv6); err != nil {
-		t.Fatalf("LocalAction failed: %v", err)
-	}
+	errCh := make(chan error, 2)
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		if err := test.LocalAction(ctx, ip, ipv6); err != nil && !errors.Is(err, context.Canceled) {
+			errCh <- fmt.Errorf("LocalAction failed: %v", err)
+		} else {
+			errCh <- nil
+		}
+		if test.LocalSufficient() {
+			errCh <- nil
+		}
+	}()
 
-	// Wait for the final statement. This structure has the side effect
-	// that all container logs will appear within the individual test
-	// context.
-	if _, err := d.WaitForOutput(ctx, TerminalStatement, TestTimeout); err != nil {
-		t.Fatalf("test failed: %v", err)
+	// Run the container side.
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		// Wait for the final statement. This structure has the side
+		// effect that all container logs will appear within the
+		// individual test context.
+		if _, err := d.WaitForOutput(ctx, TerminalStatement, TestTimeout); err != nil && !errors.Is(err, context.Canceled) {
+			errCh <- fmt.Errorf("ContainerAction failed: %v", err)
+		} else {
+			errCh <- nil
+		}
+		if test.ContainerSufficient() {
+			errCh <- nil
+		}
+	}()
+
+	for i := 0; i < 2; i++ {
+		select {
+		case err := <-errCh:
+			if err != nil {
+				t.Fatal(err)
+			}
+		}
 	}
 }
 

test/iptables/iptables_util.go

@@ -15,6 +15,7 @@
 package iptables
 
 import (
+	"context"
 	"encoding/binary"
 	"errors"
 	"fmt"
@@ -70,7 +71,7 @@ func tableRules(ipv6 bool, table string, argsList [][]string) error {
 
 // listenUDP listens on a UDP port and returns the value of net.Conn.Read() for
 // the first read on that port.
-func listenUDP(port int, timeout time.Duration) error {
+func listenUDP(ctx context.Context, port int) error {
 	localAddr := net.UDPAddr{
 		Port: port,
 	}
@@ -79,68 +80,53 @@ func listenUDP(port int, timeout time.Duration) error {
 		return err
 	}
 	defer conn.Close()
-	conn.SetDeadline(time.Now().Add(timeout))
-	_, err = conn.Read([]byte{0})
-	return err
+
+	ch := make(chan error)
+	go func() {
+		_, err = conn.Read([]byte{0})
+		ch <- err
+	}()
+
+	select {
+	case err := <-ch:
+		return err
+	case <-ctx.Done():
+		return ctx.Err()
+	}
 }
 
 // sendUDPLoop sends 1 byte UDP packets repeatedly to the IP and port specified
 // over a duration.
-func sendUDPLoop(ip net.IP, port int, duration time.Duration) error {
-	conn, err := connectUDP(ip, port)
-	if err != nil {
-		return err
-	}
-	defer conn.Close()
-	loopUDP(conn, duration)
-	return nil
-}
-
-// spawnUDPLoop works like sendUDPLoop, but returns immediately and sends
-// packets in another goroutine.
-func spawnUDPLoop(ip net.IP, port int, duration time.Duration) error {
-	conn, err := connectUDP(ip, port)
-	if err != nil {
-		return err
-	}
-	go func() {
-		defer conn.Close()
-		loopUDP(conn, duration)
-	}()
-	return nil
-}
-
-func connectUDP(ip net.IP, port int) (net.Conn, error) {
+func sendUDPLoop(ctx context.Context, ip net.IP, port int) error {
 	remote := net.UDPAddr{
 		IP:   ip,
 		Port: port,
 	}
 	conn, err := net.DialUDP("udp", nil, &remote)
 	if err != nil {
-		return nil, err
+		return err
 	}
-	return conn, nil
-}
+	defer conn.Close()
 
-func loopUDP(conn net.Conn, duration time.Duration) {
-	to := time.After(duration)
-	for timedOut := false; !timedOut; {
+	for {
 		// This may return an error (connection refused) if the remote
 		// hasn't started listening yet or they're dropping our
 		// packets. So we ignore Write errors and depend on the remote
 		// to report a failure if it doesn't get a packet it needs.
 		conn.Write([]byte{0})
 		select {
-		case <-to:
-			timedOut = true
-		default:
-			time.Sleep(200 * time.Millisecond)
+		case <-ctx.Done():
+			// Being cancelled or timing out isn't an error, as we
+			// cannot tell with UDP whether we succeeded.
+			return nil
+		// Continue looping.
+		case <-time.After(200 * time.Millisecond):
 		}
 	}
 }
 
 // listenTCP listens for connections on a TCP port.
-func listenTCP(port int, timeout time.Duration) error {
+func listenTCP(ctx context.Context, port int) error {
 	localAddr := net.TCPAddr{
 		Port: port,
 	}
@@ -153,17 +139,23 @@ func listenTCP(port int, timeout time.Duration) error {
 	defer lConn.Close()
 
 	// Accept connections on port.
-	lConn.SetDeadline(time.Now().Add(timeout))
-	conn, err := lConn.AcceptTCP()
-	if err != nil {
+	ch := make(chan error)
+	go func() {
+		conn, err := lConn.AcceptTCP()
+		ch <- err
+		conn.Close()
+	}()
+
+	select {
+	case err := <-ch:
 		return err
+	case <-ctx.Done():
+		return fmt.Errorf("timed out waiting for a connection at %#v: %w", localAddr, ctx.Err())
 	}
-	conn.Close()
-	return nil
 }
 
 // connectTCP connects to the given IP and port from an ephemeral local address.
-func connectTCP(ip net.IP, port int, timeout time.Duration) error {
+func connectTCP(ctx context.Context, ip net.IP, port int) error {
 	contAddr := net.TCPAddr{
 		IP:   ip,
 		Port: port,
@@ -171,13 +163,14 @@ func connectTCP(ip net.IP, port int, timeout time.Duration) error {
 	// The container may not be listening when we first connect, so retry
 	// upon error.
 	callback := func() error {
-		conn, err := net.DialTimeout("tcp", contAddr.String(), timeout)
+		var d net.Dialer
+		conn, err := d.DialContext(ctx, "tcp", contAddr.String())
 		if conn != nil {
 			conn.Close()
 		}
 		return err
 	}
-	if err := testutil.Poll(callback, timeout); err != nil {
+	if err := testutil.PollContext(ctx, callback); err != nil {
 		return fmt.Errorf("timed out waiting to connect IP on port %v, most recent error: %v", port, err)
 	}
 

test/iptables/runner/main.go

@@ -16,6 +16,7 @@
 package main
 
 import (
+	"context"
 	"flag"
 	"fmt"
 	"log"
@@ -46,7 +47,9 @@ func main() {
 	}
 
 	// Run the test.
-	if err := test.ContainerAction(ip, *ipv6); err != nil {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	if err := test.ContainerAction(ctx, ip, *ipv6); err != nil {
 		log.Fatalf("Failed running test %q: %v", *name, err)
 	}