feat(signal): add P2P-first hole punching with relay fallback (#157)

Previously the signal server answered the initiator immediately with the target's still-unpunched address, so direct P2P rarely had a chance and connections fell back to relay. The server now forwards PunchHole to the target and waits for its PunchHoleSent before delivering the genuine PunchHoleResponse, giving direct P2P a fair chance. A configurable grace period (P2PFallbackMs, default 2000ms) schedules a relay-capable fallback response so TCP-signaling clients never hang (preserves the Phase 7 fix).

Experimental: enabled by default but fully reversible via P2P_FIRST=N / -p2p-first=false. New config: P2PFirst, P2PFallbackMs; env P2P_FIRST, P2P_FALLBACK_MS; flags -p2p-first, -p2p-fallback-ms.

This commit was made possible thanks to Insolve.

betterdesk-server/config/config.go

@@ -81,6 +81,21 @@ type Config struct {
 	// times out (issue #121).  Default: enabled.
 	SameNATRelay bool
 
+	// P2PFirst enables the classic RustDesk hole-punching handshake: instead
+	// of immediately answering the initiator with the target's (still
+	// un-punched) address, the server forwards PunchHole to the target and
+	// waits for its PunchHoleSent before delivering the genuine
+	// PunchHoleResponse. This gives direct P2P a real chance to succeed
+	// (issue #157). If the target does not complete hole punching within
+	// P2PFallbackMs, a best-effort response is delivered so the client can
+	// fall back to relay instead of hanging. Default: enabled.
+	P2PFirst bool
+
+	// P2PFallbackMs is the grace period (milliseconds) the server waits for a
+	// target's PunchHoleSent before sending the relay-capable fallback
+	// response. Only used when P2PFirst is enabled. Default: 2000.
+	P2PFallbackMs int
+
 	// WebSocket security (M3)
 	AllowedWSOrigins    string // Comma-separated allowed WebSocket origins (empty = allow all)
 	APIAllowedWSOrigins string // Comma-separated allowed WebSocket origins for HTTP API events endpoint
@@ -129,6 +144,8 @@ func DefaultConfig() *Config {
 		CDAPRateLimit:        30,
 		SignalRateLimitPerIP: IPRateLimitRegistrations,
 		SameNATRelay:         true, // issue #121: auto-fallback to relay on shared public IP
+		P2PFirst:             true, // issue #157: give direct P2P a real chance before relay
+		P2PFallbackMs:        2000, // grace period for target hole punch before relay fallback
 	}
 }
 
@@ -245,6 +262,22 @@ func (c *Config) LoadEnv() {
 			c.SameNATRelay = false
 		}
 	}
+	// Issue #157: P2P-first hole punching. Enabled by default so direct
+	// connections are attempted before relay. Set P2P_FIRST=N to restore the
+	// legacy behavior of answering the initiator immediately (always relay).
+	if v := os.Getenv("P2P_FIRST"); v != "" {
+		switch strings.ToUpper(v) {
+		case "Y", "YES", "1", "TRUE", "ON":
+			c.P2PFirst = true
+		case "N", "NO", "0", "FALSE", "OFF":
+			c.P2PFirst = false
+		}
+	}
+	if v := os.Getenv("P2P_FALLBACK_MS"); v != "" {
+		if n, err := strconv.Atoi(v); err == nil && n >= 0 {
+			c.P2PFallbackMs = n
+		}
+	}
 	if v := os.Getenv("INIT_ADMIN_USER"); v != "" {
 		c.InitAdminUser = v
 	}

betterdesk-server/main.go

@@ -565,6 +565,8 @@ func parseFlags() *config.Config {
 	flag.IntVar(&cfg.RelayMaxConnsIP, "relay-max-conns-ip", cfg.RelayMaxConnsIP, "Max relay connections per IP (0 = unlimited)")
 	flag.IntVar(&cfg.SignalRateLimitPerIP, "signal-rate-limit-per-ip", cfg.SignalRateLimitPerIP, "Max signal registrations per IP per minute (0 = unlimited; raise for large NAT deployments — issue #122)")
 	flag.BoolVar(&cfg.SameNATRelay, "same-nat-relay", cfg.SameNATRelay, "Auto-fallback to relay when both peers share the same public IP (avoids NAT hairpin failures — issue #121)")
+	flag.BoolVar(&cfg.P2PFirst, "p2p-first", cfg.P2PFirst, "Wait for the target's hole punch before answering the initiator so direct P2P can succeed (issue #157; disable to always answer immediately)")
+	flag.IntVar(&cfg.P2PFallbackMs, "p2p-fallback-ms", cfg.P2PFallbackMs, "Grace period (ms) to wait for the target's PunchHoleSent before sending the relay fallback response (only with --p2p-first)")
 	flag.StringVar(&cfg.InitAdminUser, "init-admin-user", cfg.InitAdminUser, "Initial admin username (default: admin)")
 	flag.StringVar(&cfg.InitAdminPass, "init-admin-pass", cfg.InitAdminPass, "Initial admin password (auto-generated if empty)")
 	flag.BoolVar(&cfg.TLSSignal, "tls-signal", cfg.TLSSignal, "Enable TLS on signal TCP/WS ports (requires --tls-cert and --tls-key)")

betterdesk-server/signal/handler.go

@@ -567,6 +567,10 @@ func (s *Server) handlePunchHoleRequest(msg *pb.PunchHoleRequest, raddr *net.UDP
 
 	if target.UDPAddr != nil {
 		s.sendUDP(punchHole, target.UDPAddr)
+	} else {
+		// Target is connected via TCP/WS (e.g. logged in): forward PunchHole
+		// over its active connection so it can still open its NAT for P2P.
+		s.sendToPeer(targetID, punchHole)
 	}
 
 	// Send PunchHoleResponse to the INITIATOR with signed PK for E2E.
@@ -608,6 +612,24 @@ func (s *Server) handlePunchHoleRequest(msg *pb.PunchHoleRequest, raddr *net.UDP
 			PunchHoleResponse: phr,
 		},
 	}
+
+	// P2P-first (issue #157): when the peers are not on the same LAN, defer
+	// this response and wait for the target's PunchHoleSent, which carries the
+	// target's actual punched address and lets direct P2P succeed. If the
+	// target stays silent past the grace period, the scheduled fallback sends
+	// this relay-capable response so the client can fall back to relay instead
+	// of hanging. handlePunchHoleSent cancels the fallback once the genuine
+	// response is forwarded.
+	if s.cfg.P2PFirst && !sameNetwork {
+		raddrCopy := *raddr
+		s.schedulePunchFallback(normalizeAddrKey(raddr.String()), func() {
+			log.Printf("[signal] P2P-first: target %s did not complete hole punch in time, sending relay fallback to %s",
+				targetID, raddrCopy.String())
+			s.sendUDP(resp, &raddrCopy)
+		})
+		return
+	}
+
 	s.sendUDP(resp, raddr)
 }
 
@@ -770,11 +792,31 @@ func (s *Server) handlePunchHoleRequestTCP(msg *pb.PunchHoleRequest, raddr *net.
 		phr.Union = &pb.PunchHoleResponse_NatType{NatType: pb.NatType(target.NATType)}
 	}
 
-	return &pb.RendezvousMessage{
+	resp := &pb.RendezvousMessage{
 		Union: &pb.RendezvousMessage_PunchHoleResponse{
 			PunchHoleResponse: phr,
 		},
 	}
+
+	// P2P-first (issue #157): for non-LAN peers, defer this response and wait
+	// for the target's PunchHoleSent, which carries the target's actual punched
+	// address and lets direct P2P succeed. The TCP connection is already kept
+	// alive (keepAlive via logAndCheckKeepAlive) and registered in
+	// tcpPunchConns, so handlePunchHoleSent can forward the genuine response
+	// over it. If the target stays silent past the grace period, the scheduled
+	// fallback forwards this relay-capable response so the client can fall back
+	// to relay instead of hanging (preserving the Phase 7 timeout fix).
+	if s.cfg.P2PFirst && !sameNetwork {
+		initiatorKey := normalizeAddrKey(raddr.String())
+		s.schedulePunchFallback(initiatorKey, func() {
+			log.Printf("[signal] P2P-first (TCP): target %s did not complete hole punch in time, forwarding relay fallback to %s",
+				targetID, initiatorKey)
+			s.forwardToTCPInitiator(initiatorKey, resp)
+		})
+		return nil
+	}
+
+	return resp
 }
 
 // handlePunchHoleSent processes a PunchHoleSent message from the target peer.
@@ -875,6 +917,13 @@ func (s *Server) handlePunchHoleSent(phs *pb.PunchHoleSent, senderAddr *net.UDPA
 
 	addrStr := normalizeAddrKey(initiatorAddr.String())
 
+	// P2P-first (issue #157): the target completed hole punching, so cancel any
+	// scheduled relay fallback for this initiator before delivering the genuine
+	// PunchHoleResponse (which carries the target's real punched address).
+	if s.cancelPunchFallback(addrStr) {
+		log.Printf("[signal] P2P-first: cancelled relay fallback for %s — direct P2P response incoming", addrStr)
+	}
+
 	// Try TCP delivery first (initiator may have an open TCP connection).
 	if s.forwardToTCPInitiator(addrStr, resp) {
 		log.Printf("[signal] PunchHoleResponse forwarded via TCP to %s (target=%s)", addrStr, phs.Id)

betterdesk-server/signal/server.go

@@ -49,6 +49,19 @@ type pendingUUID struct {
 	createdAt time.Time
 }
 
+// pendingPunch tracks a scheduled P2P-first fallback for an initiator that is
+// waiting for its target to complete hole punching (issue #157). When the
+// target responds with PunchHoleSent in time, handlePunchHoleSent cancels the
+// fallback and the genuine PunchHoleResponse (carrying the target's punched
+// address) is delivered. Otherwise the fallback fires, delivering a
+// relay-capable response so the client can fall back to relay instead of
+// hanging.
+type pendingPunch struct {
+	timer     *time.Timer
+	fired     atomic.Bool
+	createdAt time.Time
+}
+
 // writeProto sends a protobuf message, using encryption if the connection is secure.
 func (pc *tcpPunchConn) writeProto(msg *pb.RendezvousMessage) error {
 	pc.writeMu.Lock()
@@ -89,6 +102,10 @@ type Server struct {
 	// so relay pairing succeeds. Key=targetID, Value=*pendingUUID.
 	pendingRelayUUIDs sync.Map // map[string]*pendingUUID
 
+	// pendingPunches tracks P2P-first fallback timers per initiator address
+	// (issue #157). Key=normalizeAddrKey(initiatorAddr), Value=*pendingPunch.
+	pendingPunches sync.Map // map[string]*pendingPunch
+
 	// localIP is the server's detected public IP address (via external service).
 	// Used to build the relay server address when -relay-servers is not set.
 	localIP atomic.Value // stores string
@@ -799,6 +816,27 @@ func (s *Server) cleanupTCPPunchConns() {
 			if uuidEvicted > 0 {
 				log.Printf("[signal] Pending relay UUIDs cleanup: evicted %d stale entries", uuidEvicted)
 			}
+
+			// Safety sweep for pendingPunches (issue #157). Entries normally
+			// self-remove when the fallback fires or is cancelled; this only
+			// reclaims orphans left by abnormal shutdown paths.
+			punchEvicted := 0
+			s.pendingPunches.Range(func(key, value any) bool {
+				pp := value.(*pendingPunch)
+				if now.Sub(pp.createdAt) > maxTTL {
+					if val, ok := s.pendingPunches.LoadAndDelete(key); ok {
+						if stale, ok := val.(*pendingPunch); ok {
+							stale.fired.Store(true)
+							stale.timer.Stop()
+						}
+					}
+					punchEvicted++
+				}
+				return true
+			})
+			if punchEvicted > 0 {
+				log.Printf("[signal] Pending punches cleanup: evicted %d stale entries", punchEvicted)
+			}
 		}
 	}
 }
@@ -835,6 +873,50 @@ func (s *Server) getPendingUUID(targetID string) string {
 	return ""
 }
 
+// schedulePunchFallback registers a delayed P2P-first fallback for an
+// initiator (issue #157). After the configured grace period, fallback() runs
+// unless cancelPunchFallback is called first (i.e. the target completed hole
+// punching and the genuine PunchHoleResponse was delivered). Any pre-existing
+// pending punch for the same initiator is replaced.
+func (s *Server) schedulePunchFallback(initiatorKey string, fallback func()) {
+	delay := time.Duration(s.cfg.P2PFallbackMs) * time.Millisecond
+	if delay <= 0 {
+		delay = 2 * time.Second
+	}
+	pp := &pendingPunch{createdAt: time.Now()}
+	pp.timer = time.AfterFunc(delay, func() {
+		if pp.fired.Swap(true) {
+			return
+		}
+		s.pendingPunches.Delete(initiatorKey)
+		fallback()
+	})
+	if old, loaded := s.pendingPunches.Swap(initiatorKey, pp); loaded {
+		if op, ok := old.(*pendingPunch); ok {
+			op.fired.Store(true)
+			op.timer.Stop()
+		}
+	}
+}
+
+// cancelPunchFallback cancels a scheduled P2P-first fallback for an initiator.
+// Returns true if the fallback was cancelled before it could fire.
+func (s *Server) cancelPunchFallback(initiatorKey string) bool {
+	val, ok := s.pendingPunches.LoadAndDelete(initiatorKey)
+	if !ok {
+		return false
+	}
+	pp, ok := val.(*pendingPunch)
+	if !ok {
+		return false
+	}
+	if pp.fired.Swap(true) {
+		return false // already fired
+	}
+	pp.timer.Stop()
+	return true
+}
+
 // isNormalClose returns true if the error represents a normal connection close
 // (EOF, timeout, or connection reset by peer). These are expected during
 // TCP connection lifecycle and should not be logged as errors.