Initial commit

protocol.go

@@ -0,0 +1,183 @@
+package proxyproto
+
+import (
+	"bufio"
+	"bytes"
+	"fmt"
+	"log"
+	"net"
+	"strconv"
+	"strings"
+	"sync"
+	"time"
+)
+
+var (
+	// prefix is the string we look for at the start of a connection
+	// to check if this connection is using the proxy protocol
+	prefix    = []byte("PROXY ")
+	prefixLen = len(prefix)
+)
+
+// Listener is used to wrap an underlying listener,
+// whose connections may be using the HAProxy Proxy Protocol (version 1).
+// If the connection is using the protocol, the RemoteAddr() will return
+// the correct client address.
+type Listener struct {
+	Listener net.Listener
+}
+
+// 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
+	dstAddr   *net.TCPAddr
+	srcAddr   *net.TCPAddr
+	once      sync.Once
+}
+
+// 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), 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) *Conn {
+	pConn := &Conn{
+		bufReader: bufio.NewReader(conn),
+		conn:      conn,
+	}
+	return pConn
+}
+
+func (p *Conn) Read(b []byte) (int, error) {
+	var err error
+	p.once.Do(func() { err = p.checkPrefix() })
+	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 {
+	return p.conn.LocalAddr()
+}
+
+func (p *Conn) RemoteAddr() net.Addr {
+	p.once.Do(func() {
+		if err := p.checkPrefix(); err != nil {
+			log.Printf("[ERR] Failed to read proxy prefix: %v", err)
+		}
+	})
+	if p.srcAddr != nil {
+		return p.srcAddr
+	}
+	return p.conn.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) checkPrefix() error {
+	// Incrementally check each byte of the prefix
+	for i := 1; i <= prefixLen; i++ {
+		inp, err := p.bufReader.Peek(i)
+		if err != nil {
+			return err
+		}
+
+		// Check for a prefix mis-match, quit early
+		if !bytes.Equal(inp, prefix[:i]) {
+			return nil
+		}
+	}
+
+	// Read the header line
+	header, err := p.bufReader.ReadString('\n')
+	if err != nil {
+		p.conn.Close()
+		return err
+	}
+
+	// Strip the carriage return and new line
+	header = header[:len(header)-2]
+
+	// Split on spaces, should be (PROXY <type> <src addr> <dst addr> <src port> <dst port>)
+	parts := strings.Split(header, " ")
+	if len(parts) != 6 {
+		p.conn.Close()
+		return fmt.Errorf("Invalid header line: %s", header)
+	}
+
+	// Verify the type is known
+	switch parts[1] {
+	case "TCP4":
+	case "TCP6":
+	default:
+		p.conn.Close()
+		return fmt.Errorf("Unhandled address type: %s", parts[1])
+	}
+
+	// Parse out the source address
+	ip := net.ParseIP(parts[2])
+	if ip == nil {
+		p.conn.Close()
+		return fmt.Errorf("Invalid source ip: %s", parts[2])
+	}
+	port, err := strconv.Atoi(parts[4])
+	if err != nil {
+		p.conn.Close()
+		return fmt.Errorf("Invalid source port: %s", parts[4])
+	}
+	p.srcAddr = &net.TCPAddr{IP: ip, Port: port}
+
+	// Parse out the destination address
+	ip = net.ParseIP(parts[3])
+	if ip == nil {
+		p.conn.Close()
+		return fmt.Errorf("Invalid destination ip: %s", parts[3])
+	}
+	port, err = strconv.Atoi(parts[5])
+	if err != nil {
+		p.conn.Close()
+		return fmt.Errorf("Invalid destination port: %s", parts[5])
+	}
+	p.dstAddr = &net.TCPAddr{IP: ip, Port: port}
+
+	return nil
+}