node.go

package main

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"log"
	"math/rand"
	"net"
	"sync"
	"time"
)

const heartbeatInterval = 5 * time.Second

type ChatMsg struct {
	From    uint16
	Time    time.Time
	Content string
}

type Node struct {
	id        uint16
	port      uint16
	conn      *net.UDPConn
	peers     *PeerManager
	quit      chan struct{}
	onChange  func()
	bootstrap Addr

	chatMu   sync.Mutex
	chatMsgs []ChatMsg
	chatBlm  *BloomFilterManager
	chatSeq  uint16
}

func NewNode(port uint16, onChange func()) (*Node, error) {
	conn, err := net.ListenUDP("udp", &net.UDPAddr{Port: int(port)})
	if err != nil {
		return nil, fmt.Errorf("listen: %w", err)
	}
	id := uint16(rand.Intn(65536))

	blm := &BloomFilterManager{}
	blm.start()

	return &Node{
		id:       id,
		port:     port,
		conn:     conn,
		peers:    NewPeerManager(),
		quit:     make(chan struct{}),
		onChange: onChange,
		chatBlm:  blm,
	}, nil
}

func (n *Node) ID() uint16    { return n.id }
func (n *Node) Port() uint16  { return n.port }
func (n *Node) Peers() []Peer { return n.peers.All() }

func (n *Node) Connect(addrStr string) error {
	addr, err := ResolveAddr(addrStr)
	if err != nil {
		return err
	}
	n.sendHello(addr)
	return nil
}

func (n *Node) ChatMessages() []ChatMsg {
	n.chatMu.Lock()
	defer n.chatMu.Unlock()
	msgs := make([]ChatMsg, len(n.chatMsgs))
	copy(msgs, n.chatMsgs)
	return msgs
}

func (n *Node) BroadcastChat(content string) {
	n.chatSeq++
	msgID := n.chatSeq
	msg := EncodeChat(n.id, msgID, 0, content)

	// Add to own chat
	n.chatMu.Lock()
	n.chatMsgs = append(n.chatMsgs, ChatMsg{From: n.id, Time: time.Now(), Content: content})
	if len(n.chatMsgs) > 200 {
		n.chatMsgs = n.chatMsgs[len(n.chatMsgs)-200:]
	}
	n.chatMu.Unlock()

	// Mark as relayed so we don't echo it back
	n.chatBlm.markRelayed(n.id, msgID)

	// Send to all peers
	for _, p := range n.peers.All() {
		n.sendTo(msg, p.Addr)
	}
}

func (n *Node) Start(bootstrap string) {
	if bootstrap != "" {
		addr, err := ResolveAddr(bootstrap)
		if err != nil {
			log.Printf("invalid bootstrap address %q: %v", bootstrap, err)
		} else {
			n.bootstrap = addr
			n.sendHello(addr)
		}
	}

	go n.recvLoop()
	go n.heartbeatLoop()
	go n.cleanLoop()
	if n.bootstrap.Type != 0 {
		go n.bootstrapLoop()
	}
}

func (n *Node) Stop() {
	leave := EncodeLeave(n.id)
	for _, p := range n.peers.All() {
		n.sendTo(leave, p.Addr)
	}
	close(n.quit)
	n.chatBlm.stop()
	n.conn.Close()
}

func (n *Node) recvLoop() {
	buf := make([]byte, 2048)
	for {
		select {
		case <-n.quit:
			return
		default:
		}
		n.conn.SetReadDeadline(time.Now().Add(time.Second))
		nr, remote, err := n.conn.ReadFromUDP(buf)
		if err != nil {
			if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
				continue
			}
			continue
		}
		if n.handlePacket(buf[:nr], remote) {
			if n.onChange != nil {
				n.onChange()
			}
		}
	}
}

func (n *Node) handlePacket(data []byte, remote *net.UDPAddr) (changed bool) {
	if len(data) < 1 {
		return false
	}
	switch data[0] {
	case MsgHello:
		return n.handleHello(data, remote)
	case MsgLeave:
		return n.handleLeave(data)
	case MsgChat:
		return n.handleChat(data, remote)
	}
	return false
}

func (n *Node) handleHello(data []byte, remote *net.UDPAddr) bool {
	if len(data) < 7 {
		return false
	}
	id := binary.BigEndian.Uint16(data[1:3])
	port := binary.BigEndian.Uint16(data[3:5])
	cnt := binary.BigEndian.Uint16(data[5:7])

	senderAddr := addrFromUDP(remote, port)

	if id == n.id {
		return false
	}

	changed := n.peers.AddOrUpdate(id, senderAddr)

	if changed {
		msg := EncodeHello(n.id, n.port, n.peers.Addrs())
		n.sendTo(msg, senderAddr)
	}

	r := bytes.NewReader(data[7:])
	for i := uint16(0); i < cnt; i++ {
		addr, err := DecodeAddr(r)
		if err != nil {
			return changed
		}
		key := addr.String()
		if key != senderAddr.String() && !n.isSelf(addr) {
			if n.peers.AddOrUpdate(0, addr) {
				changed = true
			}
		}
	}
	return changed
}

func (n *Node) handleLeave(data []byte) bool {
	if len(data) < 3 {
		return false
	}
	id := binary.BigEndian.Uint16(data[1:3])
	for _, p := range n.peers.All() {
		if p.ID == id {
			n.peers.RemoveByAddr(p.Addr)
			return true
		}
	}
	return false
}

func (n *Node) handleChat(data []byte, remote *net.UDPAddr) bool {
	if len(data) < 8 {
		return false
	}
	senderID := binary.BigEndian.Uint16(data[1:3])
	msgID := binary.BigEndian.Uint16(data[3:5])
	hops := data[5]
	contentLen := binary.BigEndian.Uint16(data[6:8])

	if len(data) < int(8+contentLen) {
		return false
	}
	content := string(data[8 : 8+contentLen])

	if senderID == n.id {
		return false
	}

	// Dedup
	if n.chatBlm.hasRelayed(senderID, msgID) {
		return false
	}
	n.chatBlm.markRelayed(senderID, msgID)

	// Add to chat history
	n.chatMu.Lock()
	n.chatMsgs = append(n.chatMsgs, ChatMsg{From: senderID, Time: time.Now(), Content: content})
	if len(n.chatMsgs) > 200 {
		n.chatMsgs = n.chatMsgs[len(n.chatMsgs)-200:]
	}
	n.chatMu.Unlock()

	// Relay to all peers; bloom filter dedups the echo
	if hops < MaxRelayHops {
		relay := EncodeChat(senderID, msgID, hops+1, content)
		for _, p := range n.peers.All() {
			n.sendTo(relay, p.Addr)
		}
	}

	return true
}

func (n *Node) heartbeatLoop() {
	ticker := time.NewTicker(heartbeatInterval)
	defer ticker.Stop()
	for {
		select {
		case <-n.quit:
			return
		case <-ticker.C:
			peers := n.peers.Addrs()
			msg := EncodeHello(n.id, n.port, peers)
			for _, p := range n.peers.All() {
				n.sendTo(msg, p.Addr)
			}
		}
	}
}

func (n *Node) cleanLoop() {
	ticker := time.NewTicker(cleanInterval)
	defer ticker.Stop()
	for {
		select {
		case <-n.quit:
			return
		case <-ticker.C:
			if n.peers.Clean() > 0 && n.onChange != nil {
				n.onChange()
			}
		}
	}
}

func (n *Node) bootstrapLoop() {
	ticker := time.NewTicker(10 * time.Second)
	defer ticker.Stop()
	for {
		select {
		case <-n.quit:
			return
		case <-ticker.C:
			if !n.peers.HasAddr(n.bootstrap) {
				n.sendHello(n.bootstrap)
			}
		}
	}
}

func (n *Node) sendHello(addr Addr) {
	msg := EncodeHello(n.id, n.port, n.peers.Addrs())
	n.sendTo(msg, addr)
}

func (n *Node) sendTo(data []byte, addr Addr) {
	var udpAddr *net.UDPAddr
	switch addr.Type {
	case AddrIPv4, AddrIPv6:
		udpAddr = &net.UDPAddr{IP: net.ParseIP(addr.Host), Port: int(addr.Port)}
	case AddrDomain:
		resolved, err := net.ResolveUDPAddr("udp", addr.String())
		if err != nil {
			return
		}
		udpAddr = resolved
	}
	if udpAddr == nil {
		return
	}
	n.conn.WriteToUDP(data, udpAddr)
}

func (n *Node) isSelf(addr Addr) bool {
	if addr.Port != n.port {
		return false
	}
	ip := net.ParseIP(addr.Host)
	if ip == nil {
		return false
	}
	if ip.IsLoopback() || ip.IsUnspecified() {
		return true
	}
	addrs, err := net.InterfaceAddrs()
	if err != nil {
		return false
	}
	for _, a := range addrs {
		if ipNet, ok := a.(*net.IPNet); ok && ipNet.IP.Equal(ip) {
			return true
		}
	}
	return false
}

func addrFromUDP(addr *net.UDPAddr, port uint16) Addr {
	ip := addr.IP
	if ip4 := ip.To4(); ip4 != nil {
		return Addr{Type: AddrIPv4, Host: ip4.String(), Port: port}
	}
	return Addr{Type: AddrIPv6, Host: ip.String(), Port: port}
}