refact: separate proxy and add bps
rabbitprincess committed
Apr 12, 2026 at 21:16 UTC
4a5d947ce0a93ee0156872f5ef5ba66261e82ce2
6 files changed
+109
-45
portal/api_server.go
+5
-3
@@ -157,6 +157,8 @@ func (s *Server) handleRelayDiscovery(w http.ResponseWriter, r *http.Request) {
157
}
158
159
now := time.Now().UTC()
160
+ activeConns := float64(s.proxy.ActiveConns())
161
+ tcpTrafficBPS := s.proxy.CurrentTCPBPS(now)
162
ingressAddr := s.identity.Name
163
if s.cfg.SNIPort != 0 && s.cfg.SNIPort != 443 {
164
ingressAddr = fmt.Sprintf("%s:%d", ingressAddr, s.cfg.SNIPort)
@@ -176,8 +178,6 @@ func (s *Server) handleRelayDiscovery(w http.ResponseWriter, r *http.Request) {
178
Identity: s.identity.Base(),
179
RelayID: s.cfg.PortalURL,
180
OwnerAddress: s.identity.Address,
179
- SignerPublicKey: s.identity.PublicKey,
180
- Sequence: uint64(now.UnixMilli()),
181
Version: 1,
182
IssuedAt: now,
183
ExpiresAt: now.Add(2 * discovery.DiscoveryPollInterval),
@@ -190,7 +190,9 @@ func (s *Server) handleRelayDiscovery(w http.ResponseWriter, r *http.Request) {
190
SupportsUDP: s.cfg.UDPEnabled && s.quicTunnel != nil,
191
SupportsTCP: s.cfg.TCPEnabled,
192
SupportsOverlayPeer: s.overlay != nil,
193
- Load: float64(s.activeConns.Load()),
193
+ Load: activeConns,
194
+ LoadScore: tcpTrafficBPS,
195
+ LastUpdated: now.UnixMilli(),
196
})
197
if err != nil {
198
utils.WriteAPIError(w, http.StatusInternalServerError, types.APIErrorCodeInternal, err.Error())
portal/proxy.go
new
+101
@@ -0,0 +1,101 @@
1
+package portal
2
+
3
+import (
4
+ "io"
5
+ "net"
6
+ "sync"
7
+ "sync/atomic"
8
+ "time"
9
+
10
+ "golang.org/x/sync/errgroup"
11
+)
12
+
13
+type proxy struct {
14
+ activeConns atomic.Int64
15
+ tcpBytes atomic.Int64
16
+ tcpLoadMu sync.Mutex
17
+ tcpLoadAt time.Time
18
+ tcpLoadBytes int64
19
+}
20
+
21
+func (p *proxy) Bridge(left, right net.Conn) {
22
+ p.activeConns.Add(1)
23
+ defer p.activeConns.Add(-1)
24
+
25
+ defer left.Close()
26
+ defer right.Close()
27
+
28
+ var group errgroup.Group
29
+ group.Go(func() error {
30
+ _, err := io.Copy(&countingConn{Conn: right, bytes: &p.tcpBytes}, left)
31
+ closeWrite(right)
32
+ return err
33
+ })
34
+ group.Go(func() error {
35
+ _, err := io.Copy(&countingConn{Conn: left, bytes: &p.tcpBytes}, right)
36
+ closeWrite(left)
37
+ return err
38
+ })
39
+ _ = group.Wait()
40
+}
41
+
42
+func (p *proxy) ActiveConns() int64 {
43
+ return p.activeConns.Load()
44
+}
45
+
46
+func (p *proxy) CurrentTCPBPS(now time.Time) float64 {
47
+ totalTCPBytes := p.tcpBytes.Load()
48
+
49
+ p.tcpLoadMu.Lock()
50
+ defer p.tcpLoadMu.Unlock()
51
+
52
+ if p.tcpLoadAt.IsZero() {
53
+ p.tcpLoadAt = now
54
+ p.tcpLoadBytes = totalTCPBytes
55
+ return 0
56
+ }
57
+
58
+ if elapsed := now.Sub(p.tcpLoadAt); elapsed > 0 {
59
+ tcpTrafficBPS := float64(totalTCPBytes-p.tcpLoadBytes) / elapsed.Seconds()
60
+ p.tcpLoadAt = now
61
+ p.tcpLoadBytes = totalTCPBytes
62
+ return tcpTrafficBPS
63
+ }
64
+
65
+ return 0
66
+}
67
+
68
+type countingConn struct {
69
+ net.Conn
70
+ bytes *atomic.Int64
71
+}
72
+
73
+func (c *countingConn) Write(p []byte) (int, error) {
74
+ n, err := c.Conn.Write(p)
75
+ if n > 0 {
76
+ c.bytes.Add(int64(n))
77
+ }
78
+ return n, err
79
+}
80
+
81
+func (c *countingConn) ReadFrom(r io.Reader) (int64, error) {
82
+ readerFrom, ok := c.Conn.(io.ReaderFrom)
83
+ if !ok {
84
+ return io.Copy(struct{ io.Writer }{Writer: c}, r)
85
+ }
86
+
87
+ n, err := readerFrom.ReadFrom(r)
88
+ if n > 0 {
89
+ c.bytes.Add(n)
90
+ }
91
+ return n, err
92
+}
93
+
94
+func closeWrite(conn net.Conn) {
95
+ type closeWriter interface {
96
+ CloseWrite() error
97
+ }
98
+ if cw, ok := conn.(closeWriter); ok {
99
+ _ = cw.CloseWrite()
100
+ }
101
+}
portal/server.go
+3
-34
@@ -10,7 +10,6 @@ import (
10
"net/http"
11
"strings"
12
"sync"
13
- "sync/atomic"
13
"time"
14
15
"github.com/gosuda/keyless_tls/relay/l4"
@@ -109,7 +108,7 @@ type Server struct {
108
cfg ServerConfig
109
identity types.RelayIdentity
110
acmeManager *acme.Manager
112
- activeConns atomic.Int64
111
+ proxy proxy
112
113
apiListener net.Listener
114
sniListener net.Listener
@@ -491,7 +490,7 @@ func (s *Server) runSNIListener(ctx context.Context) error {
490
_ = wrappedConn.Close()
491
return
492
}
494
- s.BridgeConns(wrappedConn, upstream)
493
+ s.proxy.Bridge(wrappedConn, upstream)
494
return
495
}
496
@@ -510,7 +509,7 @@ func (s *Server) runSNIListener(ctx context.Context) error {
509
return
510
}
511
513
- s.BridgeConns(wrappedConn, session)
512
+ s.proxy.Bridge(wrappedConn, session)
513
}(conn)
514
case errors.Is(err, net.ErrClosed):
515
return nil
@@ -652,33 +651,3 @@ func (s *Server) runRelayDiscoveryLoop(ctx context.Context) error {
651
}
652
}
653
}
655
-
656
-func (s *Server) BridgeConns(left, right net.Conn) {
657
- s.activeConns.Add(1)
658
- defer s.activeConns.Add(-1)
659
-
660
- defer left.Close()
661
- defer right.Close()
662
-
663
- var group errgroup.Group
664
- group.Go(func() error {
665
- _, err := io.Copy(right, left)
666
- closeWrite(right)
667
- return err
668
- })
669
- group.Go(func() error {
670
- _, err := io.Copy(left, right)
671
- closeWrite(left)
672
- return err
673
- })
674
- _ = group.Wait()
675
-}
676
-
677
-func closeWrite(conn net.Conn) {
678
- type closeWriter interface {
679
- CloseWrite() error
680
- }
681
- if cw, ok := conn.(closeWriter); ok {
682
- _ = cw.CloseWrite()
683
- }
684
-}
portal/server_test.go
-2
@@ -37,7 +37,6 @@ func mustRelayDescriptor(t *testing.T, relayURL string) types.RelayDescriptor {
37
Name: utils.PortalRootHost(relayURL),
38
},
39
RelayID: relayURL,
40
- Sequence: uint64(now.UnixMilli()),
40
Version: 1,
41
IssuedAt: now,
42
ExpiresAt: now.Add(time.Hour),
@@ -543,7 +542,6 @@ func TestServerDiscoverySkipsSelfRelayHint(t *testing.T) {
542
selfHint, err := utils.NormalizeDescriptor(types.RelayDescriptor{
543
Identity: server.identity.Base(),
544
RelayID: "https://self-mirror.example.com",
546
- Sequence: uint64(now.UnixMilli()),
545
Version: 1,
546
IssuedAt: now,
547
ExpiresAt: now.Add(time.Hour),
types/identity.go
-2
@@ -105,8 +105,6 @@ type RelayDescriptor struct {
105
106
RelayID string `json:"relay_id,omitempty"`
107
OwnerAddress string `json:"owner_address,omitempty"`
108
- SignerPublicKey string `json:"signer_public_key,omitempty"`
109
- Sequence uint64 `json:"sequence"`
108
Version uint32 `json:"version"`
109
IssuedAt time.Time `json:"issued_at"`
110
ExpiresAt time.Time `json:"expires_at"`
utils/identity.go
-4
@@ -41,7 +41,6 @@ func NormalizeDescriptor(desc types.RelayDescriptor) (types.RelayDescriptor, err
41
desc.OverlayIPv4 = strings.TrimSpace(desc.OverlayIPv4)
42
desc.OverlayCIDRs = NormalizeIPPrefixes(desc.OverlayCIDRs)
43
desc.OwnerAddress = strings.TrimSpace(desc.OwnerAddress)
44
- desc.SignerPublicKey = strings.TrimSpace(desc.SignerPublicKey)
44
if !desc.IssuedAt.IsZero() {
45
desc.IssuedAt = desc.IssuedAt.UTC()
46
}
@@ -83,9 +82,6 @@ func NormalizeDescriptor(desc types.RelayDescriptor) (types.RelayDescriptor, err
82
}
83
desc.OwnerAddress = normalized
84
}
86
- if desc.SignerPublicKey == "" {
87
- desc.SignerPublicKey = desc.PublicKey
88
- }
85
86
switch {
87
case desc.Name == "":