Skip to content

Commit 1cda0e5

Browse files
tsosunchiatsosunchia
committed
Fix context propagation and FastIP cache synchronization
1 parent 7c96e4c commit 1cda0e5

11 files changed

Lines changed: 123 additions & 26 deletions

‎cmd/cmd.go‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -786,11 +786,11 @@ func closeLeoWebsocket(leoWs *wshandle.WsConn) {
786786
}
787787
}
788788

789-
func maybeHandleGlobalping(ctx context.Context, from string, opts *trace.GlobalpingOptions, conf *trace.Config) bool {
789+
func maybeHandleGlobalping(from string, opts *trace.GlobalpingOptions, conf *trace.Config) bool {
790790
if from == "" {
791791
return false
792792
}
793-
handleGlobalpingTrace(ctx, opts, conf)
793+
handleGlobalpingTrace(opts, conf)
794794
return true
795795
}
796796

@@ -1391,7 +1391,6 @@ func Execute() {
13911391
}
13921392

13931393
if maybeHandleGlobalping(
1394-
rootCtx,
13951394
*from,
13961395
&trace.GlobalpingOptions{
13971396
Target: *str,

‎cmd/globalping_disabled.go‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,14 +3,13 @@
33
package cmd
44

55
import (
6-
"context"
76
"fmt"
87
"os"
98

109
"github.com/nxtrace/NTrace-core/trace"
1110
)
1211

13-
func handleGlobalpingTrace(_ context.Context, _ *trace.GlobalpingOptions, _ *trace.Config) {
12+
func handleGlobalpingTrace(_ *trace.GlobalpingOptions, _ *trace.Config) {
1413
fmt.Fprintf(os.Stderr, "--from (Globalping) is not available in %s; please use the full nexttrace build\n", appBinName)
1514
os.Exit(1)
1615
}

‎cmd/globalping_full.go‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,10 @@ import (
1616
"github.com/nxtrace/NTrace-core/util"
1717
)
1818

19-
func handleGlobalpingTrace(ctx context.Context, opts *trace.GlobalpingOptions, config *trace.Config) {
20-
if config != nil {
21-
config.Context = ctx
19+
func handleGlobalpingTrace(opts *trace.GlobalpingOptions, config *trace.Config) {
20+
ctx := context.Background()
21+
if config != nil && config.Context != nil {
22+
ctx = config.Context
2223
}
2324
res, measurement, err := trace.GlobalpingTraceroute(opts, config)
2425
if err != nil {

‎cmd/mtr_mode.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -243,7 +243,7 @@ func buildAPIInfo(dataOrigin string) string {
243243
if !strings.EqualFold(dataOrigin, "LeoMoeAPI") {
244244
return ""
245245
}
246-
meta := util.FastIPMetaCache
246+
meta := util.GetFastIPMetaCache()
247247
if meta.IP == "" {
248248
return ""
249249
}
@@ -258,7 +258,7 @@ func buildRawAPIInfoLine(dataOrigin string) string {
258258
if !strings.EqualFold(dataOrigin, "LeoMoeAPI") {
259259
return ""
260260
}
261-
meta := util.FastIPMetaCache
261+
meta := util.GetFastIPMetaCache()
262262
if meta.IP == "" {
263263
return ""
264264
}

‎cmd/mtr_mode_test.go‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -426,16 +426,17 @@ func TestNormalizeMTRReportConfig_WidePreservesGeoSettings(t *testing.T) {
426426
}
427427

428428
func TestBuildRawAPIInfoLine_LeoMoeAPI(t *testing.T) {
429-
old := util.FastIPMetaCache
429+
oldCache := util.GetFastIPCache()
430+
oldMeta := util.GetFastIPMetaCache()
430431
t.Cleanup(func() {
431-
util.FastIPMetaCache = old
432+
util.SetFastIPCacheState(oldCache, oldMeta)
432433
})
433434

434-
util.FastIPMetaCache = util.FastIPMeta{
435+
util.SetFastIPCacheState("", util.FastIPMeta{
435436
IP: "2403:18c0:1001:462:dd:38ff:fe48:e0c5",
436437
Latency: "21.33",
437438
NodeName: "DMIT.NRT",
438-
}
439+
})
439440

440441
got := buildRawAPIInfoLine("LeoMoeAPI")
441442
want := "[NextTrace API] preferred API IP - [2403:18c0:1001:462:dd:38ff:fe48:e0c5] - 21.33ms - DMIT.NRT"

‎pow/pow.go‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,27 @@ const (
1717

1818
var retTokenFn = powclient.RetToken
1919

20+
func resolveTokenRequestTimeout(ctx context.Context, fallback time.Duration) time.Duration {
21+
if fallback <= 0 {
22+
fallback = 5 * time.Second
23+
}
24+
if ctx == nil {
25+
return fallback
26+
}
27+
deadline, ok := ctx.Deadline()
28+
if !ok {
29+
return fallback
30+
}
31+
remaining := time.Until(deadline)
32+
if remaining <= 0 {
33+
return time.Millisecond
34+
}
35+
if remaining < fallback {
36+
return remaining
37+
}
38+
return fallback
39+
}
40+
2041
func GetToken(fastIp string, host string, port string) (string, error) {
2142
return GetTokenWithContext(context.Background(), fastIp, host, port)
2243
}
@@ -34,6 +55,7 @@ func GetTokenWithContext(ctx context.Context, fastIp string, host string, port s
3455
getTokenParams.SNI = host
3556
getTokenParams.Host = host
3657
getTokenParams.UserAgent = util.UserAgent
58+
getTokenParams.TimeoutSec = resolveTokenRequestTimeout(opCtx, getTokenParams.TimeoutSec)
3759
proxyUrl := util.GetProxy()
3860
if proxyUrl != nil {
3961
getTokenParams.Proxy = proxyUrl

‎pow/pow_test.go‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,3 +60,34 @@ func TestGetTokenWithContextReturnsCanceled(t *testing.T) {
6060
t.Fatal("GetTokenWithContext did not return promptly after cancel")
6161
}
6262
}
63+
64+
func TestGetTokenWithContextClampsRequestTimeoutToContextDeadline(t *testing.T) {
65+
oldRetTokenFn := retTokenFn
66+
defer func() { retTokenFn = oldRetTokenFn }()
67+
68+
gotTimeout := make(chan time.Duration, 1)
69+
retTokenFn = func(params *powclient.GetTokenParams) (string, error) {
70+
select {
71+
case gotTimeout <- params.TimeoutSec:
72+
default:
73+
}
74+
return "", errors.New("boom")
75+
}
76+
77+
ctx, cancel := context.WithTimeout(context.Background(), 150*time.Millisecond)
78+
defer cancel()
79+
80+
_, _ = GetTokenWithContext(ctx, "example.com", "example.com", "443")
81+
82+
select {
83+
case timeout := <-gotTimeout:
84+
if timeout <= 0 {
85+
t.Fatalf("retToken timeout = %v, want > 0", timeout)
86+
}
87+
if timeout > 150*time.Millisecond {
88+
t.Fatalf("retToken timeout = %v, want <= 150ms", timeout)
89+
}
90+
case <-time.After(100 * time.Millisecond):
91+
t.Fatal("retTokenFn was not called")
92+
}
93+
}

‎server/ws_handler.go‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,13 @@ func sanitizeLogParam(s string) string {
5555
return b.String()
5656
}
5757

58+
func newWSSessionContext(parent context.Context) (context.Context, context.CancelFunc) {
59+
if parent == nil {
60+
parent = context.Background()
61+
}
62+
return context.WithCancel(parent)
63+
}
64+
5865
type wsEnvelope struct {
5966
Type string `json:"type"`
6067
Data interface{} `json:"data,omitempty"`
@@ -209,7 +216,7 @@ func traceWebsocketHandler(c *gin.Context) {
209216
return
210217
}
211218

212-
sessionCtx, cancel := context.WithCancel(context.Background())
219+
sessionCtx, cancel := newWSSessionContext(c.Request.Context())
213220
defer cancel()
214221
var sessionRef atomic.Pointer[wsTraceSession]
215222
go func() {

‎server/ws_handler_test.go‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package server
22

33
import (
4+
"context"
45
"encoding/json"
56
"errors"
67
"io"
@@ -163,6 +164,23 @@ func TestReadWSInitMessage_ReturnsClearDeadlineError(t *testing.T) {
163164
}
164165
}
165166

167+
func TestNewWSSessionContextInheritsParentCancellation(t *testing.T) {
168+
parent, cancelParent := context.WithCancel(context.Background())
169+
ctx, cancel := newWSSessionContext(parent)
170+
defer cancel()
171+
172+
cancelParent()
173+
174+
select {
175+
case <-ctx.Done():
176+
if !errors.Is(ctx.Err(), context.Canceled) {
177+
t.Fatalf("ctx.Err() = %v, want context.Canceled", ctx.Err())
178+
}
179+
case <-time.After(100 * time.Millisecond):
180+
t.Fatal("session context did not inherit parent cancellation")
181+
}
182+
}
183+
166184
func TestWSTraceSessionSend_QueueOverflowReturnsErrSlowConsumer(t *testing.T) {
167185
conn := newFakeWSConn(true)
168186
session := newWSTraceSession(conn, "cn", 1)

‎util/latency.go‎

Lines changed: 26 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import (
99
"net"
1010
"net/http"
1111
"strings"
12+
"sync"
1213
"time"
1314

1415
"github.com/fatih/color"
@@ -21,7 +22,8 @@ type ResponseInfo struct {
2122
}
2223

2324
var (
24-
timeout = 5 * time.Second
25+
timeout = 5 * time.Second
26+
fastIPCacheMu sync.RWMutex
2527
)
2628
var FastIpCache = ""
2729

@@ -53,6 +55,25 @@ func GetFastIP(domain string, port string, enableOutput bool) string {
5355
return ip
5456
}
5557

58+
func GetFastIPCache() string {
59+
fastIPCacheMu.RLock()
60+
defer fastIPCacheMu.RUnlock()
61+
return FastIpCache
62+
}
63+
64+
func GetFastIPMetaCache() FastIPMeta {
65+
fastIPCacheMu.RLock()
66+
defer fastIPCacheMu.RUnlock()
67+
return FastIPMetaCache
68+
}
69+
70+
func SetFastIPCacheState(ip string, meta FastIPMeta) {
71+
fastIPCacheMu.Lock()
72+
FastIpCache = ip
73+
FastIPMetaCache = meta
74+
fastIPCacheMu.Unlock()
75+
}
76+
5677
func GetFastIPWithContext(ctx context.Context, domain string, port string, enableOutput bool) (string, error) {
5778
if ctx == nil {
5879
ctx = context.Background()
@@ -61,8 +82,8 @@ func GetFastIPWithContext(ctx context.Context, domain string, port string, enabl
6182
if proxyUrl != nil {
6283
return "api.nxtrace.org", nil
6384
}
64-
if FastIpCache != "" {
65-
return FastIpCache, nil
85+
if cachedIP := GetFastIPCache(); cachedIP != "" {
86+
return cachedIP, nil
6687
}
6788

6889
var ips []net.IP
@@ -107,11 +128,12 @@ func GetFastIPWithContext(ctx context.Context, domain string, port string, enabl
107128
result.IP = defaultFastIP()
108129
}
109130

110-
FastIPMetaCache = FastIPMeta{
131+
meta := FastIPMeta{
111132
IP: result.IP,
112133
Latency: result.Latency,
113134
NodeName: strings.TrimSpace(result.Content),
114135
}
136+
SetFastIPCacheState(result.IP, meta)
115137

116138
if enableOutput && !SuppressFastIPOutput {
117139
_, _ = fmt.Fprintf(color.Output, "%s preferred API IP - %s - %s - %s",
@@ -122,7 +144,6 @@ func GetFastIPWithContext(ctx context.Context, domain string, port string, enabl
122144
)
123145
}
124146

125-
FastIpCache = result.IP
126147
return result.IP, nil
127148
}
128149

0 commit comments

Comments
 (0)