ginuerzh
148 lines
3.3 KiB
Go
148 lines
3.3 KiB
Go
package tungo
|
|
|
|
import (
|
|
"context"
|
|
"io"
|
|
"sync"
|
|
|
|
"github.com/go-gost/core/logger"
|
|
"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
|
|
|
|
log logger.Logger
|
|
}
|
|
|
|
// New returns stack.LinkEndpoint(.*Endpoint) and error.
|
|
func newEndpoint(rw io.ReadWriter, mtu int, log logger.Logger) *endpoint {
|
|
return &endpoint{
|
|
Endpoint: channel.New(defaultOutQueueLen, uint32(mtu), ""),
|
|
rw: rw,
|
|
mtu: uint32(mtu),
|
|
log: log,
|
|
}
|
|
}
|
|
|
|
// 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 {
|
|
e.log.Error(err)
|
|
break
|
|
}
|
|
|
|
if e.log.IsLevelEnabled(logger.TraceLevel) {
|
|
e.log.Tracef("read: (%d) % x", n, data[:n])
|
|
}
|
|
|
|
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 == nil {
|
|
break
|
|
}
|
|
e.writePacket(pkt)
|
|
}
|
|
}
|
|
|
|
// writePacket writes outbound packets to the io.Writer.
|
|
func (e *endpoint) writePacket(pkt *stack.PacketBuffer) 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
|
|
}
|