110 lines
2.3 KiB
Go
110 lines
2.3 KiB
Go
package peer
|
|
|
|
import (
|
|
"log"
|
|
"net/netip"
|
|
"sync"
|
|
"sync/atomic"
|
|
)
|
|
|
|
type ConnWriter struct {
|
|
wLock sync.Mutex // Lock around for sending on UDP Conn.
|
|
|
|
// Output.
|
|
writeToUDPAddrPort func([]byte, netip.AddrPort) (int, error)
|
|
|
|
// Shared state.
|
|
rt *atomic.Pointer[RoutingTable]
|
|
|
|
// For sending control packets.
|
|
cBuf1 []byte
|
|
cBuf2 []byte
|
|
|
|
// For sending data packets.
|
|
dBuf1 []byte
|
|
dBuf2 []byte
|
|
}
|
|
|
|
func NewConnWriter(
|
|
writeToUDPAddrPort func([]byte, netip.AddrPort) (int, error),
|
|
rt *atomic.Pointer[RoutingTable],
|
|
) *ConnWriter {
|
|
return &ConnWriter{
|
|
writeToUDPAddrPort: writeToUDPAddrPort,
|
|
rt: rt,
|
|
cBuf1: newBuf(),
|
|
cBuf2: newBuf(),
|
|
dBuf1: newBuf(),
|
|
dBuf2: newBuf(),
|
|
}
|
|
}
|
|
|
|
// Called by ConnReader to forward already encrypted bytes to another peer.
|
|
func (w *ConnWriter) Forward(ip byte, pkt []byte) {
|
|
peer := w.rt.Load().Peers[ip]
|
|
if !(peer.Up && peer.Direct) {
|
|
w.logf("Failed to forward to %d.", ip)
|
|
return
|
|
}
|
|
w.writeTo(pkt, peer.DirectAddr)
|
|
}
|
|
|
|
// Called by IFReader to send data. Encryption will be applied, and packet will
|
|
// be relayed if appropriate.
|
|
func (w *ConnWriter) WriteData(ip byte, pkt []byte) {
|
|
rt := w.rt.Load()
|
|
peer := rt.Peers[ip]
|
|
if !peer.Up {
|
|
w.logf("Failed to send data to %d.", ip)
|
|
return
|
|
}
|
|
|
|
enc := peer.EncryptDataPacket(ip, pkt, w.dBuf1)
|
|
|
|
if peer.Direct {
|
|
w.writeTo(enc, peer.DirectAddr)
|
|
return
|
|
}
|
|
|
|
relay, ok := rt.GetRelay()
|
|
if !ok {
|
|
w.logf("Failed to send data to %d. No relay.", ip)
|
|
return
|
|
}
|
|
|
|
enc = relay.EncryptDataPacket(ip, enc, w.dBuf2)
|
|
w.writeTo(enc, relay.DirectAddr)
|
|
}
|
|
|
|
// Called by Supervisor to send control packets.
|
|
func (w *ConnWriter) WriteControl(peer RemotePeer, pkt Marshaller) {
|
|
enc := peer.EncryptControlPacket(pkt, w.cBuf2, w.cBuf1)
|
|
|
|
if peer.Direct {
|
|
w.writeTo(enc, peer.DirectAddr)
|
|
return
|
|
}
|
|
|
|
rt := w.rt.Load()
|
|
relay, ok := rt.GetRelay()
|
|
if !ok {
|
|
w.logf("Failed to send control to %d. No relay.", peer.IP)
|
|
return
|
|
}
|
|
|
|
enc = relay.EncryptDataPacket(peer.IP, enc, w.cBuf2)
|
|
w.writeTo(enc, relay.DirectAddr)
|
|
}
|
|
|
|
func (w *ConnWriter) writeTo(pkt []byte, addr netip.AddrPort) {
|
|
w.wLock.Lock()
|
|
if _, err := w.writeToUDPAddrPort(pkt, addr); err != nil {
|
|
w.logf("Failed to write to UDP port: %v", err)
|
|
}
|
|
w.wLock.Unlock()
|
|
}
|
|
|
|
func (w *ConnWriter) logf(s string, args ...any) {
|
|
log.Printf("[ConnWriter] "+s, args...)
|
|
}
|