inp2p/pkg/punch/punch.go

// Package punch implements UDP and TCP hole-punching.
package punch

import (
	"fmt"
	"log"
	"net"
	"time"

	"github.com/openp2p-cn/inp2p/pkg/protocol"
)

const (
	punchTimeout    = 5 * time.Second
	punchRetries    = 5
	handshakeMagic  = "INP2P-PUNCH"
	handshakeAck    = "INP2P-PUNCH-ACK"
)

// Result holds the outcome of a punch attempt.
type Result struct {
	Conn     net.Conn
	Mode     string // "udp" or "tcp"
	RTT      time.Duration
	PeerAddr string
	Error    error
}

// Config for a punch attempt.
type Config struct {
	PeerIP      string
	PeerPort    int
	PeerNAT     protocol.NATType
	SelfNAT     protocol.NATType
	SelfPort    int // local port to bind (0 = auto)
	IsInitiator bool
}

// AttemptUDP tries to establish a UDP connection via hole-punching.
// Both sides must call this simultaneously (coordinated by server).
func AttemptUDP(cfg Config) Result {
	if !protocol.CanPunch(cfg.SelfNAT, cfg.PeerNAT) {
		return Result{Error: fmt.Errorf("cannot UDP punch: self=%s peer=%s", cfg.SelfNAT, cfg.PeerNAT)}
	}

	localAddr := &net.UDPAddr{Port: cfg.SelfPort}
	conn, err := net.ListenUDP("udp", localAddr)
	if err != nil {
		return Result{Error: fmt.Errorf("listen UDP: %w", err)}
	}

	peerAddr := &net.UDPAddr{
		IP:   net.ParseIP(cfg.PeerIP),
		Port: cfg.PeerPort,
	}

	start := time.Now()

	// Send punch packets
	for i := 0; i < punchRetries; i++ {
		conn.SetWriteDeadline(time.Now().Add(time.Second))
		conn.WriteTo([]byte(handshakeMagic), peerAddr)
		time.Sleep(200 * time.Millisecond)
	}

	// Listen for response
	buf := make([]byte, 256)
	conn.SetReadDeadline(time.Now().Add(punchTimeout))
	n, from, err := conn.ReadFromUDP(buf)
	if err != nil {
		conn.Close()
		return Result{Error: fmt.Errorf("UDP punch timeout: %w", err)}
	}

	// Verify handshake
	msg := string(buf[:n])
	if msg != handshakeMagic && msg != handshakeAck {
		conn.Close()
		return Result{Error: fmt.Errorf("unexpected punch data: %q", msg)}
	}

	// Send ack
	conn.WriteTo([]byte(handshakeAck), from)

	rtt := time.Since(start)
	log.Printf("[punch] UDP punch ok: peer=%s rtt=%s", from, rtt)

	return Result{
		Conn:     conn,
		Mode:     "udp",
		RTT:      rtt,
		PeerAddr: from.String(),
	}
}

// AttemptTCP tries TCP hole-punching using simultaneous SYN.
// This works by having both sides dial each other at the same time.
func AttemptTCP(cfg Config) Result {
	if !protocol.CanPunch(cfg.SelfNAT, cfg.PeerNAT) {
		return Result{Error: fmt.Errorf("cannot TCP punch: self=%s peer=%s", cfg.SelfNAT, cfg.PeerNAT)}
	}

	peerAddr := fmt.Sprintf("%s:%d", cfg.PeerIP, cfg.PeerPort)
	start := time.Now()

	// TCP simultaneous open: keep trying to dial the peer
	var conn net.Conn
	var err error
	for i := 0; i < punchRetries*2; i++ {
		d := net.Dialer{Timeout: time.Second, LocalAddr: &net.TCPAddr{Port: cfg.SelfPort}}
		conn, err = d.Dial("tcp", peerAddr)
		if err == nil {
			break
		}
		time.Sleep(300 * time.Millisecond)
	}

	if err != nil {
		return Result{Error: fmt.Errorf("TCP punch failed: %w", err)}
	}

	// TCP handshake for INP2P
	conn.SetWriteDeadline(time.Now().Add(5 * time.Second))
	conn.Write([]byte(handshakeMagic))

	buf := make([]byte, 256)
	conn.SetReadDeadline(time.Now().Add(5 * time.Second))
	n, err := conn.Read(buf)
	if err != nil {
		conn.Close()
		return Result{Error: fmt.Errorf("TCP handshake read: %w", err)}
	}

	msg := string(buf[:n])
	if msg != handshakeMagic && msg != handshakeAck {
		conn.Close()
		return Result{Error: fmt.Errorf("TCP unexpected handshake: %q", msg)}
	}

	conn.SetWriteDeadline(time.Now().Add(5 * time.Second))
	conn.Write([]byte(handshakeAck))

	rtt := time.Since(start)
	log.Printf("[punch] TCP punch ok: peer=%s rtt=%s", conn.RemoteAddr(), rtt)

	return Result{
		Conn:     conn,
		Mode:     "tcp",
		RTT:      rtt,
		PeerAddr: conn.RemoteAddr().String(),
	}
}

// AttemptDirect tries to directly connect when one side has a public IP.
func AttemptDirect(cfg Config) Result {
	addr := fmt.Sprintf("%s:%d", cfg.PeerIP, cfg.PeerPort)
	start := time.Now()

	conn, err := net.DialTimeout("tcp", addr, punchTimeout)
	if err != nil {
		return Result{Error: fmt.Errorf("direct connect failed: %w", err)}
	}

	rtt := time.Since(start)
	log.Printf("[punch] direct connect ok: peer=%s rtt=%s", addr, rtt)

	return Result{
		Conn:     conn,
		Mode:     "tcp-direct",
		RTT:      rtt,
		PeerAddr: addr,
	}
}

// Connect tries all punch methods in priority order and returns the first success.
func Connect(cfg Config) Result {
	methods := []struct {
		name string
		fn   func(Config) Result
	}{
		{"UDP-punch", AttemptUDP},
		{"TCP-punch", AttemptTCP},
	}

	// If peer has public IP, try direct first
	if cfg.PeerNAT == protocol.NATNone {
		r := AttemptDirect(cfg)
		if r.Error == nil {
			return r
		}
		log.Printf("[punch] direct failed: %v", r.Error)
	}

	for _, m := range methods {
		log.Printf("[punch] trying %s to %s:%d", m.name, cfg.PeerIP, cfg.PeerPort)
		r := m.fn(cfg)
		if r.Error == nil {
			return r
		}
		log.Printf("[punch] %s failed: %v", m.name, r.Error)
	}

	return Result{Error: fmt.Errorf("all punch methods exhausted")}
}