// 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")} }