sing-box/common/proxyproto/listener.go
2022-09-25 22:29:18 +08:00

45 lines
1 KiB
Go

package proxyproto
import (
std_bufio "bufio"
"net"
"github.com/sagernet/sing/common/buf"
"github.com/sagernet/sing/common/bufio"
M "github.com/sagernet/sing/common/metadata"
"github.com/pires/go-proxyproto"
)
type Listener struct {
net.Listener
AcceptNoHeader bool
}
func (l *Listener) Accept() (net.Conn, error) {
conn, err := l.Listener.Accept()
if err != nil {
return nil, err
}
bufReader := std_bufio.NewReader(conn)
header, err := proxyproto.Read(bufReader)
if err != nil && !(l.AcceptNoHeader && err == proxyproto.ErrNoProxyProtocol) {
return nil, err
}
if bufReader.Buffered() > 0 {
cache := buf.NewSize(bufReader.Buffered())
_, err = cache.ReadFullFrom(bufReader, cache.FreeLen())
if err != nil {
return nil, err
}
conn = bufio.NewCachedConn(conn, cache)
}
if header != nil {
return &bufio.AddrConn{Conn: conn, Metadata: M.Metadata{
Source: M.SocksaddrFromNet(header.SourceAddr).Unwrap(),
Destination: M.SocksaddrFromNet(header.DestinationAddr).Unwrap(),
}}, nil
}
return conn, nil
}