128 lines
2.8 KiB
Go
128 lines
2.8 KiB
Go
package node
|
|
|
|
import (
|
|
"log"
|
|
"net"
|
|
"net/netip"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
)
|
|
|
|
type controlPacketSender interface {
|
|
SendControlPacket(pkt marshaller, route peerRoute)
|
|
}
|
|
|
|
type dataPacketSender interface {
|
|
SendDataPacket(pkt []byte, route peerRoute)
|
|
SendEncryptedDataPacket(pkt []byte, addr netip.AddrPort)
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------
|
|
|
|
type packetSender struct {
|
|
conn *net.UDPConn
|
|
|
|
// For sending control packets.
|
|
cLock sync.Mutex
|
|
cBuf1 []byte
|
|
cBuf2 []byte
|
|
|
|
// For sending data packets.
|
|
dBuf1 []byte
|
|
dBuf2 []byte
|
|
|
|
counters [256]uint64
|
|
|
|
// Lock around for sending on UDP Conn.
|
|
wLock sync.Mutex
|
|
}
|
|
|
|
func newPacketSender(conn *net.UDPConn) *packetSender {
|
|
ps := &packetSender{
|
|
conn: conn,
|
|
cBuf1: make([]byte, bufferSize),
|
|
cBuf2: make([]byte, bufferSize),
|
|
dBuf1: make([]byte, bufferSize),
|
|
dBuf2: make([]byte, bufferSize),
|
|
}
|
|
for i := range ps.counters {
|
|
ps.counters[i] = uint64(time.Now().Unix()<<30 + 1)
|
|
}
|
|
return ps
|
|
}
|
|
|
|
// Safe for concurrent use.
|
|
func (sender *packetSender) SendControlPacket(pkt marshaller, route peerRoute) {
|
|
sender.cLock.Lock()
|
|
defer sender.cLock.Unlock()
|
|
|
|
buf := pkt.Marshal(sender.cBuf1)
|
|
h := header{
|
|
StreamID: controlStreamID,
|
|
Counter: atomic.AddUint64(&sender.counters[route.IP], 1),
|
|
SourceIP: localIP,
|
|
DestIP: route.IP,
|
|
}
|
|
buf = route.ControlCipher.Encrypt(h, buf, sender.cBuf2)
|
|
|
|
if route.Direct {
|
|
sender.writeTo(buf, route.RemoteAddr)
|
|
return
|
|
}
|
|
|
|
sender.relayPacket(route.IP, buf, sender.cBuf1)
|
|
}
|
|
|
|
// Not safe for concurrent use.
|
|
func (sender *packetSender) SendDataPacket(pkt []byte, route peerRoute) {
|
|
h := header{
|
|
StreamID: dataStreamID,
|
|
Counter: atomic.AddUint64(&sender.counters[route.IP], 1),
|
|
SourceIP: localIP,
|
|
DestIP: route.IP,
|
|
}
|
|
|
|
enc := route.DataCipher.Encrypt(h, pkt, sender.dBuf1)
|
|
|
|
if route.Direct {
|
|
sender.writeTo(enc, route.RemoteAddr)
|
|
return
|
|
}
|
|
|
|
sender.relayPacket(route.IP, enc, sender.dBuf2)
|
|
}
|
|
|
|
func (sender *packetSender) SendEncryptedDataPacket(pkt []byte, addr netip.AddrPort) {
|
|
sender.writeTo(pkt, addr)
|
|
}
|
|
|
|
func (sender *packetSender) relayPacket(destIP byte, data, buf []byte) {
|
|
ip := relayIP.Load()
|
|
if ip == nil {
|
|
return
|
|
}
|
|
relayRoute := routingTable[*ip].Load()
|
|
if relayRoute == nil || !relayRoute.Up || !relayRoute.Relay {
|
|
return
|
|
}
|
|
|
|
h := header{
|
|
StreamID: dataStreamID,
|
|
Counter: atomic.AddUint64(&sender.counters[relayRoute.IP], 1),
|
|
SourceIP: localIP,
|
|
DestIP: destIP,
|
|
}
|
|
|
|
enc := relayRoute.DataCipher.Encrypt(h, data, buf)
|
|
sender.writeTo(enc, relayRoute.RemoteAddr)
|
|
}
|
|
|
|
func (sender *packetSender) writeTo(packet []byte, addr netip.AddrPort) {
|
|
sender.wLock.Lock()
|
|
if _, err := sender.conn.WriteToUDPAddrPort(packet, addr); err != nil {
|
|
log.Printf("Failed to write to UDP port: %v", err)
|
|
}
|
|
sender.wLock.Unlock()
|
|
}
|