Add callClose to UDP Dispatcher

Fixes https://github.com/XTLS/Xray-core/issues/1611
This commit is contained in:
RPRX 2023-02-08 14:59:14 +08:00 committed by GitHub
parent f32921df30
commit 9046eda5ce
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

@ -31,6 +31,7 @@ type Dispatcher struct {
conns map[net.Destination]*connEntry conns map[net.Destination]*connEntry
dispatcher routing.Dispatcher dispatcher routing.Dispatcher
callback ResponseCallback callback ResponseCallback
callClose func() error
} }
func NewDispatcher(dispatcher routing.Dispatcher, callback ResponseCallback) *Dispatcher { func NewDispatcher(dispatcher routing.Dispatcher, callback ResponseCallback) *Dispatcher {
@ -79,7 +80,7 @@ func (v *Dispatcher) getInboundRay(ctx context.Context, dest net.Destination) (*
cancel: removeRay, cancel: removeRay,
} }
v.conns[dest] = entry v.conns[dest] = entry
go handleInput(ctx, entry, dest, v.callback) go handleInput(ctx, entry, dest, v.callback, v.callClose)
return entry, nil return entry, nil
} }
@ -102,8 +103,13 @@ func (v *Dispatcher) Dispatch(ctx context.Context, destination net.Destination,
} }
} }
func handleInput(ctx context.Context, conn *connEntry, dest net.Destination, callback ResponseCallback) { func handleInput(ctx context.Context, conn *connEntry, dest net.Destination, callback ResponseCallback, callClose func() error) {
defer conn.cancel() defer func() {
conn.cancel()
if callClose != nil {
callClose()
}
}()
input := conn.link.Reader input := conn.link.Reader
timer := conn.timer timer := conn.timer
@ -144,7 +150,12 @@ func DialDispatcher(ctx context.Context, dispatcher routing.Dispatcher) (net.Pac
done: done.New(), done: done.New(),
} }
d := NewDispatcher(dispatcher, c.callback) d := &Dispatcher{
conns: make(map[net.Destination]*connEntry),
dispatcher: dispatcher,
callback: c.callback,
callClose: c.Close,
}
c.dispatcher = d c.dispatcher = d
return c, nil return c, nil
} }
@ -162,16 +173,22 @@ func (c *dispatcherConn) callback(ctx context.Context, packet *udp.Packet) {
} }
func (c *dispatcherConn) ReadFrom(p []byte) (int, net.Addr, error) { func (c *dispatcherConn) ReadFrom(p []byte) (int, net.Addr, error) {
var packet *udp.Packet
s:
select { select {
case <-c.done.Wait(): case <-c.done.Wait():
return 0, nil, io.EOF select {
case packet := <-c.cache: case packet = <-c.cache:
n := copy(p, packet.Payload.Bytes()) break s
return n, &net.UDPAddr{ default:
IP: packet.Source.Address.IP(), return 0, nil, io.EOF
Port: int(packet.Source.Port), }
}, nil case packet = <-c.cache:
} }
return copy(p, packet.Payload.Bytes()), &net.UDPAddr{
IP: packet.Source.Address.IP(),
Port: int(packet.Source.Port),
}, nil
} }
func (c *dispatcherConn) WriteTo(p []byte, addr net.Addr) (int, error) { func (c *dispatcherConn) WriteTo(p []byte, addr net.Addr) (int, error) {