gapmind/services/dht/dht_node.go

package main

import (
	"bytes"
	"crypto/rand"
	"encoding/binary"
	"fmt"
	"net"
	"sort"
	"sync"
	"time"
)

type SniffEvent struct {
	InfoHash string
	IP       string
	Port     int
	Query    string
	NodeID   string
	At       time.Time
}

type DHTNode struct {
	conn    *net.UDPConn
	selfID  []byte
	mu      sync.Mutex
	seen    map[string]time.Time
	onSniff func(SniffEvent)
}

var bootstrapNodes = []string{
	"router.bittorrent.com:6881",
	"router.utorrent.com:6881",
	"dht.transmissionbt.com:6881",
	"router.bitcomet.com:6881",
}

func bEncode(v any) []byte {
	switch t := v.(type) {
	case string:
		return []byte(fmt.Sprintf("%d:%s", len(t), t))
	case []byte:
		return []byte(fmt.Sprintf("%d:%s", len(t), string(t)))
	case int:
		return []byte(fmt.Sprintf("i%de", t))
	case int64:
		return []byte(fmt.Sprintf("i%de", t))
	case map[string]any:
		keys := make([]string, 0, len(t))
		for k := range t {
			keys = append(keys, k)
		}
		sort.Strings(keys)
		var b bytes.Buffer
		b.WriteByte('d')
		for _, k := range keys {
			b.Write(bEncode(k))
			b.Write(bEncode(t[k]))
		}
		b.WriteByte('e')
		return b.Bytes()
	default:
		return []byte{}
	}
}

type decoder struct {
	b   []byte
	pos int
}

func bDecodeDict(b []byte) (map[string]any, bool) {
	d := &decoder{b: b}
	v, ok := d.value()
	m, ok2 := v.(map[string]any)
	return m, ok && ok2
}

func (d *decoder) value() (any, bool) {
	if d.pos >= len(d.b) {
		return nil, false
	}
	c := d.b[d.pos]
	switch {
	case c == 'd':
		d.pos++
		m := map[string]any{}
		for d.pos < len(d.b) && d.b[d.pos] != 'e' {
			ks, ok := d.bytes()
			if !ok {
				return nil, false
			}
			v, ok := d.value()
			if !ok {
				return nil, false
			}
			m[string(ks)] = v
		}
		if d.pos >= len(d.b) {
			return nil, false
		}
		d.pos++
		return m, true
	case c == 'l':
		d.pos++
		var l []any
		for d.pos < len(d.b) && d.b[d.pos] != 'e' {
			v, ok := d.value()
			if !ok {
				return nil, false
			}
			l = append(l, v)
		}
		if d.pos >= len(d.b) {
			return nil, false
		}
		d.pos++
		return l, true
	case c == 'i':
		d.pos++
		end := bytes.IndexByte(d.b[d.pos:], 'e')
		if end < 0 {
			return nil, false
		}
		var n int64
		_, _ = fmt.Sscanf(string(d.b[d.pos:d.pos+end]), "%d", &n)
		d.pos += end + 1
		return n, true
	case c >= '0' && c <= '9':
		b, ok := d.bytes()
		return b, ok
	}
	return nil, false
}

func (d *decoder) bytes() ([]byte, bool) {
	colon := bytes.IndexByte(d.b[d.pos:], ':')
	if colon < 0 {
		return nil, false
	}
	var n int
	_, _ = fmt.Sscanf(string(d.b[d.pos:d.pos+colon]), "%d", &n)
	d.pos += colon + 1
	if n < 0 || d.pos+n > len(d.b) {
		return nil, false
	}
	out := make([]byte, n)
	copy(out, d.b[d.pos:d.pos+n])
	d.pos += n
	return out, true
}

func dictStr(m map[string]any, key string) []byte {
	if a, ok := m[key].(map[string]any); ok {
		_ = a
	}
	if b, ok := m[key].([]byte); ok {
		return b
	}
	return nil
}

func ListenDHT(listen string, onSniff func(SniffEvent)) (*DHTNode, error) {
	addr, err := net.ResolveUDPAddr("udp", listen)
	if err != nil {
		return nil, err
	}
	conn, err := net.ListenUDP("udp", addr)
	if err != nil {
		return nil, err
	}
	id := make([]byte, 20)
	_, _ = rand.Read(id)
	n := &DHTNode{conn: conn, selfID: id, seen: map[string]time.Time{}, onSniff: onSniff}
	go n.loop()
	go n.bootstrap()
	go n.crawlLoop()
	return n, nil
}

func (n *DHTNode) loop() {
	buf := make([]byte, 8192)
	for {
		m, addr, err := n.conn.ReadFromUDP(buf)
		if err != nil {
			continue
		}
		pkt := make([]byte, m)
		copy(pkt, buf[:m])
		go n.handle(pkt, addr)
	}
}

func (n *DHTNode) handle(pkt []byte, addr *net.UDPAddr) {
	msg, ok := bDecodeDict(pkt)
	if !ok {
		return
	}
	y, _ := msg["y"].([]byte)
	if string(y) != "q" {
		return
	}
	q, _ := msg["q"].([]byte)
	a, _ := msg["a"].(map[string]any)
	t := dictStr(msg, "t")
	if len(t) == 0 {
		t = []byte("aa")
	}
	querierID := ""
	if a != nil {
		if id := dictStr(a, "id"); len(id) == 20 {
			querierID = fmt.Sprintf("%x", id)
		}
	}
	reply := func(extra map[string]any) {
		r := map[string]any{"id": n.selfID}
		for k, v := range extra {
			r[k] = v
		}
		resp := map[string]any{"t": string(t), "y": "r", "r": r}
		_, _ = n.conn.WriteToUDP(bEncode(resp), addr)
	}
	switch string(q) {
	case "ping":
		reply(map[string]any{})
	case "find_node":
		reply(map[string]any{"nodes": []byte{}})

		go n.findNode(addr, dictStr(a, "target"))
	case "get_peers":
		ih := dictStr(a, "info_hash")
		if len(ih) == 20 {
			n.sniff("get_peers", ih, querierID, addr)
		}

		reply(map[string]any{"nodes": []byte{}, "token": "gapmind", "values": []any{}})
	case "announce_peer":
		ih := dictStr(a, "info_hash")
		if len(ih) == 20 {
			n.sniff("announce_peer", ih, querierID, addr)
		}
		reply(map[string]any{})
	case "sample_infohashes":
		reply(map[string]any{"interval": 180, "nodes": []byte{}, "samples": []byte{}})
	}
}

func (n *DHTNode) sniff(query string, infoHash []byte, querier string, addr *net.UDPAddr) {
	hexHash := fmt.Sprintf("%x", infoHash)
	n.mu.Lock()
	last, dup := n.seen[hexHash]
	if !dup || time.Since(last) > 5*time.Minute {
		n.seen[hexHash] = time.Now()
		dup = false
	}
	n.mu.Unlock()
	if dup {
		return
	}
	if n.onSniff != nil {
		n.onSniff(SniffEvent{InfoHash: hexHash, IP: addr.IP.String(), Port: addr.Port, Query: query, NodeID: querier, At: time.Now()})
	}
}

func (n *DHTNode) bootstrap() {
	for _, b := range bootstrapNodes {
		addr, err := net.ResolveUDPAddr("udp", b)
		if err != nil {
			continue
		}
		n.findNode(addr, n.selfID)
	}
}

func (n *DHTNode) crawlLoop() {
	t := time.NewTicker(60 * time.Second)
	defer t.Stop()
	for range t.C {
		n.bootstrap()
	}
}

func (n *DHTNode) findNode(addr *net.UDPAddr, target []byte) {
	if len(target) != 20 {
		target = n.selfID
	}
	msg := map[string]any{"t": "fn", "y": "q", "q": "find_node",
		"a": map[string]any{"id": n.selfID, "target": string(target)}}
	_, _ = n.conn.WriteToUDP(bEncode(msg), addr)
}

func (n *DHTNode) Close() { _ = n.conn.Close() }

var _ = binary.BigEndian