Merge cdc20e9 into 4d51b51

README.md

@@ -25,7 +25,46 @@ $ go get -u github.com/pires/go-proxyproto
 
 ### Client (TODO)
 
-### Server (TODO)
+### Server
+
+```go
+package main
+
+import (
+	"log"
+	"net"
+
+	proxyproto "github.com/pires/go-proxyproto"
+)
+
+func main() {
+	// Create a listener
+	addr := "localhost:9876"
+	list, err := net.Listen("tcp", addr)
+	if err != nil {
+		log.Fatalf("couldn't listen to %q: %q\n", addr, err.Error())
+	}
+
+	// Wrap listener in a proxyproto listener
+	proxyListener := &proxyproto.Listener{Listener: list}
+	defer proxyListener.Close()
+
+	// Wait for a connection and accept it
+	conn, err := proxyListener.Accept()
+	defer conn.Close()
+
+	// Print connection details
+	if conn.LocalAddr() == nil {
+		log.Fatal("couldn't retrieve local address")
+	}
+	log.Printf("local address: %q", conn.LocalAddr().String())
+
+	if conn.RemoteAddr() == nil {
+		log.Fatal("couldn't retrieve remote address")
+	}
+	log.Printf("remote address: %q", conn.RemoteAddr().String())
+}
+```
 
 ## Documentation
 

header.go

@@ -41,6 +41,20 @@ type Header struct {
 	DestinationPort    uint16
 }
 
+func (header *Header) RemoteAddr() net.Addr {
+	return &net.TCPAddr{
+		IP:   header.SourceAddress,
+		Port: int(header.SourcePort),
+	}
+}
+
+func (header *Header) LocalAddr() net.Addr {
+	return &net.TCPAddr{
+		IP:   header.DestinationAddress,
+		Port: int(header.DestinationPort),
+	}
+}
+
 // EqualTo returns true if headers are equivalent, false otherwise.
 func (header *Header) EqualTo(q *Header) bool {
 	if header == nil || q == nil {
@@ -58,13 +72,23 @@ func (header *Header) EqualTo(q *Header) bool {
 
 // WriteTo renders a proxy protocol header in a format to write over the wire.
 func (header *Header) WriteTo(w io.Writer) (int64, error) {
+	buf, err := header.Format()
+	if err != nil {
+		return 0, err
+	}
+
+	return bytes.NewBuffer(buf).WriteTo(w)
+}
+
+// WriteTo renders a proxy protocol header in a format to write over the wire.
+func (header *Header) Format() ([]byte, error) {
 	switch header.Version {
 	case 1:
-		return header.writeVersion1(w)
+		return header.formatVersion1()
 	case 2:
-		return header.writeVersion2(w)
+		return header.formatVersion2()
 	default:
-		return 0, ErrUnknownProxyProtocolVersion
+		return nil, ErrUnknownProxyProtocolVersion
 	}
 }
 

protocol.go

@@ -0,0 +1,123 @@
+package proxyproto
+
+import (
+	"bufio"
+	"net"
+	"sync"
+	"time"
+)
+
+// Listener is used to wrap an underlying listener,
+// whose connections may be using the HAProxy Proxy Protocol.
+// If the connection is using the protocol, the RemoteAddr() will return
+// the correct client address.
+//
+// Optionally define ProxyHeaderTimeout to set a maximum time to
+// receive the Proxy Protocol Header. Zero means no timeout.
+type Listener struct {
+	Listener           net.Listener
+	ProxyHeaderTimeout time.Duration
+}
+
+// Conn is used to wrap and underlying connection which
+// may be speaking the Proxy Protocol. If it is, the RemoteAddr() will
+// return the address of the client instead of the proxy address.
+type Conn struct {
+	bufReader          *bufio.Reader
+	conn               net.Conn
+	header             *Header
+	once               sync.Once
+	proxyHeaderTimeout time.Duration
+}
+
+// Accept waits for and returns the next connection to the listener.
+func (p *Listener) Accept() (net.Conn, error) {
+	// Get the underlying connection
+	conn, err := p.Listener.Accept()
+	if err != nil {
+		return nil, err
+	}
+	return NewConn(conn, p.ProxyHeaderTimeout), nil
+}
+
+// Close closes the underlying listener.
+func (p *Listener) Close() error {
+	return p.Listener.Close()
+}
+
+// Addr returns the underlying listener's network address.
+func (p *Listener) Addr() net.Addr {
+	return p.Listener.Addr()
+}
+
+// NewConn is used to wrap a net.Conn that may be speaking
+// the proxy protocol into a proxyproto.Conn
+func NewConn(conn net.Conn, timeout time.Duration) *Conn {
+	pConn := &Conn{
+		bufReader:          bufio.NewReader(conn),
+		conn:               conn,
+		proxyHeaderTimeout: timeout,
+	}
+	return pConn
+}
+
+// Read is check for the proxy protocol header when doing
+// the initial scan. If there is an error parsing the header,
+// it is returned and the socket is closed.
+func (p *Conn) Read(b []byte) (int, error) {
+	var err error
+	p.once.Do(func() {
+		err = p.readHeader()
+	})
+	if err != nil {
+		return 0, err
+	}
+	return p.bufReader.Read(b)
+}
+
+func (p *Conn) Write(b []byte) (int, error) {
+	return p.conn.Write(b)
+}
+
+func (p *Conn) Close() error {
+	return p.conn.Close()
+}
+
+func (p *Conn) LocalAddr() net.Addr {
+	p.once.Do(func() { p.readHeader() })
+	if p.header == nil {
+		return p.conn.LocalAddr()
+	}
+
+	return p.header.LocalAddr()
+}
+
+// RemoteAddr returns the address of the client if the proxy
+// protocol is being used, otherwise just returns the address of
+// the socket peer.
+func (p *Conn) RemoteAddr() net.Addr {
+	p.once.Do(func() { p.readHeader() })
+	if p.header == nil {
+		return p.conn.RemoteAddr()
+	}
+
+	return p.header.RemoteAddr()
+}
+
+func (p *Conn) SetDeadline(t time.Time) error {
+	return p.conn.SetDeadline(t)
+}
+
+func (p *Conn) SetReadDeadline(t time.Time) error {
+	return p.conn.SetReadDeadline(t)
+}
+
+func (p *Conn) SetWriteDeadline(t time.Time) error {
+	return p.conn.SetWriteDeadline(t)
+}
+
+func (p *Conn) readHeader() error {
+	p.header, _ = Read(p.bufReader)
+	return nil
+
+}