103 lines
1.9 KiB
Go
103 lines
1.9 KiB
Go
// From https://github.com/goburrow/netforward
|
|
package netforward
|
|
|
|
import (
|
|
"io"
|
|
"log"
|
|
"net"
|
|
"sync"
|
|
)
|
|
|
|
type packetConn struct {
|
|
conn net.PacketConn
|
|
addr net.Addr
|
|
}
|
|
|
|
func (p *packetConn) Write(b []byte) (int, error) {
|
|
return p.conn.WriteTo(b, p.addr)
|
|
}
|
|
|
|
type syncConns struct {
|
|
mu sync.RWMutex
|
|
conns map[string]net.Conn
|
|
}
|
|
|
|
func (f *syncConns) get(addr net.Addr) (net.Conn, bool) {
|
|
var key string
|
|
if addr != nil {
|
|
key = addr.String()
|
|
}
|
|
f.mu.RLock()
|
|
defer f.mu.RUnlock()
|
|
conn, ok := f.conns[key]
|
|
return conn, ok
|
|
}
|
|
|
|
func (f *syncConns) set(addr net.Addr, conn net.Conn) {
|
|
var key string
|
|
if addr != nil {
|
|
key = addr.String()
|
|
}
|
|
f.mu.Lock()
|
|
defer f.mu.Unlock()
|
|
f.conns[key] = conn
|
|
}
|
|
|
|
func (f *syncConns) remove(addr net.Addr) {
|
|
var key string
|
|
if addr != nil {
|
|
key = addr.String()
|
|
}
|
|
f.mu.Lock()
|
|
defer f.mu.Unlock()
|
|
delete(f.conns, key)
|
|
}
|
|
|
|
func ForwardPacket(remote Dialer, local net.PacketConn) error {
|
|
conns := syncConns{
|
|
conns: make(map[string]net.Conn),
|
|
}
|
|
buf := getBuffer()
|
|
defer releaseBuffer(buf)
|
|
for {
|
|
nr, addr, err := local.ReadFrom(buf)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
remoteConn, ok := conns.get(addr)
|
|
if !ok {
|
|
remoteConn, err = remote.Dial()
|
|
if err != nil {
|
|
log.Printf("%s: dial failed: %v", addr, err)
|
|
continue
|
|
}
|
|
conns.set(addr, remoteConn)
|
|
// unix datagram does not have a network addr
|
|
if addr != nil {
|
|
go forwardPacket(remoteConn, local, addr, &conns)
|
|
}
|
|
}
|
|
nw, err := remoteConn.Write(buf[:nr])
|
|
if err != nil {
|
|
log.Printf("%s: write failed: %v", addr, err)
|
|
continue
|
|
}
|
|
if nr != nw {
|
|
log.Printf("%s: %v", addr, io.ErrShortWrite)
|
|
}
|
|
}
|
|
}
|
|
|
|
func forwardPacket(remote net.Conn, local net.PacketConn, addr net.Addr, conns *syncConns) {
|
|
defer remote.Close()
|
|
defer conns.remove(addr)
|
|
|
|
buf := getBuffer()
|
|
defer releaseBuffer(buf)
|
|
|
|
_, err := io.CopyBuffer(&packetConn{local, addr}, remote, buf)
|
|
if err != nil {
|
|
log.Printf("%s: %v", addr, err)
|
|
}
|
|
}
|