Fix DNS transports

This commit is contained in:
世界
2025-12-29 20:44:30 +08:00
parent 4e94a64dcc
commit aa8dd6e44f
18 changed files with 754 additions and 220 deletions

View File

@@ -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{}
}