sdk: tidy codes

Kim committed Mar 25, 2026 at 16:20 UTC 4654f99fc164add8d8f04f1cde196e4f63546729
4 files changed +45 -85
sdk/expose.go
-11
@@ -42,7 +42,6 @@ type Exposure struct {
42 activeRelayURLs []string
43 bannedRelayURLs []string
44 listeners map[string]*Listener
45 - starting map[string]struct{}
45
46 closeOnce sync.Once
47 connSeq atomic.Uint64
@@ -101,7 +100,6 @@ func Expose(ctx context.Context, cfg ExposeConfig) (*Exposure, error) {
100 accepted: make(chan net.Conn, max(len(relayURLs)*defaultReadyTarget*2, 1)),
101 datagrams: make(chan types.DatagramFrame, max(len(relayURLs)*32, 1)),
102 listeners: make(map[string]*Listener, len(relayURLs)),
104 - starting: make(map[string]struct{}, len(relayURLs)),
103 }
104
105 if len(relayURLs) > 0 {
@@ -401,7 +399,6 @@ func (e *Exposure) banRelayURL(relayURL string) {
399 e.activeRelayURLs = utils.RemoveRelayURL(e.activeRelayURLs, relayURL)
400 e.bannedRelayURLs = utils.AppendUniqueRelayURL(e.bannedRelayURLs, relayURL)
401 delete(e.listeners, relayURL)
404 - delete(e.starting, relayURL)
402 bannedRelayURLs := append([]string(nil), e.bannedRelayURLs...)
403 e.mu.Unlock()
404
@@ -418,10 +415,6 @@ func (e *Exposure) syncListeners(failOnError bool) error {
415 if _, ok := e.listeners[relayURL]; ok {
416 continue
417 }
421 - if _, ok := e.starting[relayURL]; ok {
422 - continue
423 - }
424 - e.starting[relayURL] = struct{}{}
418 missing = append(missing, relayURL)
419 }
420 e.mu.Unlock()
@@ -429,9 +422,6 @@ func (e *Exposure) syncListeners(failOnError bool) error {
422 for _, relayURL := range missing {
423 listener, err := e.newListener(relayURL)
424 if err != nil {
432 - e.mu.Lock()
433 - delete(e.starting, relayURL)
434 - e.mu.Unlock()
425 if failOnError {
426 return fmt.Errorf("listen %q: %w", relayURL, err)
427 }
@@ -468,7 +458,6 @@ func (e *Exposure) installListener(relayURL string, listener *Listener) {
458
459 shouldClose := false
460 e.mu.Lock()
471 - delete(e.starting, relayURL)
461 if e.closed() {
462 shouldClose = true
463 } else if _, exists := e.listeners[relayURL]; exists {
sdk/expose_test.go
-8
@@ -29,9 +29,6 @@ func TestExposureBanRelayURLMovesRelay(t *testing.T) {
29 relayA: listener,
30 relayB: {},
31 },
32 - starting: map[string]struct{}{
33 - relayA: {},
34 - },
32 }
33
34 exposure.banRelayURL(relayA)
@@ -48,14 +45,10 @@ func TestExposureBanRelayURLMovesRelay(t *testing.T) {
45
46 exposure.mu.RLock()
47 _, listenerExists := exposure.listeners[relayA]
51 - _, startingExists := exposure.starting[relayA]
48 exposure.mu.RUnlock()
49 if listenerExists {
50 t.Fatal("banned relay listener still exists in exposure.listeners")
51 }
56 - if startingExists {
57 - t.Fatal("banned relay still exists in exposure.starting")
58 - }
52 }
53
54 func TestExposureApplyRelayURLsSkipsBannedRelay(t *testing.T) {
@@ -69,7 +62,6 @@ func TestExposureApplyRelayURLsSkipsBannedRelay(t *testing.T) {
62 listeners: map[string]*Listener{
63 relayA: {},
64 },
72 - starting: make(map[string]struct{}),
65 }
66
67 added, err := exposure.applyRelayURLs([]string{relayA, relayB}, false)
sdk/listener.go
+6 -14
@@ -59,11 +59,11 @@ type Listener struct {
59 renewBefore time.Duration
60 registerBootstraps []string
61
62 - stream *transport.ClientStream
63 - datagram *transport.ClientDatagram
62 + stream *transport.ClientStream
63 + datagram *transport.ClientDatagram
64 + mitmManager *mitmManager
65
66 registered chan struct{}
66 - banOnce sync.Once
67 closeOnce sync.Once
68 registerOnce sync.Once
69
@@ -75,8 +75,6 @@ type Listener struct {
75 metadata types.LeaseMetadata
76 tlsConfig *tls.Config
77 tlsCloser io.Closer
78 -
79 - mitmManager *mitmManager
78 }
79
80 // NewListener creates one relay listener and its dedicated relay transport for one relay URL.
@@ -111,7 +109,7 @@ func NewListener(ctx context.Context, relayURL string, cfg ListenerConfig) (*Lis
109 retryWait: retryWait,
110 leaseTTL: leaseTTL,
111 renewBefore: renewBefore,
114 - registerBootstraps: append([]string(nil), initialBootstraps...),
112 + registerBootstraps: initialBootstraps,
113 metadata: cfg.Metadata.Copy(),
114 }
115 l.mitmManager = newMITMManager(listenerCtx, l)
@@ -124,9 +122,6 @@ func NewListener(ctx context.Context, relayURL string, cfg ListenerConfig) (*Lis
122 Str("lease_id", l.LeaseID()).
123 Msg("quic datagram plane disconnected; waiting to reconnect")
124 })
127 - }
128 -
129 - if l.datagram != nil {
125 go l.datagram.RunLoop(listenerCtx, l.currentDatagramState, func(ctx context.Context, state transport.ClientDatagramState) (*quic.Conn, error) {
126 return l.api.openQUICSession(ctx, state.LeaseID, state.ReverseToken)
127 })
@@ -548,11 +543,8 @@ func (l *Listener) ban() {
543 if l == nil {
544 return
545 }
551 -
552 - l.banOnce.Do(func() {
553 - l.setStartupStatus(listenerStatusBanned)
554 - _ = l.Close()
555 - })
546 + l.setStartupStatus(listenerStatusBanned)
547 + _ = l.Close()
548 }
549
550 func (l *Listener) StartupStatus() listenerStatus {
sdk/mitm.go
+39 -52
@@ -35,7 +35,6 @@ type MITMProbeReport struct {
35 RelayURL string
36 PublicURL string
37 LeaseID string
38 - CheckedAt time.Time
38 Detected bool
39 Reason string
40 }
@@ -77,22 +76,25 @@ func (m *mitmManager) reset() {
76 }
77
78 func (m *mitmManager) probeTLSPassthrough(ctx context.Context) (MITMProbeReport, error) {
80 - report := MITMProbeReport{
81 - CheckedAt: time.Now(),
82 - }
79 l := m.listener
80 if l == nil || l.api == nil || l.api.baseURL == nil {
85 - return report, errors.New("listener is not ready")
81 + return MITMProbeReport{}, errors.New("listener is not ready")
82 }
83
84 publicURL := l.PublicURL()
85 if publicURL == "" {
90 - return report, errors.New("listener is not registered")
86 + return MITMProbeReport{}, errors.New("listener is not registered")
87 }
88
89 hostname := l.Hostname()
90 if hostname == "" {
95 - return report, errors.New("listener hostname is unavailable")
91 + return MITMProbeReport{}, errors.New("listener hostname is unavailable")
92 + }
93 +
94 + report := MITMProbeReport{
95 + RelayURL: l.api.baseURL.String(),
96 + PublicURL: publicURL,
97 + LeaseID: l.LeaseID(),
98 }
99
100 probeCtx, cancel := context.WithTimeout(ctx, defaultMITMProbeTimeout)
@@ -104,18 +106,28 @@ func (m *mitmManager) probeTLSPassthrough(ctx context.Context) (MITMProbeReport,
106 }
107 nonceHex := hex.EncodeToString(nonceRaw)
108
107 - report.RelayURL = l.api.baseURL.String()
108 - report.PublicURL = publicURL
109 - report.LeaseID = l.LeaseID()
110 -
109 dialAddr, err := m.probeDialAddress(publicURL)
110 if err != nil {
111 return report, err
112 }
113
114 + probeTLSConf := &tls.Config{
115 + ServerName: hostname,
116 + InsecureSkipVerify: true,
117 + }
118 + l.mu.Lock()
119 + if tlsConfig := l.tlsConfig; tlsConfig != nil {
120 + probeTLSConf.MinVersion = tlsConfig.MinVersion
121 + probeTLSConf.MaxVersion = tlsConfig.MaxVersion
122 + if len(tlsConfig.NextProtos) > 0 {
123 + probeTLSConf.NextProtos = append([]string(nil), tlsConfig.NextProtos...)
124 + }
125 + }
126 + l.mu.Unlock()
127 +
128 dialer := &tls.Dialer{
129 NetDialer: &net.Dialer{Timeout: l.api.dialTimeout},
118 - Config: m.clientTLSConfig(hostname),
130 + Config: probeTLSConf,
131 }
132 conn, err := dialer.DialContext(probeCtx, "tcp", dialAddr)
133 if err != nil {
@@ -137,19 +149,19 @@ func (m *mitmManager) probeTLSPassthrough(ctx context.Context) (MITMProbeReport,
149 defer cleanupProbe()
150
151 paddingLen := mitmProbePaddingMin
140 - if mitmProbePaddingMax > mitmProbePaddingMin {
152 + if paddingRange := mitmProbePaddingMax - mitmProbePaddingMin; paddingRange > 0 {
153 var paddingSeed [1]byte
154 if _, err := io.ReadFull(rand.Reader, paddingSeed[:]); err != nil {
155 return report, fmt.Errorf("generate probe padding length: %w", err)
156 }
145 - paddingLen += int(paddingSeed[0]) % (mitmProbePaddingMax - mitmProbePaddingMin + 1)
157 + paddingLen += int(paddingSeed[0]) % (paddingRange + 1)
158 }
159
160 frame := make([]byte, len(nonceRaw)+paddingLen)
149 - copy(frame, nonceRaw)
150 - if _, err := io.ReadFull(rand.Reader, frame[len(nonceRaw):]); err != nil {
151 - return report, fmt.Errorf("generate probe padding: %w", err)
161 + if _, err := io.ReadFull(rand.Reader, frame); err != nil {
162 + return report, fmt.Errorf("generate probe frame: %w", err)
163 }
164 + copy(frame, nonceRaw)
165 if _, err := conn.Write(frame); err != nil {
166 return report, fmt.Errorf("write mitm probe: %w", err)
167 }
@@ -158,15 +170,13 @@ func (m *mitmManager) probeTLSPassthrough(ctx context.Context) (MITMProbeReport,
170 case result := <-resultCh:
171 report.Detected = !result.matched
172 report.Reason = result.reason
161 - return report, nil
173 case <-probeCtx.Done():
163 - report.Detected = false
174 report.Reason = types.MITMProbeReasonProbeTimeout
165 - if errors.Is(probeCtx.Err(), context.DeadlineExceeded) {
166 - return report, nil
175 + if !errors.Is(probeCtx.Err(), context.DeadlineExceeded) {
176 + return report, probeCtx.Err()
177 }
168 - return report, probeCtx.Err()
178 }
179 + return report, nil
180 }
181
182 func (m *mitmManager) probeDialAddress(publicURL string) (string, error) {
@@ -183,26 +193,6 @@ func (m *mitmManager) probeDialAddress(publicURL string) (string, error) {
193 return utils.EnsurePort(dialHost), nil
194 }
195
186 -func (m *mitmManager) clientTLSConfig(hostname string) *tls.Config {
187 - probeTLSConf := &tls.Config{
188 - ServerName: hostname,
189 - InsecureSkipVerify: true,
190 - }
191 -
192 - l := m.listener
193 - l.mu.Lock()
194 - defer l.mu.Unlock()
195 -
196 - if l.tlsConfig != nil {
197 - probeTLSConf.MinVersion = l.tlsConfig.MinVersion
198 - probeTLSConf.MaxVersion = l.tlsConfig.MaxVersion
199 - if len(l.tlsConfig.NextProtos) > 0 {
200 - probeTLSConf.NextProtos = append([]string(nil), l.tlsConfig.NextProtos...)
201 - }
202 - }
203 - return probeTLSConf
204 -}
205 -
196 func (m *mitmManager) maybeStart() {
197 l := m.listener
198 if l.closed() {
@@ -219,20 +209,17 @@ func (m *mitmManager) maybeStart() {
209
210 go func() {
211 report, err := m.probeTLSPassthrough(m.ctx)
222 - m.finish(err == nil && report.Reason != types.MITMProbeReasonProbeTimeout)
212 + success := err == nil && report.Reason != types.MITMProbeReasonProbeTimeout
213 + m.mu.Lock()
214 + m.inFlight = false
215 + if success {
216 + m.lastAt = time.Now()
217 + }
218 + m.mu.Unlock()
219 m.logResult(report, err)
220 }()
221 }
222
227 -func (m *mitmManager) finish(success bool) {
228 - m.mu.Lock()
229 - m.inFlight = false
230 - if success {
231 - m.lastAt = time.Now()
232 - }
233 - m.mu.Unlock()
234 -}
235 -
223 func (m *mitmManager) logResult(report MITMProbeReport, err error) {
224 l := m.listener
225 switch {