add proxy proto support

tls
dingjun 6 years ago
parent bfe65cf0e6
commit f6b305c18e

@ -0,0 +1,50 @@
package main
import (
"bufio"
"net"
proxyproto "github.com/pires/go-proxyproto"
)
type protoListener struct {
net.Listener
}
type protoConn struct {
net.Conn
headerDone bool
r *bufio.Reader
proxy *proxyproto.Header
}
func (l *protoListener) Accept() (net.Conn, error) {
c, err := l.Listener.Accept()
if err != nil {
return nil, err
}
return &protoConn{Conn: c}, err
}
func (c *protoConn) Read(buf []byte) (int, error) {
var err error
if !c.headerDone {
c.r = bufio.NewReader(c.Conn)
c.proxy, err = proxyproto.Read(c.r)
if err != nil && err != proxyproto.ErrNoProxyProtocol {
return 0, err
}
c.headerDone = true
return c.r.Read(buf)
}
return c.r.Read(buf)
}
func (c *protoConn) RemoteAddr() net.Addr {
if c.proxy == nil {
return c.Conn.RemoteAddr()
}
return &net.TCPAddr{
IP: c.proxy.SourceAddress,
Port: int(c.proxy.SourcePort)}
}

@ -115,23 +115,23 @@ func main() {
go func(lst listen) { go func(lst listen) {
var l net.Listener var l net.Listener
var err error var err error
if lst.Key == "" || lst.Cert == "" {
l, err = net.Listen("tcp", fmt.Sprintf(":%d", lst.Port)) l, err = net.Listen("tcp", fmt.Sprintf(":%d", lst.Port))
} else { if err != nil {
log.Fatal(err)
}
defer l.Close()
if lst.Key != "" && lst.Cert != "" {
cert, err := tls.LoadX509KeyPair(lst.Cert, lst.Key) cert, err := tls.LoadX509KeyPair(lst.Cert, lst.Key)
if err != nil { if err != nil {
log.Fatal(err) log.Fatal(err)
} }
l, err = tls.Listen("tcp", fmt.Sprintf(":%d", lst.Port), &tls.Config{ l = tls.NewListener(&protoListener{l}, &tls.Config{
Certificates: []tls.Certificate{cert}, Certificates: []tls.Certificate{cert},
}) })
} }
if err != nil {
log.Fatal(err)
}
defer l.Close()
for { for {
c, err := l.Accept() c, err := l.Accept()
if err != nil { if err != nil {

Loading…
Cancel
Save