From d0248823ce10ca9e464fd56ac590ddd983e3cf07 Mon Sep 17 00:00:00 2001
From: rkonfj <rkonfj@gmail.com>
Date: Sun, 27 Apr 2025 10:49:25 +0800
Subject: [PATCH] peermap: use `X-Secret` header to transmit the secret

---
 disco/ws/ws.go     |  3 ++-
 peermap/peermap.go | 33 +++++++++++++++++++--------------
 2 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/disco/ws/ws.go b/disco/ws/ws.go
index 8ec1ecc..cadec95 100644
--- a/disco/ws/ws.go
+++ b/disco/ws/ws.go
@@ -205,7 +205,8 @@ func (c *WSConn) dial(ctx context.Context, server string) error {
 		return fmt.Errorf("get network secret failed: %w", err)
 	}
 	handshake := http.Header{}
-	handshake.Set("X-Network", networkSecret.Secret)
+	handshake.Set("X-Secret", networkSecret.Secret)
+	handshake.Set("X-Network", networkSecret.Secret) // deprecated, will be removed in v0.13
 	handshake.Set("X-PeerID", c.peerID.String())
 	handshake.Set("X-Nonce", langs.NewNonce())
 	handshake.Set("X-Metadata", c.metadata.Encode())
diff --git a/peermap/peermap.go b/peermap/peermap.go
index e155c65..c68dd63 100644
--- a/peermap/peermap.go
+++ b/peermap/peermap.go
@@ -1,6 +1,7 @@
 package peermap
 
 import (
+	"cmp"
 	"context"
 	"encoding/base64"
 	"encoding/json"
@@ -45,7 +46,7 @@ type peerStat struct {
 }
 type peerConn struct {
 	conn      *websocket.Conn
-	exitSig   chan struct{}
+	closeChan chan struct{}
 	closeOnce sync.Once
 	peerMap   *PeerMap
 
@@ -113,7 +114,7 @@ func (p *peerConn) Close() error {
 		_ = p.conn.WriteControl(websocket.CloseMessage,
 			websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), time.Now().Add(2*time.Second))
 		p.conn.Close()
-		close(p.exitSig)
+		close(p.closeChan)
 		close(p.connData)
 		p.broadcastLeave()
 	})
@@ -188,6 +189,17 @@ func (p *peerConn) leadDisco(target *peerConn) {
 	p.write(b1)
 }
 
+func (p *peerConn) relayTo(target *peerConn, b []byte) {
+	data := b[b[1]+2:]
+	bb := make([]byte, 2+len(p.id)+len(data))
+	bb[0] = b[0]
+	bb[1] = p.id.Len()
+	copy(bb[2:p.id.Len()+2], p.id.Bytes())
+	copy(bb[p.id.Len()+2:], data)
+	_ = target.write(bb)
+	p.stat.RelayRx += uint64(len(b))
+}
+
 func (p *peerConn) broadcast(b []byte) {
 	ctx, _ := p.peerMap.getNetwork(p.networkSecret.Network)
 	var peers []*peerConn
@@ -224,7 +236,7 @@ func (p *peerConn) broadcastLeave() {
 func (p *peerConn) readMessageLoop() {
 	for {
 		select {
-		case <-p.exitSig:
+		case <-p.closeChan:
 			return
 		default:
 		}
@@ -267,14 +279,7 @@ func (p *peerConn) readMessageLoop() {
 			p.leadDisco(tgtPeer)
 			continue
 		}
-		data := b[b[1]+2:]
-		bb := make([]byte, 2+len(p.id)+len(data))
-		bb[0] = b[0]
-		bb[1] = p.id.Len()
-		copy(bb[2:p.id.Len()+2], p.id.Bytes())
-		copy(bb[p.id.Len()+2:], data)
-		_ = tgtPeer.write(bb)
-		p.stat.RelayRx += uint64(len(b))
+		p.relayTo(tgtPeer, b)
 	}
 }
 
@@ -302,7 +307,7 @@ func (p *peerConn) keepalive() {
 	ticker := time.NewTicker(12 * time.Second)
 	for {
 		select {
-		case <-p.exitSig:
+		case <-p.closeChan:
 			ticker.Stop()
 			return
 		case <-ticker.C:
@@ -600,7 +605,7 @@ func (pm *PeerMap) HandlePutNetworkMeta(w http.ResponseWriter, r *http.Request)
 }
 
 func (pm *PeerMap) HandlePeerPacketConnect(w http.ResponseWriter, r *http.Request) {
-	networkSecrest := r.Header.Get("X-Network")
+	networkSecrest := cmp.Or(r.Header.Get("X-Secret"), r.Header.Get("X-Network")) //  "X-Network" is deprecated, will be removed in v0.13
 	jsonSecret := auth.JSONSecret{
 		Network:  networkSecrest,
 		Deadline: math.MaxInt64,
@@ -650,7 +655,7 @@ func (pm *PeerMap) HandlePeerPacketConnect(w http.ResponseWriter, r *http.Reques
 		srLimiter = rate.NewLimiter(rate.Limit(pm.cfg.RateLimiter.StreamW.Limit), pm.cfg.RateLimiter.StreamW.Burst)
 	}
 	peer := peerConn{
-		exitSig:          make(chan struct{}),
+		closeChan:        make(chan struct{}),
 		peerMap:          pm,
 		networkSecret:    jsonSecret,
 		networkContext:   networkCtx,