github
/
v2ray-core
mirror of https://github.com/v2ray/v2ray-core
1
0
Fork 0
Code Issues Releases Wiki Activity

fix udp dispatcher

pull/432/head
Darien Raymond 2017-02-04 21:55:59 +01:00
parent 4812e66229
commit 9978cf07e6
No known key found for this signature in database
GPG Key ID: 7251FFA14BB18169
1 changed files with 10 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
transport/internet/udp/dispatcher.go
View File

@ -15,57 +15,56 @@ type ResponseCallback func(payload *buf.Buffer)
type Dispatcher struct { type Dispatcher struct {
sync.RWMutex sync.RWMutex
conns map[string]ray.InboundRay conns map[v2net.Destination]ray.InboundRay
dispatcher dispatcher.Interface dispatcher dispatcher.Interface
} }
func NewDispatcher(dispatcher dispatcher.Interface) *Dispatcher { func NewDispatcher(dispatcher dispatcher.Interface) *Dispatcher {
return &Dispatcher{ return &Dispatcher{
conns: make(map[string]ray.InboundRay), conns: make(map[v2net.Destination]ray.InboundRay),
dispatcher: dispatcher, dispatcher: dispatcher,
} }
} }
func (v *Dispatcher) RemoveRay(name string) { func (v *Dispatcher) RemoveRay(dest v2net.Destination) {
v.Lock() v.Lock()
defer v.Unlock() defer v.Unlock()
if conn, found := v.conns[name]; found { if conn, found := v.conns[dest]; found {
conn.InboundInput().Close() conn.InboundInput().Close()
conn.InboundOutput().Close() conn.InboundOutput().Close()
delete(v.conns, name) delete(v.conns, dest)
} }
} }
func (v *Dispatcher) getInboundRay(ctx context.Context, dest v2net.Destination) (ray.InboundRay, bool) { func (v *Dispatcher) getInboundRay(ctx context.Context, dest v2net.Destination) (ray.InboundRay, bool) {
destString := dest.String()
v.Lock() v.Lock()
defer v.Unlock() defer v.Unlock()
if entry, found := v.conns[destString]; found { if entry, found := v.conns[dest]; found {
return entry, true return entry, true
} }
log.Info("UDP|Server: establishing new connection for ", dest) log.Info("UDP|Server: establishing new connection for ", dest)
inboundRay, _ := v.dispatcher.Dispatch(ctx, dest) inboundRay, _ := v.dispatcher.Dispatch(ctx, dest)
v.conns[dest] = inboundRay
return inboundRay, false return inboundRay, false
} }
func (v *Dispatcher) Dispatch(ctx context.Context, destination v2net.Destination, payload *buf.Buffer, callback ResponseCallback) { func (v *Dispatcher) Dispatch(ctx context.Context, destination v2net.Destination, payload *buf.Buffer, callback ResponseCallback) {
// TODO: Add user to destString // TODO: Add user to destString
destString := destination.String() log.Debug("UDP|Server: Dispatch request: ", destination)
log.Debug("UDP|Server: Dispatch request: ", destString)
inboundRay, existing := v.getInboundRay(ctx, destination) inboundRay, existing := v.getInboundRay(ctx, destination)
outputStream := inboundRay.InboundInput() outputStream := inboundRay.InboundInput()
if outputStream != nil { if outputStream != nil {
if err := outputStream.Write(payload); err != nil { if err := outputStream.Write(payload); err != nil {
v.RemoveRay(destString) v.RemoveRay(destination)
} }
} }
if !existing { if !existing {
go func() { go func() {
handleInput(inboundRay.InboundOutput(), callback) handleInput(inboundRay.InboundOutput(), callback)
v.RemoveRay(destString) v.RemoveRay(destination)
}() }()
} }
} }