feat(wireguard): skip stale peers in updateConnectedPeers and add tests · PasarGuard/node@13be1dc · GitHub
Skip to content

Commit 13be1dc

Browse files
ImMohammad20000ImMohammad20000
committed
feat(wireguard): skip stale peers in updateConnectedPeers and add tests
feat(api): implement stat name parsing and response building with tests
1 parent 0fa541d commit 13be1dc

4 files changed

Lines changed: 191 additions & 107 deletions

File tree

‎backend/wireguard/jobs.go‎

Lines changed: 5 additions & 0 deletions

‎backend/wireguard/jobs_test.go‎

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
package wireguard
2+
3+
import (
4+
"context"
5+
"testing"
6+
"time"
7+
8+
"github.com/pasarguard/node/config"
9+
pkgstats "github.com/pasarguard/node/pkg/stats"
10+
"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
11+
)
12+
13+
func TestUpdateConnectedPeersSkipsStaleHandshakePeers(t *testing.T) {
14+
_, recentPub, err := GenerateKeyPair()
15+
if err != nil {
16+
t.Fatalf("failed to generate recent key pair: %v", err)
17+
}
18+
_, stalePub, err := GenerateKeyPair()
19+
if err != nil {
20+
t.Fatalf("failed to generate stale key pair: %v", err)
21+
}
22+
23+
recentKey, err := wgtypes.ParseKey(recentPub)
24+
if err != nil {
25+
t.Fatalf("failed to parse recent public key: %v", err)
26+
}
27+
staleKey, err := wgtypes.ParseKey(stalePub)
28+
if err != nil {
29+
t.Fatalf("failed to parse stale public key: %v", err)
30+
}
31+
32+
now := time.Now()
33+
manager := &Manager{
34+
iFaceName: "wg-test",
35+
client: &fakeWGClient{
36+
deviceFn: func(_ string) (*wgtypes.Device, error) {
37+
return &wgtypes.Device{Peers: []wgtypes.Peer{
38+
{
39+
PublicKey: recentKey,
40+
LastHandshakeTime: now,
41+
ReceiveBytes: 120,
42+
TransmitBytes: 80,
43+
},
44+
{
45+
PublicKey: staleKey,
46+
LastHandshakeTime: now.Add(-onlineActivityThreshold - time.Second),
47+
ReceiveBytes: 300,
48+
TransmitBytes: 200,
49+
},
50+
}}, nil
51+
},
52+
},
53+
}
54+
55+
peerStore := NewPeerStore()
56+
peerStore.Init([]*PeerInfo{
57+
{Email: "recent@example.com", PublicKey: recentKey},
58+
{Email: "stale@example.com", PublicKey: staleKey},
59+
})
60+
61+
wg := &WireGuard{
62+
manager: manager,
63+
cfg: &config.Config{},
64+
config: &Config{},
65+
peerStore: peerStore,
66+
statsTracker: pkgstats.New(),
67+
}
68+
69+
emailByKey := wg.peerStore.GetEmailMap()
70+
if got := emailByKey[recentKey.String()]; got != "recent@example.com" {
71+
t.Fatalf("unexpected recent peer mapping: got %q", got)
72+
}
73+
if got := emailByKey[staleKey.String()]; got != "stale@example.com" {
74+
t.Fatalf("unexpected stale peer mapping: got %q", got)
75+
}
76+
77+
wg.updateConnectedPeers(context.Background())
78+
79+
entries := wg.statsTracker.GetStatsEntries([]string{recentKey.String(), staleKey.String()})
80+
if len(entries) != 1 {
81+
t.Fatalf("expected one tracked peer entry, got %d", len(entries))
82+
}
83+
if _, ok := entries[recentKey.String()]; !ok {
84+
t.Fatalf("expected recent peer to be tracked, but it was not")
85+
}
86+
if _, ok := entries[staleKey.String()]; ok {
87+
t.Fatalf("expected stale peer to be excluded from tracking")
88+
}
89+
90+
resp := wg.statsTracker.GetUsersStats(context.Background(), false)
91+
92+
if len(resp.GetStats()) != 2 {
93+
t.Fatalf("expected stats only for one active peer (2 entries), got %d", len(resp.GetStats()))
94+
}
95+
96+
for _, stat := range resp.GetStats() {
97+
if stat.GetName() != "recent@example.com" {
98+
t.Fatalf("unexpected user in stats: got %s, expected only recent@example.com", stat.GetName())
99+
}
100+
}
101+
}

‎backend/xray/api/stats.go‎

Lines changed: 28 additions & 107 deletions
Original file line numberDiff line numberDiff line change
@@ -76,26 +76,7 @@ func (x *XrayHandler) GetUsersStats(ctx context.Context, reset bool) (*common.St
7676
return nil, err
7777
}
7878

79-
stats := &common.StatResponse{}
80-
for _, stat := range resp.GetStat() {
81-
data := stat.GetName()
82-
value := stat.GetValue()
83-
84-
// Extract the type from the name (e.g., "traffic")
85-
parts := strings.Split(data, ">>>")
86-
name := parts[1]
87-
link := parts[2]
88-
statType := parts[3]
89-
90-
stats.Stats = append(stats.Stats, &common.Stat{
91-
Name: name,
92-
Type: statType,
93-
Link: link,
94-
Value: value,
95-
})
96-
}
97-
98-
return stats, nil
79+
return buildStatResponse(resp.GetStat()), nil
9980
}
10081

10182
func (x *XrayHandler) GetInboundsStats(ctx context.Context, reset bool) (*common.StatResponse, error) {
@@ -104,26 +85,7 @@ func (x *XrayHandler) GetInboundsStats(ctx context.Context, reset bool) (*common
10485
return nil, err
10586
}
10687

107-
stats := &common.StatResponse{}
108-
for _, stat := range resp.GetStat() {
109-
data := stat.GetName()
110-
value := stat.GetValue()
111-
112-
// Extract the type from the name (e.g., "traffic")
113-
parts := strings.Split(data, ">>>")
114-
name := parts[1]
115-
link := parts[2]
116-
statType := parts[3]
117-
118-
stats.Stats = append(stats.Stats, &common.Stat{
119-
Name: name,
120-
Type: statType,
121-
Link: link,
122-
Value: value,
123-
})
124-
}
125-
126-
return stats, nil
88+
return buildStatResponse(resp.GetStat()), nil
12789
}
12890

12991
func (x *XrayHandler) GetOutboundsStats(ctx context.Context, reset bool) (*common.StatResponse, error) {
@@ -132,25 +94,7 @@ func (x *XrayHandler) GetOutboundsStats(ctx context.Context, reset bool) (*commo
13294
return nil, err
13395
}
13496

135-
stats := &common.StatResponse{}
136-
for _, stat := range resp.GetStat() {
137-
data := stat.GetName()
138-
value := stat.GetValue()
139-
140-
parts := strings.Split(data, ">>>")
141-
name := parts[1]
142-
link := parts[2]
143-
statType := parts[3]
144-
145-
stats.Stats = append(stats.Stats, &common.Stat{
146-
Name: name,
147-
Type: statType,
148-
Link: link,
149-
Value: value,
150-
})
151-
}
152-
153-
return stats, nil
97+
return buildStatResponse(resp.GetStat()), nil
15498
}
15599

156100
func (x *XrayHandler) GetUserStats(ctx context.Context, email string, reset bool) (*common.StatResponse, error) {
@@ -162,25 +106,7 @@ func (x *XrayHandler) GetUserStats(ctx context.Context, email string, reset bool
162106
return nil, err
163107
}
164108

165-
stats := &common.StatResponse{}
166-
for _, stat := range resp.GetStat() {
167-
data := stat.GetName()
168-
value := stat.GetValue()
169-
170-
parts := strings.Split(data, ">>>")
171-
name := parts[1]
172-
statType := parts[2]
173-
link := parts[3]
174-
175-
stats.Stats = append(stats.Stats, &common.Stat{
176-
Name: name,
177-
Type: statType,
178-
Link: link,
179-
Value: value,
180-
})
181-
}
182-
183-
return stats, nil
109+
return buildStatResponse(resp.GetStat()), nil
184110
}
185111

186112
func (x *XrayHandler) GetInboundStats(ctx context.Context, tag string, reset bool) (*common.StatResponse, error) {
@@ -192,25 +118,7 @@ func (x *XrayHandler) GetInboundStats(ctx context.Context, tag string, reset boo
192118
return nil, err
193119
}
194120

195-
stats := &common.StatResponse{}
196-
for _, stat := range resp.GetStat() {
197-
data := stat.GetName()
198-
value := stat.GetValue()
199-
200-
parts := strings.Split(data, ">>>")
201-
name := parts[1]
202-
statType := parts[2]
203-
link := parts[3]
204-
205-
stats.Stats = append(stats.Stats, &common.Stat{
206-
Name: name,
207-
Type: statType,
208-
Link: link,
209-
Value: value,
210-
})
211-
}
212-
213-
return stats, nil
121+
return buildStatResponse(resp.GetStat()), nil
214122
}
215123

216124
func (x *XrayHandler) GetOutboundStats(ctx context.Context, tag string, reset bool) (*common.StatResponse, error) {
@@ -222,23 +130,36 @@ func (x *XrayHandler) GetOutboundStats(ctx context.Context, tag string, reset bo
222130
return nil, err
223131
}
224132

225-
stats := &common.StatResponse{}
226-
for _, stat := range resp.GetStat() {
227-
data := stat.GetName()
228-
value := stat.GetValue()
133+
return buildStatResponse(resp.GetStat()), nil
134+
}
229135

230-
parts := strings.Split(data, ">>>")
231-
name := parts[1]
232-
statType := parts[2]
233-
link := parts[3]
136+
func buildStatResponse(queryStats []*command.Stat) *common.StatResponse {
137+
stats := &common.StatResponse{}
138+
for _, stat := range queryStats {
139+
name, link, statType, ok := parseStatName(stat.GetName())
140+
if !ok {
141+
continue
142+
}
234143

235144
stats.Stats = append(stats.Stats, &common.Stat{
236145
Name: name,
237146
Type: statType,
238147
Link: link,
239-
Value: value,
148+
Value: stat.GetValue(),
240149
})
241150
}
242151

243-
return stats, nil
152+
return stats
153+
}
154+
155+
func parseStatName(raw string) (name, link, statType string, ok bool) {
156+
parts := strings.Split(raw, ">>>")
157+
if len(parts) < 4 {
158+
return "", "", "", false
159+
}
160+
if parts[1] == "" || parts[2] == "" || parts[3] == "" {
161+
return "", "", "", false
162+
}
163+
164+
return parts[1], parts[2], parts[3], true
244165
}

‎backend/xray/api/stats_test.go‎

Lines changed: 57 additions & 0 deletions

0 commit comments

Comments
 (0)