feat(dns): implement DNS forwarding and resolution with fragmentation… (#234)

This commit is contained in:
Daniel Lavrushin 2026-06-01 20:55:50 +02:00 committed by GitHub
parent c63f5357c1
commit 8bdf134f7d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
15 changed files with 597 additions and 304 deletions

View file

@ -1,7 +1,8 @@
# B4 - Bye Bye Big Bro
## [1.64.0] - 2026-05-31
## [1.64.0] - 2026-06-01
- FIXED: **DNS redirect didn't work** - when a set sent its DNS lookups to a chosen server (the set's DNS tab), names failed to resolve and the internet seemed to hang. Redirected DNS now works as expected.
- IMPROVED: **MTProto Telegram routing is more reliable, especially for media** - the WebSocket routing (bridge mode) now works on phones as well as desktop, loads media more reliably, and always prefers WebSocket regardless of the proxy server's transport setting. Data centers without a WebSocket edge no longer stall when the shared fallback domains get rate-limited. New "Cloudflare Worker" support adds a free, per-user relay (no domain to buy) for accounts where media, reactions, or stickers still fail to load - see Settings -> MTProto Proxy.
- FIXED: **Copy buttons did nothing when the Web UI was opened over plain HTTP** - copying commands, set exports, or system info silently failed on installs reached by LAN IP (for example `http://192.168.1.1`), because browsers only allow the modern clipboard API over HTTPS or localhost. The fallback copy path now works in those cases too.
- FIXED: **Memory and CPU slowly climbed when the watchdog repeatedly repaired a site** - each watchdog repair left behind a background task that never stopped, piling up over days until b4 was killed. [#227](https://github.com/DanielLavrushin/b4/issues/227)

168
src/dns/forward.go Normal file
View file

@ -0,0 +1,168 @@
package dns
import (
"fmt"
"net"
"strconv"
"syscall"
"time"
"github.com/daniellavrushin/b4/sock"
"golang.org/x/sys/unix"
)
type ForwardOptions struct {
Sender *sock.Sender
Fragment bool
Seg2Delay int
ReverseOrder bool
Timeout time.Duration
Port int
Mark int
}
func ResolveUpstream(query []byte, target net.IP, opts ForwardOptions) ([]byte, error) {
port := opts.Port
if port == 0 {
port = 53
}
d := net.Dialer{}
if opts.Mark != 0 {
mark := opts.Mark
d.Control = func(_, _ string, c syscall.RawConn) error {
var serr error
if cerr := c.Control(func(fd uintptr) {
serr = unix.SetsockoptInt(int(fd), unix.SOL_SOCKET, unix.SO_MARK, mark)
}); cerr != nil {
return cerr
}
return serr
}
}
c, err := d.Dial("udp", net.JoinHostPort(target.String(), strconv.Itoa(port)))
if err != nil {
return nil, err
}
defer c.Close()
conn, ok := c.(*net.UDPConn)
if !ok {
return nil, fmt.Errorf("unexpected conn type %T", c)
}
sent := false
if opts.Fragment && opts.Sender != nil {
sent = sendFragmentedQuery(conn, query, target, port, opts)
}
if !sent {
if _, err := conn.Write(query); err != nil {
return nil, err
}
}
timeout := opts.Timeout
if timeout <= 0 {
timeout = 5 * time.Second
}
if err := conn.SetReadDeadline(time.Now().Add(timeout)); err != nil {
return nil, err
}
buf := make([]byte, 65535)
n, err := conn.Read(buf)
if err != nil {
return nil, err
}
return buf[:n], nil
}
func sendFragmentedQuery(conn *net.UDPConn, query []byte, target net.IP, port int, opts ForwardOptions) bool {
la, ok := conn.LocalAddr().(*net.UDPAddr)
if !ok {
return false
}
splitPos := findDNSSplitPoint(query)
if splitPos <= 0 {
splitPos = len(query) / 2
}
if t4 := target.To4(); t4 != nil {
pkt := sock.BuildUDPPacketV4(la.IP, target, uint16(la.Port), uint16(port), query)
if pkt == nil {
return false
}
frags, ok := sock.IPv4FragmentUDP(pkt, splitPos)
if !ok {
return false
}
return sendTwoFragments(opts.Sender, true, frags, target, opts.Seg2Delay, opts.ReverseOrder) == nil
}
if !opts.Sender.IPv6Ready() {
return false
}
pkt := sock.BuildUDPPacketV6(la.IP, target, uint16(la.Port), uint16(port), query)
if pkt == nil {
return false
}
frags, ok := sock.IPv6FragmentUDP(pkt, splitPos)
if !ok {
return false
}
return sendTwoFragments(opts.Sender, false, frags, target, opts.Seg2Delay, opts.ReverseOrder) == nil
}
func sendTwoFragments(s *sock.Sender, v4 bool, frags [][]byte, dst net.IP, delay int, reverse bool) error {
send := func(p []byte) error {
if v4 {
return s.SendIPv4(p, dst)
}
return s.SendIPv6(p, dst)
}
a, b := frags[0], frags[1]
if reverse {
a, b = b, a
}
if err := send(a); err != nil {
return err
}
if delay > 0 {
time.Sleep(time.Duration(delay) * time.Millisecond)
}
return send(b)
}
func findDNSSplitPoint(dnsPayload []byte) int {
if len(dnsPayload) < 13 {
return -1
}
pos := 12
qnameStart := pos
qnameEnd := pos
for pos < len(dnsPayload) {
labelLen := int(dnsPayload[pos])
if labelLen == 0 {
qnameEnd = pos + 1
break
}
if labelLen > 63 || pos+1+labelLen > len(dnsPayload) {
return len(dnsPayload) / 2
}
pos += 1 + labelLen
}
if qnameEnd <= qnameStart {
return len(dnsPayload) / 2
}
qnameLen := qnameEnd - qnameStart
if qnameLen > 4 {
return qnameStart + qnameLen/2
}
return len(dnsPayload) / 2
}

60
src/dns/forward_test.go Normal file
View file

@ -0,0 +1,60 @@
package dns
import (
"bytes"
"net"
"testing"
"time"
)
func TestResolveUpstreamNonFragment(t *testing.T) {
pc, err := net.ListenPacket("udp", "127.0.0.1:0")
if err != nil {
t.Fatalf("listen: %v", err)
}
defer pc.Close()
query := []byte{0x12, 0x34, 0x01, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00}
want := []byte{0x12, 0x34, 0x81, 0x80, 0x00, 0x01, 0x00, 0x01, 0xde, 0xad, 0xbe, 0xef}
go func() {
buf := make([]byte, 2048)
n, addr, err := pc.ReadFrom(buf)
if err != nil {
return
}
if !bytes.Equal(buf[:n], query) {
return
}
_, _ = pc.WriteTo(want, addr)
}()
la := pc.LocalAddr().(*net.UDPAddr)
resp, err := ResolveUpstream(query, net.IPv4(127, 0, 0, 1), ForwardOptions{
Port: la.Port,
Timeout: 2 * time.Second,
})
if err != nil {
t.Fatalf("ResolveUpstream: %v", err)
}
if !bytes.Equal(resp, want) {
t.Fatalf("response mismatch: got %x want %x", resp, want)
}
}
func TestResolveUpstreamTimeout(t *testing.T) {
pc, err := net.ListenPacket("udp", "127.0.0.1:0")
if err != nil {
t.Fatalf("listen: %v", err)
}
defer pc.Close()
la := pc.LocalAddr().(*net.UDPAddr)
_, err = ResolveUpstream([]byte{0, 0}, net.IPv4(127, 0, 0, 1), ForwardOptions{
Port: la.Port,
Timeout: 150 * time.Millisecond,
})
if err == nil {
t.Fatal("expected timeout error, got nil")
}
}

View file

@ -1,66 +0,0 @@
package dns
import (
"fmt"
"net"
"sync"
"time"
)
type dnsNATEntry struct {
originalDst net.IP
timestamp time.Time
}
var (
dnsNATTable = make(map[string]dnsNATEntry)
dnsNATMu sync.RWMutex
)
func dnsNATKey(ip net.IP, port uint16) string {
return fmt.Sprintf("%s:%d", ip.String(), port)
}
func DnsNATSet(clientIP net.IP, clientPort uint16, originalDst net.IP) {
dnsNATMu.Lock()
dnsNATTable[dnsNATKey(clientIP, clientPort)] = dnsNATEntry{
originalDst: originalDst,
timestamp: time.Now(),
}
dnsNATMu.Unlock()
}
func DnsNATGet(clientIP net.IP, clientPort uint16) (net.IP, bool) {
dnsNATMu.RLock()
entry, ok := dnsNATTable[dnsNATKey(clientIP, clientPort)]
dnsNATMu.RUnlock()
if !ok || time.Since(entry.timestamp) > 10*time.Second {
return nil, false
}
return entry.originalDst, true
}
func DnsNATDelete(clientIP net.IP, clientPort uint16) {
dnsNATMu.Lock()
delete(dnsNATTable, dnsNATKey(clientIP, clientPort))
dnsNATMu.Unlock()
}
func init() {
go dnsNATCleanupLoop()
}
func dnsNATCleanupLoop() {
ticker := time.NewTicker(30 * time.Second)
defer ticker.Stop()
for range ticker.C {
dnsNATMu.Lock()
now := time.Now()
for k, entry := range dnsNATTable {
if now.Sub(entry.timestamp) > 10*time.Second {
delete(dnsNATTable, k)
}
}
dnsNATMu.Unlock()
}
}

View file

@ -1,7 +1,6 @@
package nfq
import (
"encoding/binary"
"net"
"strings"
@ -9,7 +8,6 @@ import (
"github.com/daniellavrushin/b4/dns"
"github.com/daniellavrushin/b4/log"
"github.com/daniellavrushin/b4/sock"
"github.com/daniellavrushin/b4/utils"
"github.com/florianl/go-nfqueue"
)
@ -81,31 +79,10 @@ func (w *Worker) processDnsPacket(ipVersion byte, sport uint16, dport uint16, pa
dnsRouteKeyRequest(ipVersion, clientIP, sport, dnsServerIP, dport, txid, domain),
set.Id,
)
}
}
if set.DNS.Enabled && set.DNS.TargetDNS != "" {
if redirectIP := net.ParseIP(set.DNS.TargetDNS); redirectIP != nil {
storeDNSPendingRoute(
dnsRouteKeyRequest(ipVersion, clientIP, sport, redirectIP, dport, txid, domain),
set.Id,
)
}
}
}
}
dnsRedirect := set.DNS.Enabled && set.DNS.TargetDNS != ""
if dnsRedirect {
var dstIP net.IP
switch ipVersion {
case IPv4:
dstIP = net.IP(raw[16:20])
case IPv6:
dstIP = net.IP(raw[24:40])
}
if dstIP != nil && utils.IsPrivateIP(dstIP) {
dnsRedirect = false
}
}
if !dnsRedirect {
if !(set.DNS.Enabled && set.DNS.TargetDNS != "") {
if err := w.q.SetVerdict(id, nfqueue.NfAccept); err != nil {
log.Tracef("failed to set verdict on packet %d: %v", id, err)
}
@ -120,73 +97,38 @@ func (w *Worker) processDnsPacket(ipVersion byte, sport uint16, dport uint16, pa
return 0
}
if ipVersion == IPv4 {
targetDNS := targetIP.To4()
if targetDNS == nil {
if err := w.q.SetVerdict(id, nfqueue.NfAccept); err != nil {
log.Tracef("failed to set verdict on packet %d: %v", id, err)
}
return 0
if ipVersion == IPv6 && !cfg.Queue.IPv6Enabled {
if err := w.q.SetVerdict(id, nfqueue.NfAccept); err != nil {
log.Tracef("failed to set verdict on packet %d: %v", id, err)
}
originalDst := make(net.IP, 4)
copy(originalDst, raw[16:20])
if !cfg.Queue.IsDiscovery {
dns.DnsNATSet(net.IP(raw[12:16]), sport, originalDst)
}
copy(raw[16:20], targetDNS)
sock.FixIPv4Checksum(raw[:ihl])
sock.FixUDPChecksum(raw, ihl)
if set.DNS.FragmentQuery {
w.sendFragmentedDNSQueryV4(set, raw, ihl, targetDNS)
} else {
_ = w.sock.SendIPv4(raw, targetDNS)
}
if err := w.q.SetVerdict(id, nfqueue.NfDrop); err != nil {
log.Tracef("failed to set drop verdict on packet %d: %v", id, err)
}
log.Tracef("DNS redirect: %s -> %s (set: %s)", domain, set.DNS.TargetDNS, set.Name)
return 0
} else {
cfg := w.getConfig()
if !cfg.Queue.IPv6Enabled {
if err := w.q.SetVerdict(id, nfqueue.NfAccept); err != nil {
log.Tracef("failed to set verdict on packet %d: %v", id, err)
}
return 0
}
targetDNS := targetIP.To16()
if targetDNS == nil {
if err := w.q.SetVerdict(id, nfqueue.NfAccept); err != nil {
log.Tracef("failed to set verdict on packet %d: %v", id, err)
}
return 0
}
originalDst := make(net.IP, 16)
copy(originalDst, raw[24:40])
if !cfg.Queue.IsDiscovery {
dns.DnsNATSet(net.IP(raw[8:24]), sport, originalDst)
}
copy(raw[24:40], targetDNS)
sock.FixUDPChecksumV6(raw)
if set.DNS.FragmentQuery {
w.sendFragmentedDNSQueryV6(set, raw, targetDNS)
} else {
_ = w.sock.SendIPv6(raw, targetDNS)
}
if err := w.q.SetVerdict(id, nfqueue.NfDrop); err != nil {
log.Tracef("failed to set drop verdict on packet %d: %v", id, err)
}
log.Tracef("DNS redirect (IPv6): %s -> %s (set: %s)", domain, set.DNS.TargetDNS, set.Name)
return 0
}
var clientIP, originalDst net.IP
switch ipVersion {
case IPv4:
clientIP = append(net.IP(nil), raw[12:16]...)
originalDst = append(net.IP(nil), raw[16:20]...)
default:
clientIP = append(net.IP(nil), raw[8:24]...)
originalDst = append(net.IP(nil), raw[24:40]...)
}
query := append([]byte(nil), payload...)
ver := ipVersion
clientPort := sport
delay := config.ResolveSeg2Delay(set.UDP.Seg2Delay, set.UDP.Seg2DelayMax)
if err := w.q.SetVerdict(id, nfqueue.NfDrop); err != nil {
log.Tracef("failed to set drop verdict on packet %d: %v", id, err)
}
w.wg.Add(1)
go func(s *config.SetConfig, c *config.Config) {
defer w.wg.Done()
w.resolveDNSRedirect(ver, s, c, query, clientIP, clientPort, originalDst, targetIP, delay)
}(set, cfg)
return 0
}
}
}
@ -223,34 +165,6 @@ func (w *Worker) processDnsPacket(ipVersion byte, sport uint16, dport uint16, pa
}
}
}
if ipVersion == IPv4 {
if originalDst, ok := dns.DnsNATGet(net.IP(raw[16:20]), dport); ok {
copy(raw[12:16], originalDst.To4())
sock.FixIPv4Checksum(raw[:ihl])
sock.FixUDPChecksum(raw, ihl)
dns.DnsNATDelete(net.IP(raw[16:20]), dport)
_ = w.sock.SendIPv4(raw, net.IP(raw[16:20]))
if err := w.q.SetVerdict(id, nfqueue.NfDrop); err != nil {
log.Tracef("failed to set drop verdict on packet %d: %v", id, err)
}
return 0
}
} else {
cfg := w.getConfig()
if cfg.Queue.IPv6Enabled {
if originalDst, ok := dns.DnsNATGet(net.IP(raw[24:40]), dport); ok {
copy(raw[8:24], originalDst.To16())
sock.FixUDPChecksumV6(raw)
dns.DnsNATDelete(net.IP(raw[24:40]), dport)
_ = w.sock.SendIPv6(raw, net.IP(raw[24:40]))
if err := w.q.SetVerdict(id, nfqueue.NfDrop); err != nil {
log.Tracef("failed to set drop verdict on packet %d: %v", id, err)
}
return 0
}
}
}
}
if err := w.q.SetVerdict(id, nfqueue.NfAccept); err != nil {
@ -259,111 +173,34 @@ func (w *Worker) processDnsPacket(ipVersion byte, sport uint16, dport uint16, pa
return 0
}
func (w *Worker) sendFragmentedDNSQueryV4(cfg *config.SetConfig, raw []byte, ihl int, dst net.IP) {
udpOffset := ihl
if len(raw) < ihl+8 {
_ = w.sock.SendIPv4(raw, dst)
return
}
udpLen := int(binary.BigEndian.Uint16(raw[udpOffset+4 : udpOffset+6]))
if udpLen < 20 {
_ = w.sock.SendIPv4(raw, dst)
func (w *Worker) resolveDNSRedirect(ipVersion byte, set *config.SetConfig, cfg *config.Config, query []byte, clientIP net.IP, clientPort uint16, originalDst, targetIP net.IP, delay int) {
resp, err := dns.ResolveUpstream(query, targetIP, dns.ForwardOptions{
Sender: w.sock,
Fragment: set.DNS.FragmentQuery,
Seg2Delay: delay,
ReverseOrder: set.Fragmentation.ReverseOrder,
Mark: int(cfg.Queue.Mark),
})
if err != nil {
log.Tracef("DNS redirect: upstream %s failed: %v", set.DNS.TargetDNS, err)
return
}
dnsPayload := raw[udpOffset+8:]
if len(dnsPayload) < 12 {
_ = w.sock.SendIPv4(raw, dst)
return
}
splitPos := findDNSSplitPoint(dnsPayload)
if splitPos <= 0 {
splitPos = len(dnsPayload) / 2
}
frags, ok := sock.IPv4FragmentUDP(raw, splitPos)
if !ok {
log.Tracef("DNS frag: IP fragmentation failed, sending original")
_ = w.sock.SendIPv4(raw, dst)
return
}
seg2d := config.ResolveSeg2Delay(cfg.UDP.Seg2Delay, cfg.UDP.Seg2DelayMax)
w.SendTwoSegmentsV4(frags[0], frags[1], dst, seg2d, cfg.Fragmentation.ReverseOrder)
log.Tracef("DNS frag: sent %d fragments for query", len(frags))
}
func (w *Worker) sendFragmentedDNSQueryV6(cfg *config.SetConfig, raw []byte, dst net.IP) {
ipv6HdrLen := 40
if len(raw) < ipv6HdrLen+8 {
_ = w.sock.SendIPv6(raw, dst)
return
}
udpLen := int(binary.BigEndian.Uint16(raw[ipv6HdrLen+4 : ipv6HdrLen+6]))
if udpLen < 20 {
_ = w.sock.SendIPv6(raw, dst)
return
}
dnsPayload := raw[ipv6HdrLen+8:]
if len(dnsPayload) < 12 {
_ = w.sock.SendIPv6(raw, dst)
return
}
splitPos := findDNSSplitPoint(dnsPayload)
if splitPos <= 0 {
splitPos = len(dnsPayload) / 2
}
frags, ok := sock.IPv6FragmentUDP(raw, splitPos)
if !ok {
log.Tracef("DNS frag v6: fragmentation failed, sending original")
_ = w.sock.SendIPv6(raw, dst)
return
}
seg2d := config.ResolveSeg2Delay(cfg.UDP.Seg2Delay, cfg.UDP.Seg2DelayMax)
w.SendTwoSegmentsV6(frags[0], frags[1], dst, seg2d, cfg.Fragmentation.ReverseOrder)
log.Tracef("DNS frag v6: sent %d fragments", len(frags))
}
func findDNSSplitPoint(dnsPayload []byte) int {
if len(dnsPayload) < 13 {
return -1
}
pos := 12
qnameStart := pos
qnameEnd := pos
for pos < len(dnsPayload) {
labelLen := int(dnsPayload[pos])
if labelLen == 0 {
qnameEnd = pos + 1
break
if ipVersion == IPv4 {
if pkt := sock.BuildUDPPacketV4(originalDst, clientIP, 53, clientPort, resp); pkt != nil {
_ = w.sock.SendIPv4(pkt, clientIP)
}
if labelLen > 63 || pos+1+labelLen > len(dnsPayload) {
return len(dnsPayload) / 2
} else {
if pkt := sock.BuildUDPPacketV6(originalDst, clientIP, 53, clientPort, resp); pkt != nil {
_ = w.sock.SendIPv6(pkt, clientIP)
}
pos += 1 + labelLen
}
if qnameEnd <= qnameStart {
return len(dnsPayload) / 2
if set.Routing.Enabled && !cfg.Queue.IsDiscovery && RoutingHandleDNSFunc != nil {
if ips := dns.ParseResponseIPs(resp); len(ips) > 0 {
RoutingHandleDNSFunc(cfg, set, ips)
}
}
qnameLen := qnameEnd - qnameStart
if qnameLen > 4 {
return qnameStart + qnameLen/2
}
return len(dnsPayload) / 2
log.Tracef("DNS redirect: %s -> %s answered for %s (set: %s)", originalDst, set.DNS.TargetDNS, clientIP, set.Name)
}

View file

@ -9,6 +9,7 @@ import (
"github.com/daniellavrushin/b4/config"
"github.com/daniellavrushin/b4/log"
"github.com/daniellavrushin/b4/sock"
"github.com/daniellavrushin/b4/utils"
)
const (
@ -226,7 +227,7 @@ func (w *Worker) reorderExtensions(packet []byte, cfg *config.SetConfig) []byte
// Random shuffle other extensions
for i := len(otherExts) - 1; i > 0; i-- {
j := int(randomUint32() % uint32(i+1))
j := int(utils.RandUint32() % uint32(i+1))
otherExts[i], otherExts[j] = otherExts[j], otherExts[i]
}
@ -721,12 +722,3 @@ func (w *Worker) updatePacketLengths(packet []byte) {
sock.FixIPv4Checksum(packet[:ipHdrLen])
sock.FixTCPChecksum(packet)
}
// randomUint32 generates a random uint32 with thread safety
func randomUint32() uint32 {
var b [4]byte
randMutex.Lock()
rand.Read(b[:])
randMutex.Unlock()
return binary.BigEndian.Uint32(b[:])
}

View file

@ -7,6 +7,7 @@ import (
"github.com/daniellavrushin/b4/config"
"github.com/daniellavrushin/b4/log"
"github.com/daniellavrushin/b4/utils"
)
const ipv6HdrLen = 40
@ -182,7 +183,7 @@ func (w *Worker) reorderExtensionsv6(packet []byte, cfg *config.SetConfig) []byt
}
for i := len(otherExts) - 1; i > 0; i-- {
j := int(randomUint32() % uint32(i+1))
j := int(utils.RandUint32() % uint32(i+1))
otherExts[i], otherExts[j] = otherExts[j], otherExts[i]
}

View file

@ -1,15 +1,10 @@
package sock
import (
"crypto/rand"
"encoding/binary"
)
func generateFragmentID() uint32 {
var buf [4]byte
rand.Read(buf[:])
return binary.BigEndian.Uint32(buf[:])
}
"github.com/daniellavrushin/b4/utils"
)
// IPv6FragmentPacket creates IPv6 fragments using Fragment extension headers
// This implements true IPv6 fragmentation (IP-level)
@ -41,7 +36,7 @@ func IPv6FragmentPacket(packet []byte, splitPos int) ([][]byte, bool) {
}
fragHdrLen := 8
var identification uint32 = generateFragmentID()
var identification uint32 = utils.RandUint32()
// First fragment
frag1Len := ipv6HdrLen + fragHdrLen + splitPos

View file

@ -72,6 +72,10 @@ func (s *Sender) SendIPv6(packet []byte, destIP net.IP) error {
return syscall.Sendto(s.fd6, packet, 0, &addr)
}
func (s *Sender) IPv6Ready() bool {
return s.fd6 >= 0
}
func (s *Sender) Close() {
if s.fd4 >= 0 {
_ = syscall.Close(s.fd4)

View file

@ -2,8 +2,42 @@ package sock
import (
"encoding/binary"
"net"
"github.com/daniellavrushin/b4/utils"
)
func BuildUDPPacketV4(srcIP, dstIP net.IP, srcPort, dstPort uint16, payload []byte) []byte {
src := srcIP.To4()
dst := dstIP.To4()
if src == nil || dst == nil {
return nil
}
if len(payload) > 0xffff-28 {
return nil
}
total := 20 + 8 + len(payload)
pkt := make([]byte, total)
pkt[0] = 0x45
binary.BigEndian.PutUint16(pkt[2:4], uint16(total))
binary.BigEndian.PutUint16(pkt[4:6], utils.RandUint16())
pkt[8] = 64
pkt[9] = 17
copy(pkt[12:16], src)
copy(pkt[16:20], dst)
binary.BigEndian.PutUint16(pkt[20:22], srcPort)
binary.BigEndian.PutUint16(pkt[22:24], dstPort)
binary.BigEndian.PutUint16(pkt[24:26], uint16(8+len(payload)))
copy(pkt[28:], payload)
FixIPv4Checksum(pkt[:20])
FixUDPChecksum(pkt, 20)
return pkt
}
func udpChecksumIPv4(pkt []byte) {
ihl := int((pkt[0] & 0x0f) << 2)
udpo := ihl

View file

@ -2,9 +2,72 @@ package sock
import (
"encoding/binary"
"net"
"testing"
)
func onesSum(b []byte) uint16 {
var sum uint32
for i := 0; i+1 < len(b); i += 2 {
sum += uint32(binary.BigEndian.Uint16(b[i : i+2]))
}
if len(b)%2 == 1 {
sum += uint32(b[len(b)-1]) << 8
}
for sum > 0xffff {
sum = (sum >> 16) + (sum & 0xffff)
}
return uint16(sum)
}
func TestBuildUDPPacketV4PayloadBounds(t *testing.T) {
src := net.IPv4(8, 8, 8, 8)
dst := net.IPv4(192, 168, 1, 50)
if pkt := BuildUDPPacketV4(src, dst, 53, 40000, make([]byte, 0xffff-28)); pkt == nil {
t.Error("expected non-nil at max payload size")
}
if pkt := BuildUDPPacketV4(src, dst, 53, 40000, make([]byte, 0xffff-27)); pkt != nil {
t.Error("expected nil for oversized payload (would overflow IPv4 total length)")
}
}
func TestBuildUDPPacketV4Checksums(t *testing.T) {
src := net.IPv4(8, 8, 8, 8)
dst := net.IPv4(192, 168, 1, 50)
payload := []byte{0x12, 0x34, 0x81, 0x80, 0x00, 0x01, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x07}
pkt := BuildUDPPacketV4(src, dst, 53, 40000, payload)
if pkt == nil {
t.Fatal("BuildUDPPacketV4 returned nil")
}
if len(pkt) != 20+8+len(payload) {
t.Fatalf("unexpected length %d", len(pkt))
}
if pkt[9] != 17 {
t.Fatalf("protocol = %d, want 17", pkt[9])
}
if !net.IP(pkt[12:16]).Equal(src.To4()) {
t.Fatalf("src mismatch: %v", net.IP(pkt[12:16]))
}
if !net.IP(pkt[16:20]).Equal(dst.To4()) {
t.Fatalf("dst mismatch: %v", net.IP(pkt[16:20]))
}
if got := onesSum(pkt[:20]); got != 0xffff {
t.Fatalf("ip header checksum invalid: folded sum = %#x", got)
}
pseudo := make([]byte, 12)
copy(pseudo[0:4], pkt[12:16])
copy(pseudo[4:8], pkt[16:20])
pseudo[9] = 17
binary.BigEndian.PutUint16(pseudo[10:12], uint16(8+len(payload)))
if got := onesSum(append(pseudo, pkt[20:]...)); got != 0xffff {
t.Fatalf("udp checksum invalid: folded sum = %#x", got)
}
}
func buildMinimalIPv4UDPPacket(payloadSize int) []byte {
ipHdrLen := 20
udpHdrLen := 8

View file

@ -2,8 +2,39 @@ package sock
import (
"encoding/binary"
"net"
"github.com/daniellavrushin/b4/utils"
)
func BuildUDPPacketV6(srcIP, dstIP net.IP, srcPort, dstPort uint16, payload []byte) []byte {
src := srcIP.To16()
dst := dstIP.To16()
if src == nil || dst == nil || srcIP.To4() != nil || dstIP.To4() != nil {
return nil
}
if len(payload) > 0xffff-8 {
return nil
}
pkt := make([]byte, 40+8+len(payload))
pkt[0] = 0x60
binary.BigEndian.PutUint16(pkt[4:6], uint16(8+len(payload)))
pkt[6] = 17
pkt[7] = 64
copy(pkt[8:24], src)
copy(pkt[24:40], dst)
binary.BigEndian.PutUint16(pkt[40:42], srcPort)
binary.BigEndian.PutUint16(pkt[42:44], dstPort)
binary.BigEndian.PutUint16(pkt[44:46], uint16(8+len(payload)))
copy(pkt[48:], payload)
FixUDPChecksumV6(pkt)
return pkt
}
func udpChecksumIPv6(pkt []byte) {
if len(pkt) < 48 {
return
@ -113,7 +144,7 @@ func IPv6FragmentUDP(orig []byte, split int) ([][]byte, bool) {
fragHdrLen := 8
var identification uint32 = generateFragmentID()
var identification uint32 = utils.RandUint32()
frag1Len := ipv6HdrLen + fragHdrLen + firstDataAligned
frag1 := make([]byte, frag1Len)

View file

@ -2,9 +2,54 @@ package sock
import (
"encoding/binary"
"net"
"testing"
)
func TestBuildUDPPacketV6PayloadBounds(t *testing.T) {
src := net.ParseIP("2001:db8::1")
dst := net.ParseIP("2001:db8::2")
if pkt := BuildUDPPacketV6(src, dst, 53, 40000, make([]byte, 0xffff-8)); pkt == nil {
t.Error("expected non-nil at max payload size")
}
if pkt := BuildUDPPacketV6(src, dst, 53, 40000, make([]byte, 0xffff-7)); pkt != nil {
t.Error("expected nil for oversized payload (would overflow UDP/payload length)")
}
}
func TestBuildUDPPacketV6Checksums(t *testing.T) {
src := net.ParseIP("2001:db8::1")
dst := net.ParseIP("2001:db8::2")
payload := []byte{0xab, 0xcd, 0x01, 0x00, 0x00, 0x01}
pkt := BuildUDPPacketV6(src, dst, 53, 40000, payload)
if pkt == nil {
t.Fatal("BuildUDPPacketV6 returned nil")
}
if len(pkt) != 40+8+len(payload) {
t.Fatalf("unexpected length %d", len(pkt))
}
if pkt[6] != 17 {
t.Fatalf("next header = %d, want 17", pkt[6])
}
pseudo := make([]byte, 40)
copy(pseudo[0:16], pkt[8:24])
copy(pseudo[16:32], pkt[24:40])
binary.BigEndian.PutUint32(pseudo[32:36], uint32(8+len(payload)))
pseudo[39] = 17
if got := onesSum(append(pseudo, pkt[40:]...)); got != 0xffff {
t.Fatalf("udp6 checksum invalid: folded sum = %#x", got)
}
}
func TestBuildUDPPacketV6RejectsV4(t *testing.T) {
if pkt := BuildUDPPacketV6(net.IPv4(1, 1, 1, 1), net.ParseIP("2001:db8::2"), 53, 40000, nil); pkt != nil {
t.Fatal("expected nil for v4 source")
}
}
func buildMinimalIPv6UDPPacket(payloadSize int) []byte {
ipv6HdrLen := 40
udpHdrLen := 8

View file

@ -1,6 +1,8 @@
package utils
import (
crand "crypto/rand"
"encoding/binary"
"math/rand"
"time"
)
@ -8,3 +10,15 @@ import (
func NewRand() *rand.Rand {
return rand.New(rand.NewSource(time.Now().UnixNano()))
}
func RandUint16() uint16 {
var b [2]byte
_, _ = crand.Read(b[:])
return binary.BigEndian.Uint16(b[:])
}
func RandUint32() uint32 {
var b [4]byte
_, _ = crand.Read(b[:])
return binary.BigEndian.Uint32(b[:])
}

View file

@ -1,6 +1,9 @@
package utils
import "testing"
import (
"net"
"testing"
)
func TestFilterUniqueStrings(t *testing.T) {
tests := []struct {
@ -119,6 +122,117 @@ func TestValidatePorts(t *testing.T) {
}
}
func TestRandUint16(t *testing.T) {
const samples = 2000
var orBits, andBits uint16 = 0x0000, 0xffff
first := RandUint16()
allSame := true
for i := 0; i < samples; i++ {
v := RandUint16()
orBits |= v
andBits &= v
if v != first {
allSame = false
}
}
if allSame {
t.Fatal("RandUint16 returned the same value every time")
}
if orBits != 0xffff {
t.Errorf("RandUint16 never set some bits across %d samples: OR=%#04x (byte-packing bug?)", samples, orBits)
}
if andBits != 0x0000 {
t.Errorf("RandUint16 had bits always set across %d samples: AND=%#04x", samples, andBits)
}
}
func TestRandUint32(t *testing.T) {
const samples = 4000
var orBits, andBits uint32 = 0x00000000, 0xffffffff
first := RandUint32()
allSame := true
for i := 0; i < samples; i++ {
v := RandUint32()
orBits |= v
andBits &= v
if v != first {
allSame = false
}
}
if allSame {
t.Fatal("RandUint32 returned the same value every time")
}
if orBits != 0xffffffff {
t.Errorf("RandUint32 never set some bits across %d samples: OR=%#08x (byte-packing bug?)", samples, orBits)
}
if andBits != 0x00000000 {
t.Errorf("RandUint32 had bits always set across %d samples: AND=%#08x", samples, andBits)
}
}
func TestSlicesAreEqual(t *testing.T) {
tests := []struct {
name string
a, b []string
expected bool
}{
{name: "both empty", a: []string{}, b: []string{}, expected: true},
{name: "both nil", a: nil, b: nil, expected: true},
{name: "identical", a: []string{"a", "b", "c"}, b: []string{"a", "b", "c"}, expected: true},
{name: "same elements different order", a: []string{"a", "b", "c"}, b: []string{"c", "a", "b"}, expected: true},
{name: "different length", a: []string{"a", "b"}, b: []string{"a", "b", "c"}, expected: false},
{name: "same length different element", a: []string{"a", "b"}, b: []string{"a", "x"}, expected: false},
{name: "empty vs non-empty", a: []string{}, b: []string{"a"}, expected: false},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if got := SlicesAreEqual(tt.a, tt.b); got != tt.expected {
t.Errorf("SlicesAreEqual(%v, %v) = %v, want %v", tt.a, tt.b, got, tt.expected)
}
})
}
}
func TestIsPrivateIP(t *testing.T) {
tests := []struct {
name string
ip string
expected bool
}{
{name: "10.x", ip: "10.0.0.1", expected: true},
{name: "172.16 low boundary", ip: "172.16.0.1", expected: true},
{name: "172.31 high boundary", ip: "172.31.255.255", expected: true},
{name: "172.15 below range", ip: "172.15.0.1", expected: false},
{name: "172.32 above range", ip: "172.32.0.1", expected: false},
{name: "192.168", ip: "192.168.1.1", expected: true},
{name: "192.167 not private", ip: "192.167.1.1", expected: false},
{name: "loopback v4", ip: "127.0.0.1", expected: true},
{name: "public v4 8.8.8.8", ip: "8.8.8.8", expected: false},
{name: "public v4 1.1.1.1", ip: "1.1.1.1", expected: false},
{name: "loopback v6", ip: "::1", expected: true},
{name: "link-local v6", ip: "fe80::1", expected: true},
{name: "ULA v6", ip: "fd00::1", expected: true},
{name: "public v6", ip: "2001:4860:4860::8888", expected: false},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
ip := net.ParseIP(tt.ip)
if ip == nil {
t.Fatalf("could not parse %q", tt.ip)
}
if got := IsPrivateIP(ip); got != tt.expected {
t.Errorf("IsPrivateIP(%s) = %v, want %v", tt.ip, got, tt.expected)
}
})
}
if IsPrivateIP(nil) {
t.Error("IsPrivateIP(nil) = true, want false")
}
}
func BenchmarkFilterUniqueStrings(b *testing.B) {
input := make([]string, 1000)
for i := 0; i < 1000; i++ {