76 lines
1.7 KiB
Go
76 lines
1.7 KiB
Go
package wireguard
|
|
|
|
import (
|
|
"github.com/sagernet/sing-tun"
|
|
"github.com/sagernet/sing/common/buf"
|
|
)
|
|
|
|
var _ Device = (*natDeviceWrapper)(nil)
|
|
|
|
type natDeviceWrapper struct {
|
|
Device
|
|
outbound chan *buf.Buffer
|
|
mapping *tun.NatMapping
|
|
writer *tun.NatWriter
|
|
}
|
|
|
|
func NewNATDevice(upstream Device, ipRewrite bool) NatDevice {
|
|
wrapper := &natDeviceWrapper{
|
|
Device: upstream,
|
|
outbound: make(chan *buf.Buffer, 256),
|
|
mapping: tun.NewNatMapping(ipRewrite),
|
|
}
|
|
if ipRewrite {
|
|
wrapper.writer = tun.NewNatWriter(upstream.Inet4Address(), upstream.Inet6Address())
|
|
}
|
|
return wrapper
|
|
}
|
|
|
|
func (d *natDeviceWrapper) Read(p []byte, offset int) (int, error) {
|
|
select {
|
|
case packet := <-d.outbound:
|
|
defer packet.Release()
|
|
return copy(p[offset:], packet.Bytes()), nil
|
|
default:
|
|
}
|
|
return d.Device.Read(p, offset)
|
|
}
|
|
|
|
func (d *natDeviceWrapper) Write(p []byte, offset int) (int, error) {
|
|
packet := p[offset:]
|
|
handled, err := d.mapping.WritePacket(packet)
|
|
if handled {
|
|
return len(packet), err
|
|
}
|
|
return d.Device.Write(p, offset)
|
|
}
|
|
|
|
func (d *natDeviceWrapper) CreateDestination(session tun.RouteSession, conn tun.RouteContext) tun.DirectDestination {
|
|
d.mapping.CreateSession(session, conn)
|
|
return &natDestinationWrapper{d, session}
|
|
}
|
|
|
|
var _ tun.DirectDestination = (*natDestinationWrapper)(nil)
|
|
|
|
type natDestinationWrapper struct {
|
|
device *natDeviceWrapper
|
|
session tun.RouteSession
|
|
}
|
|
|
|
func (d *natDestinationWrapper) WritePacket(buffer *buf.Buffer) error {
|
|
if d.device.writer != nil {
|
|
d.device.writer.RewritePacket(buffer.Bytes())
|
|
}
|
|
d.device.outbound <- buffer
|
|
return nil
|
|
}
|
|
|
|
func (d *natDestinationWrapper) Close() error {
|
|
d.device.mapping.DeleteSession(d.session)
|
|
return nil
|
|
}
|
|
|
|
func (d *natDestinationWrapper) Timeout() bool {
|
|
return false
|
|
}
|