147 lines
3.6 KiB
Go
147 lines
3.6 KiB
Go
package node
|
|
|
|
import (
|
|
"log"
|
|
"net/netip"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
)
|
|
|
|
// ----------------------------------------------------------------------------
|
|
|
|
type peerRoute struct {
|
|
IP byte
|
|
Up bool // True if data can be sent on the route.
|
|
Relay bool // True if the peer is a relay.
|
|
Direct bool // True if this is a direct connection.
|
|
PubSignKey []byte
|
|
ControlCipher *controlCipher
|
|
DataCipher *dataCipher
|
|
RemoteAddr netip.AddrPort // Remote address if directly connected.
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------
|
|
|
|
type udpAddrPortWriter interface {
|
|
WriteToUDPAddrPort([]byte, netip.AddrPort) (int, error)
|
|
}
|
|
|
|
type marshaller interface {
|
|
Marshal([]byte) []byte
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------
|
|
|
|
type connWriter struct {
|
|
localIP byte
|
|
conn udpAddrPortWriter
|
|
|
|
// For sending control packets.
|
|
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 newConnWriter(conn udpAddrPortWriter, localIP byte) *connWriter {
|
|
w := &connWriter{
|
|
localIP: localIP,
|
|
conn: conn,
|
|
cBuf1: make([]byte, bufferSize),
|
|
cBuf2: make([]byte, bufferSize),
|
|
dBuf1: make([]byte, bufferSize),
|
|
dBuf2: make([]byte, bufferSize),
|
|
}
|
|
for i := range w.counters {
|
|
w.counters[i] = uint64(time.Now().Unix()<<30 + 1)
|
|
}
|
|
return w
|
|
}
|
|
|
|
// Not safe for concurrent use. Should only be called by supervisor.
|
|
func (w *connWriter) SendControlPacket(pkt marshaller, route *peerRoute) {
|
|
buf := pkt.Marshal(w.cBuf1)
|
|
h := header{
|
|
StreamID: controlStreamID,
|
|
Counter: atomic.AddUint64(&w.counters[route.IP], 1),
|
|
SourceIP: w.localIP,
|
|
DestIP: route.IP,
|
|
}
|
|
buf = route.ControlCipher.Encrypt(h, buf, w.cBuf2)
|
|
w.writeTo(buf, route.RemoteAddr)
|
|
}
|
|
|
|
func (w *connWriter) RelayControlPacket(pkt marshaller, route, relay *peerRoute) {
|
|
buf := pkt.Marshal(w.cBuf1)
|
|
h := header{
|
|
StreamID: controlStreamID,
|
|
Counter: atomic.AddUint64(&w.counters[route.IP], 1),
|
|
SourceIP: w.localIP,
|
|
DestIP: route.IP,
|
|
}
|
|
buf = route.ControlCipher.Encrypt(h, buf, w.cBuf2)
|
|
w.relayPacket(buf, w.cBuf1, route, relay)
|
|
}
|
|
|
|
// Not safe for concurrent use. Should only be called by ifReader.
|
|
func (w *connWriter) SendDataPacket(pkt []byte, route, relay *peerRoute) {
|
|
h := header{
|
|
StreamID: dataStreamID,
|
|
Counter: atomic.AddUint64(&w.counters[route.IP], 1),
|
|
SourceIP: w.localIP,
|
|
DestIP: route.IP,
|
|
}
|
|
|
|
enc := route.DataCipher.Encrypt(h, pkt, w.dBuf1)
|
|
|
|
if route.Direct {
|
|
w.writeTo(enc, route.RemoteAddr)
|
|
return
|
|
}
|
|
|
|
w.relayPacket(enc, w.dBuf2, route, relay)
|
|
}
|
|
|
|
// TODO: RelayDataPacket
|
|
|
|
// Safe for concurrent use. Should only be called by connReader.
|
|
//
|
|
// This function will send pkt to the peer directly. This is used when a peer
|
|
// is acting as a relay and is forwarding already encrypted data for another
|
|
// peer.
|
|
func (w *connWriter) SendEncryptedDataPacket(pkt []byte, route *peerRoute) {
|
|
w.writeTo(pkt, route.RemoteAddr)
|
|
}
|
|
|
|
func (w *connWriter) relayPacket(data, buf []byte, route, relay *peerRoute) {
|
|
if relay == nil || !relay.Up {
|
|
return
|
|
}
|
|
|
|
h := header{
|
|
StreamID: dataStreamID,
|
|
Counter: atomic.AddUint64(&w.counters[relay.IP], 1),
|
|
SourceIP: w.localIP,
|
|
DestIP: route.IP,
|
|
}
|
|
|
|
enc := relay.DataCipher.Encrypt(h, data, buf)
|
|
w.writeTo(enc, relay.RemoteAddr)
|
|
}
|
|
|
|
func (w *connWriter) writeTo(packet []byte, addr netip.AddrPort) {
|
|
w.wLock.Lock()
|
|
if _, err := w.conn.WriteToUDPAddrPort(packet, addr); err != nil {
|
|
log.Printf("Failed to write to UDP port: %v", err)
|
|
}
|
|
w.wLock.Unlock()
|
|
}
|