forked from v2fly/v2ray-core
-
Notifications
You must be signed in to change notification settings - Fork 0
/
dispatcher.go
82 lines (70 loc) · 1.86 KB
/
dispatcher.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
package udp
import (
"context"
"sync"
"v2ray.com/core/app/dispatcher"
"v2ray.com/core/app/log"
"v2ray.com/core/common/buf"
"v2ray.com/core/common/net"
"v2ray.com/core/transport/ray"
)
type ResponseCallback func(payload *buf.Buffer)
type Dispatcher struct {
sync.RWMutex
conns map[net.Destination]ray.InboundRay
dispatcher dispatcher.Interface
}
func NewDispatcher(dispatcher dispatcher.Interface) *Dispatcher {
return &Dispatcher{
conns: make(map[net.Destination]ray.InboundRay),
dispatcher: dispatcher,
}
}
func (v *Dispatcher) RemoveRay(dest net.Destination) {
v.Lock()
defer v.Unlock()
if conn, found := v.conns[dest]; found {
conn.InboundInput().Close()
conn.InboundOutput().Close()
delete(v.conns, dest)
}
}
func (v *Dispatcher) getInboundRay(ctx context.Context, dest net.Destination) (ray.InboundRay, bool) {
v.Lock()
defer v.Unlock()
if entry, found := v.conns[dest]; found {
return entry, true
}
log.Trace(newError("establishing new connection for ", dest))
inboundRay, _ := v.dispatcher.Dispatch(ctx, dest)
v.conns[dest] = inboundRay
return inboundRay, false
}
func (v *Dispatcher) Dispatch(ctx context.Context, destination net.Destination, payload *buf.Buffer, callback ResponseCallback) {
// TODO: Add user to destString
log.Trace(newError("dispatch request to: ", destination).AtDebug())
inboundRay, existing := v.getInboundRay(ctx, destination)
outputStream := inboundRay.InboundInput()
if outputStream != nil {
if err := outputStream.Write(buf.NewMultiBufferValue(payload)); err != nil {
v.RemoveRay(destination)
}
}
if !existing {
go func() {
handleInput(inboundRay.InboundOutput(), callback)
v.RemoveRay(destination)
}()
}
}
func handleInput(input ray.InputStream, callback ResponseCallback) {
for {
mb, err := input.Read()
if err != nil {
break
}
for _, b := range mb {
callback(b)
}
}
}