testing: add TCPProxy

Author: rogpeppe

mgo_unix.go

@@ -13,4 +13,3 @@ import "os"
 func (inst *MgoInstance) DestroyWithLog() {
 	inst.killAndCleanup(os.Interrupt)
 }
-

mgo_windows.go

@@ -16,4 +16,3 @@ import "os"
 func (inst *MgoInstance) DestroyWithLog() {
 	inst.killAndCleanup(os.Kill)
 }
-

osenv.go

@@ -67,25 +67,25 @@ var testingVariables = []string{
 }
 
 func (s *OsEnvSuite) setEnviron() {
-	var isWhitelisted func (string) bool
+	var isWhitelisted func(string) bool
 	switch runtime.GOOS {
 	case "windows":
 		// Lowercase variable names for comparison as they are case
 		// insenstive on windows. Fancy folding not required for ascii.
 		lowerEnv := make(map[string]struct{},
-			len(windowsVariables) + len(testingVariables))
+			len(windowsVariables)+len(testingVariables))
 		for _, envVar := range windowsVariables {
 			lowerEnv[strings.ToLower(envVar)] = struct{}{}
 		}
 		for _, envVar := range testingVariables {
 			lowerEnv[strings.ToLower(envVar)] = struct{}{}
 		}
-		isWhitelisted = func (envVar string) bool {
+		isWhitelisted = func(envVar string) bool {
 			_, ok := lowerEnv[strings.ToLower(envVar)]
 			return ok
 		}
 	default:
-		isWhitelisted = func (envVar string) bool {
+		isWhitelisted = func(envVar string) bool {
 			for _, testingVar := range testingVariables {
 				if testingVar == envVar {
 					return true

tcpproxy.go

@@ -0,0 +1,104 @@
+// Copyright 2015 Canonical Ltd.
+// Licensed under the LGPLv3, see LICENCE file for details.
+
+package testing
+
+import (
+	"io"
+	"net"
+	"sync"
+
+	jc "github.com/juju/testing/checkers"
+	gc "gopkg.in/check.v1"
+)
+
+// TCPProxy is a simple TCP proxy that can be used
+// to deliberately break TCP connections.
+type TCPProxy struct {
+	listener net.Listener
+	// mu guards the fields below it.
+	mu sync.Mutex
+	// closed holds whether the proxy has been closed.
+	closed bool
+	// conns holds all connections that have been made.
+	conns []io.Closer
+}
+
+// NewTCPProxy runs a proxy that copies to and from
+// the given remote TCP address. When the proxy
+// is closed, its listener and all connections will be closed.
+func NewTCPProxy(c *gc.C, remoteAddr string) *TCPProxy {
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
+	c.Assert(err, jc.ErrorIsNil)
+	p := &TCPProxy{
+		listener: listener,
+	}
+	go func() {
+		for {
+			client, err := p.listener.Accept()
+			if err != nil {
+				if !p.isClosed() {
+					c.Error("cannot accept: %v", err)
+				}
+				return
+			}
+			p.addConn(client)
+			server, err := net.Dial("tcp", remoteAddr)
+			if err != nil {
+				if !p.isClosed() {
+					c.Error("cannot dial remote address: %v", err)
+				}
+				return
+			}
+			p.addConn(server)
+			go stream(client, server)
+			go stream(server, client)
+		}
+	}()
+	return p
+}
+
+func (p *TCPProxy) addConn(c net.Conn) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	if p.closed {
+		c.Close()
+	} else {
+		p.conns = append(p.conns, c)
+	}
+}
+
+// Close closes the TCPProxy and any connections that
+// are currently active.
+func (p *TCPProxy) Close() error {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	p.closed = true
+	p.listener.Close()
+	for _, c := range p.conns {
+		c.Close()
+	}
+	return nil
+}
+
+// Addr returns the TCP address of the proxy. Dialing
+// this address will cause a connection to be made
+// to the remote address; any data written will be
+// written there, and any data read from the remote
+// address will be available to read locally.
+func (p *TCPProxy) Addr() string {
+	// Note: this only works because we explicitly listen on 127.0.0.1 rather
+	// than the wildcard address.
+	return p.listener.Addr().String()
+}
+func (p *TCPProxy) isClosed() bool {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	return p.closed
+}
+
+func stream(dst io.WriteCloser, src io.ReadCloser) {
+	defer dst.Close()
+	defer src.Close()
+	io.Copy(dst, src)
+}

tcpproxy_test.go

@@ -0,0 +1,89 @@
+// Copyright 2015 Canonical Ltd.
+// Licensed under the LGPLv3, see LICENCE file for details.
+
+package testing_test
+
+import (
+	"fmt"
+	"io"
+	"net"
+	"sync"
+
+	"github.com/juju/testing"
+	gc "gopkg.in/check.v1"
+)
+
+var _ = gc.Suite(&tcpProxySuite{})
+
+type tcpProxySuite struct{}
+
+func (*tcpProxySuite) TestTCPProxy(c *gc.C) {
+	var wg sync.WaitGroup
+
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
+	c.Assert(err, gc.IsNil)
+	defer listener.Close()
+	wg.Add(1)
+	go tcpEcho(&wg, listener)
+
+	p := testing.NewTCPProxy(c, listener.Addr().String())
+	c.Assert(p.Addr(), gc.Not(gc.Equals), listener.Addr().String())
+
+	// Dial the proxy and check that we see the text echoed correctly.
+	conn, err := net.Dial("tcp", p.Addr())
+	c.Assert(err, gc.IsNil)
+	defer conn.Close()
+	txt := "hello, world\n"
+	fmt.Fprint(conn, txt)
+
+	buf := make([]byte, len(txt))
+	n, err := io.ReadFull(conn, buf)
+	c.Assert(err, gc.IsNil)
+	c.Assert(string(buf[0:n]), gc.Equals, txt)
+
+	// Close the connection and check that we see
+	// the connection closed for read.
+	conn.(*net.TCPConn).CloseWrite()
+	n, err = conn.Read(buf)
+	c.Assert(err, gc.Equals, io.EOF)
+	c.Assert(n, gc.Equals, 0)
+
+	// Make another connection and close the proxy,
+	// which should close down the proxy and cause us
+	// to get an error.
+	conn, err = net.Dial("tcp", p.Addr())
+	c.Assert(err, gc.IsNil)
+	defer conn.Close()
+
+	p.Close()
+	_, err = conn.Read(buf)
+	c.Assert(err, gc.Equals, io.EOF)
+
+	// Make sure that we cannot dial the proxy address either.
+	conn, err = net.Dial("tcp", p.Addr())
+	c.Assert(err, gc.ErrorMatches, ".*connection refused")
+
+	listener.Close()
+	// Make sure that all our connections have gone away too.
+	wg.Wait()
+}
+
+// tcpEcho listens on the given listener for TCP connections,
+// writes all traffic received back to the sender, and calls
+// wg.Done when all its goroutines have completed.
+func tcpEcho(wg *sync.WaitGroup, listener net.Listener) {
+	defer wg.Done()
+	for {
+		conn, err := listener.Accept()
+		if err != nil {
+			return
+		}
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			defer conn.Close()
+			// Echo anything that was written.
+			io.Copy(conn, conn)
+		}()
+	}
+}