waste-go/internal/mesh/peer.go

// Package mesh/peer handles individual peer TCP connections.
package mesh

import (
	"bufio"
	"encoding/json"
	"fmt"
	"log"
	"net"
	"time"

	"github.com/waste-go/internal/crypto"
	"github.com/waste-go/internal/proto"
)

// HandleConn runs the full lifecycle of one peer connection:
//
//  1. Handshake (Hello / HelloAck)
//  2. ECDH → session key
//  3. Register in mesh
//  4. Concurrent read + write loops
//  5. Unregister on disconnect
//
// Call this in a goroutine for both inbound and outbound connections.
func HandleConn(conn net.Conn, m *Mesh, weInitiated bool) {
	defer conn.Close()
	addr := conn.RemoteAddr().String()
	log.Printf("peer: connected to %s (we initiated: %v)", addr, weInitiated)

	session, peerInfo, err := handshake(conn, m.Identity, weInitiated)
	if err != nil {
		log.Printf("peer: handshake with %s failed: %v", addr, err)
		return
	}
	log.Printf("peer: handshake complete with %s (%s)", peerInfo.Alias, peerInfo.ID.Short())

	// Channel for outbound messages (IPC handler writes here)
	sendCh := make(chan []byte, 64)
	peerConn := &PeerConn{Info: *peerInfo, Send: sendCh}
	m.AddPeer(peerConn)
	defer m.RemovePeer(peerInfo.ID)

	// Outbound writer goroutine
	done := make(chan struct{})
	go func() {
		defer close(done)
		w := bufio.NewWriter(conn)
		for payload := range sendCh {
			nonce, ct, err := session.Encrypt(payload)
			if err != nil {
				log.Printf("peer: encrypt error: %v", err)
				continue
			}
			env := proto.Envelope{Nonce: nonce, Payload: ct}
			line, _ := json.Marshal(env)
			line = append(line, '\n')
			if _, err := w.Write(line); err != nil {
				return
			}
			w.Flush()
		}
	}()

	// Inbound reader loop (this goroutine)
	scanner := bufio.NewScanner(conn)
	for scanner.Scan() {
		var env proto.Envelope
		if err := json.Unmarshal(scanner.Bytes(), &env); err != nil {
			log.Printf("peer: bad envelope from %s: %v", addr, err)
			continue
		}
		plaintext, err := session.Decrypt(env.Nonce, env.Payload)
		if err != nil {
			log.Printf("peer: decrypt error from %s: %v", addr, err)
			continue
		}
		var msg proto.PeerMessage
		if err := json.Unmarshal(plaintext, &msg); err != nil {
			log.Printf("peer: bad peer message from %s: %v", addr, err)
			continue
		}
		handleMessage(msg, peerInfo, m)
	}

	close(sendCh)
	<-done
	log.Printf("peer: disconnected from %s", addr)
}

// handshake performs the Ed25519-authenticated X25519 key exchange.
// Returns the symmetric session and the remote peer's info.
func handshake(conn net.Conn, id *crypto.Identity, weInitiated bool) (*crypto.Session, *proto.PeerInfo, error) {
	conn.SetDeadline(time.Now().Add(10 * time.Second))
	defer conn.SetDeadline(time.Time{})

	ek, err := crypto.GenerateEphemeral()
	if err != nil {
		return nil, nil, err
	}

	ourHello := proto.Hello{
		Version:      1,
		Peer:         id.PeerInfo(),
		EphemeralKey: ek.PublicKeyB64(),
		Signature:    id.Sign([]byte(ek.PublicKeyB64())),
	}

	enc := json.NewEncoder(conn)
	dec := json.NewDecoder(conn)

	if weInitiated {
		// We go first
		if err := enc.Encode(ourHello); err != nil {
			return nil, nil, fmt.Errorf("sending hello: %w", err)
		}
		var ack proto.HelloAck
		if err := dec.Decode(&ack); err != nil {
			return nil, nil, fmt.Errorf("reading hello ack: %w", err)
		}
		if err := crypto.Verify(ack.Peer.PublicKey, []byte(ack.EphemeralKey), ack.Signature); err != nil {
			return nil, nil, fmt.Errorf("bad ack signature: %w", err)
		}
		secret, err := ek.SharedSecret(ack.EphemeralKey)
		if err != nil {
			return nil, nil, err
		}
		return crypto.NewSession(secret), &ack.Peer, nil
	}

	// They go first — read their Hello, then send our Ack
	var theirHello proto.Hello
	if err := dec.Decode(&theirHello); err != nil {
		return nil, nil, fmt.Errorf("reading hello: %w", err)
	}
	if err := crypto.Verify(theirHello.Peer.PublicKey, []byte(theirHello.EphemeralKey), theirHello.Signature); err != nil {
		return nil, nil, fmt.Errorf("bad hello signature: %w", err)
	}

	ack := proto.HelloAck{
		Peer:         id.PeerInfo(),
		EphemeralKey: ek.PublicKeyB64(),
		Signature:    id.Sign([]byte(ek.PublicKeyB64())),
	}
	if err := enc.Encode(ack); err != nil {
		return nil, nil, fmt.Errorf("sending ack: %w", err)
	}

	secret, err := ek.SharedSecret(theirHello.EphemeralKey)
	if err != nil {
		return nil, nil, err
	}
	return crypto.NewSession(secret), &theirHello.Peer, nil
}

// handleMessage dispatches an incoming decrypted peer message.
func handleMessage(msg proto.PeerMessage, from *proto.PeerInfo, m *Mesh) {
	switch msg.Type {
	case proto.MsgChat:
		if msg.Chat == nil {
			return
		}
		m.Emit(proto.IpcMessage{
			Type:    proto.EvtMessageReceived,
			Message: msg.Chat,
		})

	case proto.MsgPeerGossip:
		if msg.Gossip == nil {
			return
		}
		log.Printf("mesh: gossip from %s: %d peer hints", from.Alias, len(msg.Gossip.Peers))
		// TODO: attempt connections to new peers

	case proto.MsgPing:
		log.Printf("mesh: ping from %s", from.Alias)
		// TODO: send pong back through the send channel

	case proto.MsgPong:
		log.Printf("mesh: pong from %s", from.Alias)

	case proto.MsgFileOffer:
		if msg.FileOffer == nil {
			return
		}
		m.Emit(proto.IpcMessage{
			Type:   proto.EvtIncomingFile,
			PeerID: &from.ID,
			Offer:  msg.FileOffer,
		})

	default:
		log.Printf("mesh: unknown message type %q from %s", msg.Type, from.Alias)
	}
}