Fix DNS transports
This commit is contained in:
@@ -2,9 +2,8 @@ package transport
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net"
|
||||
"os"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
|
||||
"github.com/sagernet/sing-box/adapter"
|
||||
"github.com/sagernet/sing-box/common/dialer"
|
||||
@@ -28,15 +27,23 @@ func RegisterUDP(registry *dns.TransportRegistry) {
|
||||
}
|
||||
|
||||
type UDPTransport struct {
|
||||
dns.TransportAdapter
|
||||
logger logger.ContextLogger
|
||||
dialer N.Dialer
|
||||
serverAddr M.Socksaddr
|
||||
udpSize int
|
||||
tcpTransport *TCPTransport
|
||||
access sync.Mutex
|
||||
conn *dnsConnection
|
||||
done chan struct{}
|
||||
*BaseTransport
|
||||
|
||||
dialer N.Dialer
|
||||
serverAddr M.Socksaddr
|
||||
udpSize atomic.Int32
|
||||
|
||||
connector *Connector[*Connection]
|
||||
|
||||
callbackAccess sync.RWMutex
|
||||
queryId uint16
|
||||
callbacks map[uint16]*udpCallback
|
||||
}
|
||||
|
||||
type udpCallback struct {
|
||||
access sync.Mutex
|
||||
response *mDNS.Msg
|
||||
done chan struct{}
|
||||
}
|
||||
|
||||
func NewUDP(ctx context.Context, logger log.ContextLogger, tag string, options option.RemoteDNSServerOptions) (adapter.DNSTransport, error) {
|
||||
@@ -54,180 +61,198 @@ func NewUDP(ctx context.Context, logger log.ContextLogger, tag string, options o
|
||||
return NewUDPRaw(logger, dns.NewTransportAdapterWithRemoteOptions(C.DNSTypeUDP, tag, options), transportDialer, serverAddr), nil
|
||||
}
|
||||
|
||||
func NewUDPRaw(logger logger.ContextLogger, adapter dns.TransportAdapter, dialer N.Dialer, serverAddr M.Socksaddr) *UDPTransport {
|
||||
return &UDPTransport{
|
||||
TransportAdapter: adapter,
|
||||
logger: logger,
|
||||
dialer: dialer,
|
||||
serverAddr: serverAddr,
|
||||
udpSize: 2048,
|
||||
tcpTransport: &TCPTransport{
|
||||
dialer: dialer,
|
||||
serverAddr: serverAddr,
|
||||
},
|
||||
done: make(chan struct{}),
|
||||
func NewUDPRaw(logger logger.ContextLogger, adapter dns.TransportAdapter, dialerInstance N.Dialer, serverAddr M.Socksaddr) *UDPTransport {
|
||||
t := &UDPTransport{
|
||||
BaseTransport: NewBaseTransport(adapter, logger),
|
||||
dialer: dialerInstance,
|
||||
serverAddr: serverAddr,
|
||||
callbacks: make(map[uint16]*udpCallback),
|
||||
}
|
||||
t.udpSize.Store(2048)
|
||||
t.connector = NewSingleflightConnector(t.CloseContext(), t.dial)
|
||||
return t
|
||||
}
|
||||
|
||||
func (t *UDPTransport) dial(ctx context.Context) (*Connection, error) {
|
||||
rawConn, err := t.dialer.DialContext(ctx, N.NetworkUDP, t.serverAddr)
|
||||
if err != nil {
|
||||
return nil, E.Cause(err, "dial UDP connection")
|
||||
}
|
||||
conn := WrapConnection(rawConn)
|
||||
go t.recvLoop(conn)
|
||||
return conn, nil
|
||||
}
|
||||
|
||||
func (t *UDPTransport) Start(stage adapter.StartStage) error {
|
||||
if stage != adapter.StartStateStart {
|
||||
return nil
|
||||
}
|
||||
err := t.SetStarted()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return dialer.InitializeDetour(t.dialer)
|
||||
}
|
||||
|
||||
func (t *UDPTransport) Close() error {
|
||||
t.access.Lock()
|
||||
defer t.access.Unlock()
|
||||
close(t.done)
|
||||
t.done = make(chan struct{})
|
||||
return nil
|
||||
return E.Errors(t.BaseTransport.Close(), t.connector.Close())
|
||||
}
|
||||
|
||||
func (t *UDPTransport) Reset() {
|
||||
t.connector.Reset()
|
||||
}
|
||||
|
||||
func (t *UDPTransport) nextAvailableQueryId() (uint16, error) {
|
||||
start := t.queryId
|
||||
for {
|
||||
t.queryId++
|
||||
if _, exists := t.callbacks[t.queryId]; !exists {
|
||||
return t.queryId, nil
|
||||
}
|
||||
if t.queryId == start {
|
||||
return 0, E.New("no available query ID")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (t *UDPTransport) Exchange(ctx context.Context, message *mDNS.Msg) (*mDNS.Msg, error) {
|
||||
if !t.BeginQuery() {
|
||||
return nil, ErrTransportClosed
|
||||
}
|
||||
defer t.EndQuery()
|
||||
|
||||
response, err := t.exchange(ctx, message)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if response.Truncated {
|
||||
t.logger.InfoContext(ctx, "response truncated, retrying with TCP")
|
||||
return t.tcpTransport.Exchange(ctx, message)
|
||||
t.Logger.InfoContext(ctx, "response truncated, retrying with TCP")
|
||||
return t.exchangeTCP(ctx, message)
|
||||
}
|
||||
return response, nil
|
||||
}
|
||||
|
||||
func (t *UDPTransport) exchangeTCP(ctx context.Context, message *mDNS.Msg) (*mDNS.Msg, error) {
|
||||
conn, err := t.dialer.DialContext(ctx, N.NetworkTCP, t.serverAddr)
|
||||
if err != nil {
|
||||
return nil, E.Cause(err, "dial TCP connection")
|
||||
}
|
||||
defer conn.Close()
|
||||
err = WriteMessage(conn, message.Id, message)
|
||||
if err != nil {
|
||||
return nil, E.Cause(err, "write request")
|
||||
}
|
||||
response, err := ReadMessage(conn)
|
||||
if err != nil {
|
||||
return nil, E.Cause(err, "read response")
|
||||
}
|
||||
return response, nil
|
||||
}
|
||||
|
||||
func (t *UDPTransport) exchange(ctx context.Context, message *mDNS.Msg) (*mDNS.Msg, error) {
|
||||
t.access.Lock()
|
||||
if edns0Opt := message.IsEdns0(); edns0Opt != nil {
|
||||
if udpSize := int(edns0Opt.UDPSize()); udpSize > t.udpSize {
|
||||
t.udpSize = udpSize
|
||||
close(t.done)
|
||||
t.done = make(chan struct{})
|
||||
udpSize := int32(edns0Opt.UDPSize())
|
||||
for {
|
||||
current := t.udpSize.Load()
|
||||
if udpSize <= current {
|
||||
break
|
||||
}
|
||||
if t.udpSize.CompareAndSwap(current, udpSize) {
|
||||
t.connector.Reset()
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
t.access.Unlock()
|
||||
conn, err := t.open(ctx)
|
||||
|
||||
conn, err := t.connector.Get(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
buffer := buf.NewSize(1 + message.Len())
|
||||
defer buffer.Release()
|
||||
exMessage := *message
|
||||
exMessage.Compress = true
|
||||
messageId := message.Id
|
||||
callback := &dnsCallback{
|
||||
|
||||
callback := &udpCallback{
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
conn.access.Lock()
|
||||
conn.queryId++
|
||||
exMessage.Id = conn.queryId
|
||||
conn.callbacks[exMessage.Id] = callback
|
||||
conn.access.Unlock()
|
||||
|
||||
t.callbackAccess.Lock()
|
||||
queryId, err := t.nextAvailableQueryId()
|
||||
if err != nil {
|
||||
t.callbackAccess.Unlock()
|
||||
return nil, err
|
||||
}
|
||||
t.callbacks[queryId] = callback
|
||||
t.callbackAccess.Unlock()
|
||||
|
||||
defer func() {
|
||||
conn.access.Lock()
|
||||
delete(conn.callbacks, exMessage.Id)
|
||||
conn.access.Unlock()
|
||||
t.callbackAccess.Lock()
|
||||
delete(t.callbacks, queryId)
|
||||
t.callbackAccess.Unlock()
|
||||
}()
|
||||
|
||||
buffer := buf.NewSize(1 + message.Len())
|
||||
defer buffer.Release()
|
||||
|
||||
exMessage := *message
|
||||
exMessage.Compress = true
|
||||
originalId := message.Id
|
||||
exMessage.Id = queryId
|
||||
|
||||
rawMessage, err := exMessage.PackBuffer(buffer.FreeBytes())
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
_, err = conn.Write(rawMessage)
|
||||
if err != nil {
|
||||
conn.Close(err)
|
||||
return nil, err
|
||||
conn.CloseWithError(err)
|
||||
return nil, E.Cause(err, "write request")
|
||||
}
|
||||
|
||||
select {
|
||||
case <-callback.done:
|
||||
callback.message.Id = messageId
|
||||
return callback.message, nil
|
||||
case <-conn.done:
|
||||
return nil, conn.err
|
||||
case <-t.done:
|
||||
return nil, os.ErrClosed
|
||||
callback.response.Id = originalId
|
||||
return callback.response, nil
|
||||
case <-conn.Done():
|
||||
return nil, conn.CloseError()
|
||||
case <-t.CloseContext().Done():
|
||||
return nil, ErrTransportClosed
|
||||
case <-ctx.Done():
|
||||
conn.Close(ctx.Err())
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
func (t *UDPTransport) open(ctx context.Context) (*dnsConnection, error) {
|
||||
t.access.Lock()
|
||||
defer t.access.Unlock()
|
||||
if t.conn != nil {
|
||||
select {
|
||||
case <-t.conn.done:
|
||||
default:
|
||||
return t.conn, nil
|
||||
}
|
||||
}
|
||||
conn, err := t.dialer.DialContext(ctx, N.NetworkUDP, t.serverAddr)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
dnsConn := &dnsConnection{
|
||||
Conn: conn,
|
||||
done: make(chan struct{}),
|
||||
callbacks: make(map[uint16]*dnsCallback),
|
||||
}
|
||||
go t.recvLoop(dnsConn)
|
||||
t.conn = dnsConn
|
||||
return dnsConn, nil
|
||||
}
|
||||
|
||||
func (t *UDPTransport) recvLoop(conn *dnsConnection) {
|
||||
func (t *UDPTransport) recvLoop(conn *Connection) {
|
||||
for {
|
||||
buffer := buf.NewSize(t.udpSize)
|
||||
buffer := buf.NewSize(int(t.udpSize.Load()))
|
||||
_, err := buffer.ReadOnceFrom(conn)
|
||||
if err != nil {
|
||||
buffer.Release()
|
||||
conn.Close(err)
|
||||
conn.CloseWithError(err)
|
||||
return
|
||||
}
|
||||
|
||||
var message mDNS.Msg
|
||||
err = message.Unpack(buffer.Bytes())
|
||||
buffer.Release()
|
||||
if err != nil {
|
||||
conn.Close(err)
|
||||
return
|
||||
t.Logger.Debug("discarded malformed UDP response: ", err)
|
||||
continue
|
||||
}
|
||||
conn.access.RLock()
|
||||
callback, loaded := conn.callbacks[message.Id]
|
||||
conn.access.RUnlock()
|
||||
|
||||
t.callbackAccess.RLock()
|
||||
callback, loaded := t.callbacks[message.Id]
|
||||
t.callbackAccess.RUnlock()
|
||||
|
||||
if !loaded {
|
||||
continue
|
||||
}
|
||||
|
||||
callback.access.Lock()
|
||||
select {
|
||||
case <-callback.done:
|
||||
default:
|
||||
callback.message = &message
|
||||
callback.response = &message
|
||||
close(callback.done)
|
||||
}
|
||||
callback.access.Unlock()
|
||||
}
|
||||
}
|
||||
|
||||
type dnsConnection struct {
|
||||
net.Conn
|
||||
access sync.RWMutex
|
||||
done chan struct{}
|
||||
closeOnce sync.Once
|
||||
err error
|
||||
queryId uint16
|
||||
callbacks map[uint16]*dnsCallback
|
||||
}
|
||||
|
||||
func (c *dnsConnection) Close(err error) {
|
||||
c.closeOnce.Do(func() {
|
||||
c.err = err
|
||||
close(c.done)
|
||||
})
|
||||
c.Conn.Close()
|
||||
}
|
||||
|
||||
type dnsCallback struct {
|
||||
access sync.Mutex
|
||||
message *mDNS.Msg
|
||||
done chan struct{}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user