pp.go

package netpipe

import (
	"fmt"
	"net"
	"sync"

	"github.com/google/uuid"
)

// Peer is a P2P node that connects directly to other peers without a central server.
//
// Every connection performs an automatic Diffie-Hellman key exchange - all traffic
// is AES-256-GCM encrypted by default. The developer never handles keys.
//
// In a mesh, each peer maintains its own registry of connected peers. Any peer
// can send to any other peer it is directly connected to.
//
// Usage:
//
//	peer := netpipe.NewPeer(netpipe.PeerConfig{Port: 6000})
//	peer.OnData(func(peerID string, data []byte) { ... })
//	go peer.Listen()          // accept incoming connections
//	peer.Connect("192.168.1.10", 6000)  // connect to another peer
//	peer.Broadcast([]byte("hello mesh"))
type Peer struct {
	config PeerConfig
	peerID string // our own UUID in the mesh

	listener net.Listener
	peers    map[string]*peerEntry
	mu       sync.RWMutex

	// stream reassembly (shared across all peer connections)
	streams *streamBuffer

	// partial mesh routing table
	routes *routeTable

	// ATC addon - visual dashboard telemetry
	atcBus *ATCBus

	onConnect    func(peerID string)
	onDisconnect func(peerID string)
	onData       func(peerID string, data []byte)
	onStream     func(peerID string, data []byte)
}

// peerEntry tracks one connected peer.
type peerEntry struct {
	id     string
	conn   *safeConn
	crypto *peerCrypto // DH-derived encryption state for this specific peer
}

// PeerConfig controls P2P peer behaviour.
type PeerConfig struct {
	Port      int    // listening port - defaults to 6000
	Protocol  string // defaults to "tcp"
	Name      string // optional human-readable name (advertised during LAN discovery)
	NoEncrypt bool   // set to true to disable DH encryption (testing/trusted LAN only)
	ChunkSize int    // stream chunk size - defaults to 64 KB

	// Addons
	EnableATC bool // enable the ATC visual dashboard addon. defaults to false
	ATCPort   int  // ATC dashboard HTTP port. defaults to 6001
}

func (c *PeerConfig) applyDefaults() {
	if c.Port == 0 {
		c.Port = 6000
	}
	if c.Protocol == "" {
		c.Protocol = "tcp"
	}
}

// NewPeer creates a P2P peer node.
//
//	peer := netpipe.NewPeer(netpipe.PeerConfig{Port: 6000})
func NewPeer(cfg PeerConfig) *Peer {
	cfg.applyDefaults()
	return &Peer{
		config:  cfg,
		peerID:  uuid.New().String(),
		peers:   make(map[string]*peerEntry),
		streams: newStreamBuffer(),
		routes:  newRouteTable(),
	}
}

// ---------------------------------------------------------------------------
// Event registration
// ---------------------------------------------------------------------------

// OnConnect registers a callback fired when a peer joins the mesh.
func (p *Peer) OnConnect(fn func(peerID string)) { p.onConnect = fn }

// OnDisconnect registers a callback fired when a peer leaves the mesh.
func (p *Peer) OnDisconnect(fn func(peerID string)) { p.onDisconnect = fn }

// OnData registers a callback fired when data arrives from any peer.
func (p *Peer) OnData(fn func(peerID string, data []byte)) { p.onData = fn }

// OnStream registers a callback fired when a chunked stream completes from any peer.
func (p *Peer) OnStream(fn func(peerID string, data []byte)) { p.onStream = fn }

// ---------------------------------------------------------------------------
// Lifecycle
// ---------------------------------------------------------------------------

// Listen starts accepting incoming peer connections. Blocks forever.
// Run in a goroutine: go peer.Listen()
func (p *Peer) Listen() error {
	addr := fmt.Sprintf(":%d", p.config.Port)
	ln, err := net.Listen(p.config.Protocol, addr)
	if err != nil {
		return fmt.Errorf("netpipe pp: listen failed: %w", err)
	}
	p.listener = ln
	fmt.Printf("netpipe peer listening on %s %s\n", p.config.Protocol, addr)

	// ATC addon: start dashboard if enabled
	p.startATC()

	for {
		conn, err := ln.Accept()
		if err != nil {
			if p.listener == nil {
				return nil // clean shutdown
			}
			continue
		}
		go p.handleIncoming(conn)
	}
}

// Connect initiates a direct connection to another peer by address.
// Performs the DH handshake and adds the peer to the mesh.
func (p *Peer) Connect(host string, port int) (string, error) {
	addr := fmt.Sprintf("%s:%d", host, port)
	conn, err := net.Dial(p.config.Protocol, addr)
	if err != nil {
		return "", fmt.Errorf("netpipe pp: connect failed: %w", err)
	}

	sc := wrapConn(conn)
	peerID, err := p.handshakeInitiator(sc)
	if err != nil {
		conn.Close()
		return "", fmt.Errorf("netpipe pp: handshake failed: %w", err)
	}

	return peerID, nil
}

// Close shuts down the listener and disconnects all peers.
func (p *Peer) Close() error {
	if p.listener != nil {
		err := p.listener.Close()
		p.listener = nil
		p.DisconnectAll()
		return err
	}
	return nil
}

// ---------------------------------------------------------------------------
// Send methods - all traffic auto-encrypted via DH shared key
// ---------------------------------------------------------------------------

// Send sends data to a single directly-connected peer (for 1:1 connections).
// If multiple peers are connected, use SendTo or Broadcast.
func (p *Peer) Send(data []byte) error {
	p.mu.RLock()
	if len(p.peers) == 0 {
		p.mu.RUnlock()
		return fmt.Errorf("netpipe pp: no peers connected")
	}
	// grab the first (and typically only) peer
	var entry *peerEntry
	for _, e := range p.peers {
		entry = e
		break
	}
	p.mu.RUnlock()

	return p.sendToEntry(entry, data)
}

// SendTo sends data to one specific peer by UUID.
func (p *Peer) SendTo(peerID string, data []byte) error {
	entry := p.getEntry(peerID)
	if entry == nil {
		return fmt.Errorf("netpipe pp: peer %s not found", peerID)
	}
	return p.sendToEntry(entry, data)
}

// Broadcast sends data to all connected peers.
func (p *Peer) Broadcast(data []byte) {
	for _, entry := range p.entrySnapshot() {
		p.sendToEntry(entry, data)
	}
}

// SendStream sends large data to a specific peer as a chunked stream.
func (p *Peer) SendStream(peerID string, data []byte) error {
	entry := p.getEntry(peerID)
	if entry == nil {
		return fmt.Errorf("netpipe pp: peer %s not found", peerID)
	}
	return p.sendStreamToEntry(entry, data)
}

// BroadcastStream sends large data to all connected peers as a chunked stream.
func (p *Peer) BroadcastStream(data []byte) {
	for _, entry := range p.entrySnapshot() {
		p.sendStreamToEntry(entry, data)
	}
}

// ---------------------------------------------------------------------------
// Utility methods
// ---------------------------------------------------------------------------

// GetPeerID returns this peer's UUID in the mesh.
func (p *Peer) GetPeerID() string {
	return p.peerID
}

// GetPeers returns a slice of all connected peer UUIDs.
func (p *Peer) GetPeers() []string {
	p.mu.RLock()
	defer p.mu.RUnlock()
	ids := make([]string, 0, len(p.peers))
	for id := range p.peers {
		ids = append(ids, id)
	}
	return ids
}

// PeerCount returns the number of connected peers.
func (p *Peer) PeerCount() int {
	p.mu.RLock()
	defer p.mu.RUnlock()
	return len(p.peers)
}

// IsConnected returns true if there is at least one connected peer.
func (p *Peer) IsConnected() bool {
	return p.PeerCount() > 0
}

// Disconnect cleanly disconnects one specific peer.
func (p *Peer) Disconnect(peerID string) error {
	p.mu.Lock()
	entry, ok := p.peers[peerID]
	if ok {
		entry.conn.Close()
		delete(p.peers, peerID)
	}
	p.mu.Unlock()

	if !ok {
		return fmt.Errorf("netpipe pp: peer %s not found", peerID)
	}
	if p.onDisconnect != nil {
		p.onDisconnect(peerID)
	}
	return nil
}

// DisconnectAll disconnects from every peer in the mesh.
func (p *Peer) DisconnectAll() {
	p.mu.Lock()
	entries := make([]*peerEntry, 0, len(p.peers))
	for _, e := range p.peers {
		entries = append(entries, e)
	}
	p.peers = make(map[string]*peerEntry)
	p.mu.Unlock()

	for _, e := range entries {
		e.conn.Close()
		if p.onDisconnect != nil {
			p.onDisconnect(e.id)
		}
	}
}

// ---------------------------------------------------------------------------
// Internal - DH handshake (initiator side)
// ---------------------------------------------------------------------------

func (p *Peer) handshakeInitiator(sc *safeConn) (string, error) {
	// 1. generate our ephemeral key pair
	crypto, err := newPeerCrypto()
	if err != nil {
		return "", err
	}

	// 2. generate our peer ID for this mesh
	myID := uuid.New()
	myIDBytes, _ := myID.MarshalBinary()

	// 3. send our public key + ID
	body := make([]byte, handshakeBodySize)
	copy(body[:x25519PubKeySize], crypto.publicKey)
	copy(body[x25519PubKeySize:], myIDBytes)

	if err := writeMessage(sc, flagHandshake, body); err != nil {
		return "", fmt.Errorf("send handshake: %w", err)
	}

	// 4. read the remote peer's handshake
	flags, remoteBody, err := readFrameFull(sc)
	if err != nil {
		return "", fmt.Errorf("read handshake: %w", err)
	}
	if flags&flagHandshake == 0 || len(remoteBody) < handshakeBodySize {
		return "", fmt.Errorf("invalid handshake response")
	}

	remotePubKey := remoteBody[:x25519PubKeySize]
	var remoteUUID uuid.UUID
	if err := remoteUUID.UnmarshalBinary(remoteBody[x25519PubKeySize:handshakeBodySize]); err != nil {
		return "", fmt.Errorf("parse remote UUID: %w", err)
	}
	remoteID := remoteUUID.String()

	// 5. derive the shared AES key
	if !p.shouldEncrypt() {
		crypto.sharedKey = make([]byte, 32) // zero key - no encryption
	} else {
		if err := crypto.deriveSharedKey(remotePubKey); err != nil {
			return "", err
		}
	}

	// 6. register the peer and start the read loop
	entry := &peerEntry{id: remoteID, conn: sc, crypto: crypto}
	p.addEntry(entry)

	p.atcEmitConnect(remoteID)
	if p.onConnect != nil {
		p.onConnect(remoteID)
	}

	go p.readLoop(entry)
	return remoteID, nil
}

// ---------------------------------------------------------------------------
// Internal - DH handshake (listener side - incoming connection)
// ---------------------------------------------------------------------------

func (p *Peer) handleIncoming(conn net.Conn) {
	sc := wrapConn(conn)

	// 1. read the initiator's handshake
	flags, remoteBody, err := readFrameFull(conn)
	if err != nil || flags&flagHandshake == 0 || len(remoteBody) < handshakeBodySize {
		conn.Close()
		return
	}

	remotePubKey := remoteBody[:x25519PubKeySize]
	var remoteUUID uuid.UUID
	if err := remoteUUID.UnmarshalBinary(remoteBody[x25519PubKeySize:handshakeBodySize]); err != nil {
		conn.Close()
		return
	}
	remoteID := remoteUUID.String()

	// 2. generate our key pair
	crypto, err := newPeerCrypto()
	if err != nil {
		conn.Close()
		return
	}

	// 3. send our public key + ID back
	myID := uuid.New()
	myIDBytes, _ := myID.MarshalBinary()

	body := make([]byte, handshakeBodySize)
	copy(body[:x25519PubKeySize], crypto.publicKey)
	copy(body[x25519PubKeySize:], myIDBytes)

	if err := writeMessage(sc, flagHandshake, body); err != nil {
		conn.Close()
		return
	}

	// 4. derive shared key
	if !p.shouldEncrypt() {
		crypto.sharedKey = make([]byte, 32)
	} else {
		if err := crypto.deriveSharedKey(remotePubKey); err != nil {
			conn.Close()
			return
		}
	}

	// 5. register and start read loop
	entry := &peerEntry{id: remoteID, conn: sc, crypto: crypto}
	p.addEntry(entry)

	p.atcEmitConnect(remoteID)
	if p.onConnect != nil {
		p.onConnect(remoteID)
	}

	p.readLoop(entry) // blocking
}

// ---------------------------------------------------------------------------
// Internal - per-peer read loop
// ---------------------------------------------------------------------------

func (p *Peer) readLoop(entry *peerEntry) {
	defer func() {
		p.removeEntry(entry.id)
		p.routes.remove(entry.id)
		p.atcEmitDisconnect(entry.id)
		if p.onDisconnect != nil {
			p.onDisconnect(entry.id)
		}
	}()

	for {
		flags, body, err := readFrameFull(entry.conn)
		if err != nil {
			break
		}

		// ignore control frames
		if flags&flagHandshake != 0 || flags&flagReject != 0 {
			continue
		}
		if flags&flagPing != 0 || flags&flagPong != 0 {
			continue
		}

		// routed message - may need forwarding
		if flags&flagRoute != 0 {
			decrypted := body
			if flags&flagEncrypted != 0 {
				var err error
				decrypted, err = p.decryptBody(entry, body)
				if err != nil {
					continue
				}
			}
			p.handleRoutedMessage(entry.id, decrypted)
			continue
		}

		// stream chunk - reassemble
		if flags&flagStream != 0 {
			decrypted, err := p.decryptBody(entry, body)
			if err != nil {
				continue
			}

			streamID, index, total, chunk, err := parseStreamBody(decrypted)
			if err != nil {
				continue
			}

			assembled, complete := p.streams.addChunk(streamID, index, total, chunk)
			if complete && p.onStream != nil {
				p.onStream(entry.id, assembled)
			}
			continue
		}

		// regular data - decrypt
		data, err := p.decryptBody(entry, body)
		if err != nil {
			continue
		}

		if p.onData != nil {
			p.onData(entry.id, data)
		}
		p.atcEmitDataRecv(entry.id, len(data))
	}
}

// ---------------------------------------------------------------------------
// Internal - encryption helpers
// ---------------------------------------------------------------------------

func (p *Peer) shouldEncrypt() bool {
	return !p.config.NoEncrypt
}

func (p *Peer) sendToEntry(entry *peerEntry, data []byte) error {
	p.atcEmitDataSend(entry.id, len(data))
	body := data
	var flags byte = 0x00

	if p.shouldEncrypt() && entry.crypto.sharedKey != nil {
		encrypted, err := entry.crypto.encryptBytes(data)
		if err != nil {
			return err
		}
		body = encrypted
		flags = flagEncrypted
	}

	return writeMessage(entry.conn, flags, body)
}

func (p *Peer) sendStreamToEntry(entry *peerEntry, data []byte) error {
	cs := p.config.ChunkSize
	if cs <= 0 {
		cs = DefaultChunkSize
	}

	streamID := uuid.New()
	idBytes, _ := streamID.MarshalBinary()

	totalChunks := (len(data) + cs - 1) / cs
	if totalChunks == 0 {
		totalChunks = 1
	}

	for i := 0; i < totalChunks; i++ {
		start := i * cs
		end := start + cs
		if end > len(data) {
			end = len(data)
		}
		chunk := data[start:end]

		// build stream header + chunk
		header := make([]byte, streamHeaderSize)
		copy(header[:16], idBytes)
		putUint32BE(header[16:20], uint32(i))
		putUint32BE(header[20:24], uint32(totalChunks))

		payload := make([]byte, streamHeaderSize+len(chunk))
		copy(payload[:streamHeaderSize], header)
		copy(payload[streamHeaderSize:], chunk)

		// encrypt the whole stream body
		body := payload
		var flags byte = flagStream

		if p.shouldEncrypt() && entry.crypto.sharedKey != nil {
			encrypted, err := entry.crypto.encryptBytes(payload)
			if err != nil {
				return err
			}
			body = encrypted
			flags = flagStream | flagEncrypted
		}

		if err := writeMessage(entry.conn, flags, body); err != nil {
			return err
		}
	}
	return nil
}

func (p *Peer) decryptBody(entry *peerEntry, body []byte) ([]byte, error) {
	if p.shouldEncrypt() && entry.crypto.sharedKey != nil {
		return entry.crypto.decryptBytes(body)
	}
	return body, nil
}

// putUint32BE writes a uint32 in big-endian.
func putUint32BE(b []byte, v uint32) {
	b[0] = byte(v >> 24)
	b[1] = byte(v >> 16)
	b[2] = byte(v >> 8)
	b[3] = byte(v)
}

// ---------------------------------------------------------------------------
// Internal - thread-safe peer registry
// ---------------------------------------------------------------------------

func (p *Peer) addEntry(entry *peerEntry) {
	p.mu.Lock()
	p.peers[entry.id] = entry
	p.mu.Unlock()
	p.routes.addDirect(entry.id) // direct connection = self as next hop
}

func (p *Peer) removeEntry(id string) {
	p.mu.Lock()
	if entry, ok := p.peers[id]; ok {
		entry.conn.Close()
		delete(p.peers, id)
	}
	p.mu.Unlock()
}

func (p *Peer) getEntry(id string) *peerEntry {
	p.mu.RLock()
	defer p.mu.RUnlock()
	return p.peers[id]
}

func (p *Peer) entrySnapshot() []*peerEntry {
	p.mu.RLock()
	entries := make([]*peerEntry, 0, len(p.peers))
	for _, e := range p.peers {
		entries = append(entries, e)
	}
	p.mu.RUnlock()
	return entries
}