Improve udp dns close

This commit is contained in:
世界 2022-07-24 21:25:41 +08:00
parent fde33fbb30
commit 9e9e6f7ee6
No known key found for this signature in database
GPG key ID: CD109927C34A63C4
2 changed files with 57 additions and 29 deletions

View file

@ -325,7 +325,7 @@ func (a *myInboundAdapter) NewError(ctx context.Context, err error) {
func NewError(logger log.ContextLogger, ctx context.Context, err error) { func NewError(logger log.ContextLogger, ctx context.Context, err error) {
common.Close(err) common.Close(err)
if E.IsClosedOrCanceled(err) { if E.IsClosedOrCanceled(err) {
logger.DebugContext(ctx, "connection closed") logger.TraceContext(ctx, "connection closed: ", err)
return return
} }
logger.ErrorContext(ctx, err) logger.ErrorContext(ctx, err)

View file

@ -6,6 +6,8 @@ import (
"io" "io"
"net" "net"
"os" "os"
"sync"
"time"
"github.com/sagernet/sing-box/adapter" "github.com/sagernet/sing-box/adapter"
C "github.com/sagernet/sing-box/constant" C "github.com/sagernet/sing-box/constant"
@ -14,6 +16,7 @@ import (
"github.com/sagernet/sing/common/buf" "github.com/sagernet/sing/common/buf"
M "github.com/sagernet/sing/common/metadata" M "github.com/sagernet/sing/common/metadata"
N "github.com/sagernet/sing/common/network" N "github.com/sagernet/sing/common/network"
"github.com/sagernet/sing/common/task"
"golang.org/x/net/dns/dnsmessage" "golang.org/x/net/dns/dnsmessage"
) )
@ -45,6 +48,7 @@ func (d *DNS) ListenPacket(ctx context.Context, destination M.Socksaddr) (net.Pa
} }
func (d *DNS) NewConnection(ctx context.Context, conn net.Conn, metadata adapter.InboundContext) error { func (d *DNS) NewConnection(ctx context.Context, conn net.Conn, metadata adapter.InboundContext) error {
defer conn.Close()
ctx = adapter.WithContext(ctx, &metadata) ctx = adapter.WithContext(ctx, &metadata)
_buffer := buf.StackNewSize(1024) _buffer := buf.StackNewSize(1024)
defer common.KeepAlive(_buffer) defer common.KeepAlive(_buffer)
@ -97,11 +101,16 @@ func (d *DNS) NewConnection(ctx context.Context, conn net.Conn, metadata adapter
} }
func (d *DNS) NewPacketConnection(ctx context.Context, conn N.PacketConn, metadata adapter.InboundContext) error { func (d *DNS) NewPacketConnection(ctx context.Context, conn N.PacketConn, metadata adapter.InboundContext) error {
defer conn.Close()
ctx = adapter.WithContext(ctx, &metadata) ctx = adapter.WithContext(ctx, &metadata)
_buffer := buf.StackNewSize(1024) _buffer := buf.StackNewSize(1024)
defer common.KeepAlive(_buffer) defer common.KeepAlive(_buffer)
buffer := common.Dup(_buffer) buffer := common.Dup(_buffer)
defer buffer.Release() defer buffer.Release()
var wg sync.WaitGroup
fastClose, cancel := context.WithCancel(ctx)
err := task.Run(fastClose, func() error {
var count int
for { for {
buffer.FullReset() buffer.FullReset()
destination, err := conn.ReadPacket(buffer) destination, err := conn.ReadPacket(buffer)
@ -118,7 +127,9 @@ func (d *DNS) NewPacketConnection(ctx context.Context, conn N.PacketConn, metada
metadata.Domain = string(question.Name.Data[:question.Name.Length-1]) metadata.Domain = string(question.Name.Data[:question.Name.Length-1])
d.logger.DebugContext(ctx, "inbound dns query ", formatDNSQuestion(question), " from ", metadata.Source) d.logger.DebugContext(ctx, "inbound dns query ", formatDNSQuestion(question), " from ", metadata.Source)
} }
wg.Add(1)
go func() error { go func() error {
defer wg.Done()
response, err := d.router.Exchange(ctx, &message) response, err := d.router.Exchange(ctx, &message)
if err != nil { if err != nil {
return err return err
@ -135,7 +146,24 @@ func (d *DNS) NewPacketConnection(ctx context.Context, conn N.PacketConn, metada
err = conn.WritePacket(responseBuffer, destination) err = conn.WritePacket(responseBuffer, destination)
return err return err
}() }()
count++
if count == 2 {
break
} }
}
cancel()
return nil
}, func() error {
timer := time.NewTimer(5 * time.Second)
select {
case <-timer.C:
cancel()
case <-fastClose.Done():
}
return nil
})
wg.Wait()
return err
} }
func formatDNSQuestion(question dnsmessage.Question) string { func formatDNSQuestion(question dnsmessage.Question) string {