forked from xjasonlyu/tun2socks
-
Notifications
You must be signed in to change notification settings - Fork 0
/
endpoint.go
153 lines (126 loc) · 3.43 KB
/
endpoint.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
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
// Package iobased provides the implementation of io.ReadWriter
// based data-link layer endpoints.
package iobased
import (
"context"
"errors"
"io"
"sync"
"gvisor.dev/gvisor/pkg/buffer"
"gvisor.dev/gvisor/pkg/tcpip"
"gvisor.dev/gvisor/pkg/tcpip/header"
"gvisor.dev/gvisor/pkg/tcpip/link/channel"
"gvisor.dev/gvisor/pkg/tcpip/stack"
)
const (
// Queue length for outbound packet, arriving for read. Overflow
// causes packet drops.
defaultOutQueueLen = 1 << 10
)
// Endpoint implements the interface of stack.LinkEndpoint from io.ReadWriter.
type Endpoint struct {
*channel.Endpoint
// rw is the io.ReadWriter for reading and writing packets.
rw io.ReadWriter
// mtu (maximum transmission unit) is the maximum size of a packet.
mtu uint32
// offset can be useful when perform TUN device I/O with TUN_PI enabled.
offset int
// once is used to perform the init action once when attaching.
once sync.Once
// wg keeps track of running goroutines.
wg sync.WaitGroup
}
// New returns stack.LinkEndpoint(.*Endpoint) and error.
func New(rw io.ReadWriter, mtu uint32, offset int) (*Endpoint, error) {
if mtu == 0 {
return nil, errors.New("MTU size is zero")
}
if rw == nil {
return nil, errors.New("RW interface is nil")
}
if offset < 0 {
return nil, errors.New("offset must be non-negative")
}
return &Endpoint{
Endpoint: channel.New(defaultOutQueueLen, mtu, ""),
rw: rw,
mtu: mtu,
offset: offset,
}, nil
}
// Attach launches the goroutine that reads packets from io.Reader and
// dispatches them via the provided dispatcher.
func (e *Endpoint) Attach(dispatcher stack.NetworkDispatcher) {
e.Endpoint.Attach(dispatcher)
e.once.Do(func() {
ctx, cancel := context.WithCancel(context.Background())
e.wg.Add(2)
go func() {
e.outboundLoop(ctx)
e.wg.Done()
}()
go func() {
e.dispatchLoop(cancel)
e.wg.Done()
}()
})
}
func (e *Endpoint) Wait() {
e.wg.Wait()
}
// dispatchLoop dispatches packets to upper layer.
func (e *Endpoint) dispatchLoop(cancel context.CancelFunc) {
// Call cancel() to ensure (*Endpoint).outboundLoop(context.Context) exits
// gracefully after (*Endpoint).dispatchLoop(context.CancelFunc) returns.
defer cancel()
offset, mtu := e.offset, int(e.mtu)
for {
data := make([]byte, offset+mtu)
n, err := e.rw.Read(data)
if err != nil {
break
}
if n == 0 || n > mtu {
continue
}
if !e.IsAttached() {
continue /* unattached, drop packet */
}
pkt := stack.NewPacketBuffer(stack.PacketBufferOptions{
Payload: buffer.MakeWithData(data[offset : offset+n]),
})
switch header.IPVersion(data[offset:]) {
case header.IPv4Version:
e.InjectInbound(header.IPv4ProtocolNumber, pkt)
case header.IPv6Version:
e.InjectInbound(header.IPv6ProtocolNumber, pkt)
}
pkt.DecRef()
}
}
// outboundLoop reads outbound packets from channel, and then it calls
// writePacket to send those packets back to lower layer.
func (e *Endpoint) outboundLoop(ctx context.Context) {
for {
pkt := e.ReadContext(ctx)
if pkt.IsNil() {
break
}
e.writePacket(pkt)
}
}
// writePacket writes outbound packets to the io.Writer.
func (e *Endpoint) writePacket(pkt stack.PacketBufferPtr) tcpip.Error {
defer pkt.DecRef()
buf := pkt.ToBuffer()
defer buf.Release()
if e.offset != 0 {
v := buffer.NewViewWithData(make([]byte, e.offset))
_ = buf.Prepend(v)
}
if _, err := e.rw.Write(buf.Flatten()); err != nil {
return &tcpip.ErrInvalidEndpointState{}
}
return nil
}