moby/pkg/proxy/network_proxy_test.go
Michael Crosby 8c39db8f96 Move proxy into pkg
Docker-DCO-1.1-Signed-off-by: Michael Crosby <michael@crosbymichael.com> (github: crosbymichael)
2014-02-14 15:48:05 -08:00

216 lines
5.4 KiB
Go

package proxy
import (
"bytes"
"fmt"
"io"
"net"
"strings"
"testing"
"time"
)
var testBuf = []byte("Buffalo buffalo Buffalo buffalo buffalo buffalo Buffalo buffalo")
var testBufSize = len(testBuf)
type EchoServer interface {
Run()
Close()
LocalAddr() net.Addr
}
type TCPEchoServer struct {
listener net.Listener
testCtx *testing.T
}
type UDPEchoServer struct {
conn net.PacketConn
testCtx *testing.T
}
func NewEchoServer(t *testing.T, proto, address string) EchoServer {
var server EchoServer
if strings.HasPrefix(proto, "tcp") {
listener, err := net.Listen(proto, address)
if err != nil {
t.Fatal(err)
}
server = &TCPEchoServer{listener: listener, testCtx: t}
} else {
socket, err := net.ListenPacket(proto, address)
if err != nil {
t.Fatal(err)
}
server = &UDPEchoServer{conn: socket, testCtx: t}
}
return server
}
func (server *TCPEchoServer) Run() {
go func() {
for {
client, err := server.listener.Accept()
if err != nil {
return
}
go func(client net.Conn) {
if _, err := io.Copy(client, client); err != nil {
server.testCtx.Logf("can't echo to the client: %v\n", err.Error())
}
client.Close()
}(client)
}
}()
}
func (server *TCPEchoServer) LocalAddr() net.Addr { return server.listener.Addr() }
func (server *TCPEchoServer) Close() { server.listener.Addr() }
func (server *UDPEchoServer) Run() {
go func() {
readBuf := make([]byte, 1024)
for {
read, from, err := server.conn.ReadFrom(readBuf)
if err != nil {
return
}
for i := 0; i != read; {
written, err := server.conn.WriteTo(readBuf[i:read], from)
if err != nil {
break
}
i += written
}
}
}()
}
func (server *UDPEchoServer) LocalAddr() net.Addr { return server.conn.LocalAddr() }
func (server *UDPEchoServer) Close() { server.conn.Close() }
func testProxyAt(t *testing.T, proto string, proxy Proxy, addr string) {
defer proxy.Close()
go proxy.Run()
client, err := net.Dial(proto, addr)
if err != nil {
t.Fatalf("Can't connect to the proxy: %v", err)
}
defer client.Close()
client.SetDeadline(time.Now().Add(10 * time.Second))
if _, err = client.Write(testBuf); err != nil {
t.Fatal(err)
}
recvBuf := make([]byte, testBufSize)
if _, err = client.Read(recvBuf); err != nil {
t.Fatal(err)
}
if !bytes.Equal(testBuf, recvBuf) {
t.Fatal(fmt.Errorf("Expected [%v] but got [%v]", testBuf, recvBuf))
}
}
func testProxy(t *testing.T, proto string, proxy Proxy) {
testProxyAt(t, proto, proxy, proxy.FrontendAddr().String())
}
func TestTCP4Proxy(t *testing.T) {
backend := NewEchoServer(t, "tcp", "127.0.0.1:0")
defer backend.Close()
backend.Run()
frontendAddr := &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 0}
proxy, err := NewProxy(frontendAddr, backend.LocalAddr())
if err != nil {
t.Fatal(err)
}
testProxy(t, "tcp", proxy)
}
func TestTCP6Proxy(t *testing.T) {
backend := NewEchoServer(t, "tcp", "[::1]:0")
defer backend.Close()
backend.Run()
frontendAddr := &net.TCPAddr{IP: net.IPv6loopback, Port: 0}
proxy, err := NewProxy(frontendAddr, backend.LocalAddr())
if err != nil {
t.Fatal(err)
}
testProxy(t, "tcp", proxy)
}
func TestTCPDualStackProxy(t *testing.T) {
// If I understand `godoc -src net favoriteAddrFamily` (used by the
// net.Listen* functions) correctly this should work, but it doesn't.
t.Skip("No support for dual stack yet")
backend := NewEchoServer(t, "tcp", "[::1]:0")
defer backend.Close()
backend.Run()
frontendAddr := &net.TCPAddr{IP: net.IPv6loopback, Port: 0}
proxy, err := NewProxy(frontendAddr, backend.LocalAddr())
if err != nil {
t.Fatal(err)
}
ipv4ProxyAddr := &net.TCPAddr{
IP: net.IPv4(127, 0, 0, 1),
Port: proxy.FrontendAddr().(*net.TCPAddr).Port,
}
testProxyAt(t, "tcp", proxy, ipv4ProxyAddr.String())
}
func TestUDP4Proxy(t *testing.T) {
backend := NewEchoServer(t, "udp", "127.0.0.1:0")
defer backend.Close()
backend.Run()
frontendAddr := &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 0}
proxy, err := NewProxy(frontendAddr, backend.LocalAddr())
if err != nil {
t.Fatal(err)
}
testProxy(t, "udp", proxy)
}
func TestUDP6Proxy(t *testing.T) {
backend := NewEchoServer(t, "udp", "[::1]:0")
defer backend.Close()
backend.Run()
frontendAddr := &net.UDPAddr{IP: net.IPv6loopback, Port: 0}
proxy, err := NewProxy(frontendAddr, backend.LocalAddr())
if err != nil {
t.Fatal(err)
}
testProxy(t, "udp", proxy)
}
func TestUDPWriteError(t *testing.T) {
frontendAddr := &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 0}
// Hopefully, this port will be free: */
backendAddr := &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 25587}
proxy, err := NewProxy(frontendAddr, backendAddr)
if err != nil {
t.Fatal(err)
}
defer proxy.Close()
go proxy.Run()
client, err := net.Dial("udp", "127.0.0.1:25587")
if err != nil {
t.Fatalf("Can't connect to the proxy: %v", err)
}
defer client.Close()
// Make sure the proxy doesn't stop when there is no actual backend:
client.Write(testBuf)
client.Write(testBuf)
backend := NewEchoServer(t, "udp", "127.0.0.1:25587")
defer backend.Close()
backend.Run()
client.SetDeadline(time.Now().Add(10 * time.Second))
if _, err = client.Write(testBuf); err != nil {
t.Fatal(err)
}
recvBuf := make([]byte, testBufSize)
if _, err = client.Read(recvBuf); err != nil {
t.Fatal(err)
}
if !bytes.Equal(testBuf, recvBuf) {
t.Fatal(fmt.Errorf("Expected [%v] but got [%v]", testBuf, recvBuf))
}
}