* fix udp dispatcher

* fix test
This commit is contained in:
dyhkwong 2024-01-15 23:33:15 +08:00 committed by GitHub
parent 77376ed94f
commit da5a28a088
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -28,7 +28,7 @@ type connEntry struct {
type Dispatcher struct { type Dispatcher struct {
sync.RWMutex sync.RWMutex
conns map[net.Destination]*connEntry conn *connEntry
dispatcher routing.Dispatcher dispatcher routing.Dispatcher
callback ResponseCallback callback ResponseCallback
callClose func() error callClose func() error
@ -36,19 +36,18 @@ type Dispatcher struct {
func NewDispatcher(dispatcher routing.Dispatcher, callback ResponseCallback) *Dispatcher { func NewDispatcher(dispatcher routing.Dispatcher, callback ResponseCallback) *Dispatcher {
return &Dispatcher{ return &Dispatcher{
conns: make(map[net.Destination]*connEntry),
dispatcher: dispatcher, dispatcher: dispatcher,
callback: callback, callback: callback,
} }
} }
func (v *Dispatcher) RemoveRay(dest net.Destination) { func (v *Dispatcher) RemoveRay() {
v.Lock() v.Lock()
defer v.Unlock() defer v.Unlock()
if conn, found := v.conns[dest]; found { if v.conn != nil {
common.Close(conn.link.Reader) common.Close(v.conn.link.Reader)
common.Close(conn.link.Writer) common.Close(v.conn.link.Writer)
delete(v.conns, dest) v.conn = nil
} }
} }
@ -56,8 +55,8 @@ func (v *Dispatcher) getInboundRay(ctx context.Context, dest net.Destination) (*
v.Lock() v.Lock()
defer v.Unlock() defer v.Unlock()
if entry, found := v.conns[dest]; found { if v.conn != nil {
return entry, nil return v.conn, nil
} }
newError("establishing new connection for ", dest).WriteToLog() newError("establishing new connection for ", dest).WriteToLog()
@ -65,7 +64,7 @@ func (v *Dispatcher) getInboundRay(ctx context.Context, dest net.Destination) (*
ctx, cancel := context.WithCancel(ctx) ctx, cancel := context.WithCancel(ctx)
removeRay := func() { removeRay := func() {
cancel() cancel()
v.RemoveRay(dest) v.RemoveRay()
} }
timer := signal.CancelAfterInactivity(ctx, removeRay, time.Minute) timer := signal.CancelAfterInactivity(ctx, removeRay, time.Minute)
@ -79,7 +78,7 @@ func (v *Dispatcher) getInboundRay(ctx context.Context, dest net.Destination) (*
timer: timer, timer: timer,
cancel: removeRay, cancel: removeRay,
} }
v.conns[dest] = entry v.conn = entry
go handleInput(ctx, entry, dest, v.callback, v.callClose) go handleInput(ctx, entry, dest, v.callback, v.callClose)
return entry, nil return entry, nil
} }
@ -130,6 +129,9 @@ func handleInput(ctx context.Context, conn *connEntry, dest net.Destination, cal
} }
timer.Update() timer.Update()
for _, b := range mb { for _, b := range mb {
if b.UDP != nil {
dest = *b.UDP
}
callback(ctx, &udp.Packet{ callback(ctx, &udp.Packet{
Payload: b, Payload: b,
Source: dest, Source: dest,
@ -153,7 +155,6 @@ func DialDispatcher(ctx context.Context, dispatcher routing.Dispatcher) (net.Pac
} }
d := &Dispatcher{ d := &Dispatcher{
conns: make(map[net.Destination]*connEntry),
dispatcher: dispatcher, dispatcher: dispatcher,
callback: c.callback, callback: c.callback,
callClose: c.Close, callClose: c.Close,
@ -199,7 +200,9 @@ func (c *dispatcherConn) WriteTo(p []byte, addr net.Addr) (int, error) {
n := copy(raw, p) n := copy(raw, p)
buffer.Resize(0, int32(n)) buffer.Resize(0, int32(n))
c.dispatcher.Dispatch(c.ctx, net.DestinationFromAddr(addr), buffer) destination := net.DestinationFromAddr(addr)
buffer.UDP = &destination
c.dispatcher.Dispatch(c.ctx, destination, buffer)
return n, nil return n, nil
} }