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

Loading…
Cancel
Save