proxy: remove assumption that frontend is a TCPListener

We require the frontend to be a `net.Listener` and the `net.Conn`
connection which are established must be type-switched to
`Conn` to support the `CloseRead` and `CloseWrite` methods.

Signed-off-by: David Scott <dave.scott@docker.com>
This commit is contained in:
David Scott 2016-04-14 16:02:02 +01:00
parent 2cf9295602
commit 3c68728e9f
2 changed files with 11 additions and 9 deletions

View File

@ -24,13 +24,19 @@ type Proxy interface {
BackendAddr() net.Addr
}
// NewProxy creates a Proxy according to the specified frontendAddr and backendAddr.
func NewProxy(frontendAddr, backendAddr net.Addr) (Proxy, error) {
switch frontendAddr.(type) {
case *net.UDPAddr:
return NewUDPProxy(frontendAddr.(*net.UDPAddr), backendAddr.(*net.UDPAddr))
case *net.TCPAddr:
return NewTCPProxy(frontendAddr.(*net.TCPAddr), backendAddr.(*net.TCPAddr))
listener, err := net.Listen("tcp", frontendAddr.String())
if err != nil {
return nil, err
}
return NewTCPProxy(listener, backendAddr.(*net.TCPAddr))
default:
panic(fmt.Errorf("Unsupported protocol"))
}

View File

@ -19,22 +19,18 @@ type Conn interface {
// TCPProxy is a proxy for TCP connections. It implements the Proxy interface to
// handle TCP traffic forwarding between the frontend and backend addresses.
type TCPProxy struct {
listener *net.TCPListener
frontendAddr *net.TCPAddr
listener net.Listener
frontendAddr net.Addr
backendAddr *net.TCPAddr
}
// NewTCPProxy creates a new TCPProxy.
func NewTCPProxy(frontendAddr, backendAddr *net.TCPAddr) (*TCPProxy, error) {
listener, err := net.ListenTCP("tcp", frontendAddr)
if err != nil {
return nil, err
}
func NewTCPProxy(listener net.Listener, backendAddr *net.TCPAddr) (*TCPProxy, error) {
// If the port in frontendAddr was 0 then ListenTCP will have a picked
// a port to listen on, hence the call to Addr to get that actual port:
return &TCPProxy{
listener: listener,
frontendAddr: listener.Addr().(*net.TCPAddr),
frontendAddr: listener.Addr(),
backendAddr: backendAddr,
}, nil
}