adjust and refactoring quic/udp
Hee Sung Son committed
Mar 17, 2026 at 22:09 UTC
811a4c123e5ddcb89042674dcea08021568e07a0
16 files changed
+284
-148
cmd/demo-app/main.go
+1
-1
@@ -66,7 +66,7 @@ func runDemo() error {
66
return fmt.Errorf("resolve relay urls: %w", err)
67
}
68
69
- exposure, err := sdk.Expose(ctx, relayURLs, flagName, types.LeaseMetadata{
69
+ exposure, err := sdk.Expose(ctx, relayURLs, flagName, types.TransportTCP, types.LeaseMetadata{
70
Description: flagDesc,
71
Tags: utils.SplitCSV(flagTags),
72
Owner: flagOwner,
cmd/portal-tunnel/main.go
+9
-2
@@ -87,6 +87,9 @@ func runExposeCommand(args []string) error {
87
fs.StringVar(&thumbnail, "thumbnail", "", "Service thumbnail URL metadata")
88
fs.StringVar(&owner, "owner", "", "Service owner metadata")
89
fs.BoolVar(&hide, "hide", false, "Hide service from discovery")
90
+
91
+ var transport string
92
+ fs.StringVar(&transport, "transport", types.TransportBoth, "Transport mode: tcp, udp, or both (default: both)")
93
fs.Usage = func() {
94
printExposeUsage(fs.Output())
95
}
@@ -154,6 +157,7 @@ func runExposeCommand(args []string) error {
157
relayURLs,
158
target,
159
name,
160
+ transport,
161
types.LeaseMetadata{
162
Description: desc,
163
Tags: utils.SplitCSV(tags),
@@ -219,11 +223,12 @@ func runTunnel(
223
relayURLs []string,
224
target string,
225
name string,
226
+ transport string,
227
metadata types.LeaseMetadata,
228
) error {
229
logger := log.With().Str("component", "portal").Logger()
230
226
- exposure, err := sdk.Expose(ctx, relayURLs, name, metadata)
231
+ exposure, err := sdk.Expose(ctx, relayURLs, name, transport, metadata)
232
if err != nil {
233
return fmt.Errorf("service %s: failed to start relays: %w", name, err)
234
}
@@ -233,7 +238,9 @@ func runTunnel(
238
defer exposure.Close()
239
240
// UDP is best-effort — attach to existing lease, log and continue if it fails.
236
- go runUDPBestEffort(ctx, exposure, target)
241
+ if transport == types.TransportUDP || transport == types.TransportBoth {
242
+ go runUDPBestEffort(ctx, exposure, target)
243
+ }
244
245
logger.Info().
246
Str("release_version", types.ReleaseVersion).
cmd/portal-tunnel/relays.go
+8
-1
@@ -228,6 +228,7 @@ func proxyUDPRelayConnections(ctx context.Context, udpListener *sdk.UDPListener,
228
}
229
230
// Main loop: relay → local service.
231
+ logger.Info().Str("target", targetAddr).Msg("udp proxy loop started, waiting for datagrams")
232
for {
233
dg, err := udpListener.AcceptDatagram()
234
if err != nil {
@@ -237,6 +238,12 @@ func proxyUDPRelayConnections(ctx context.Context, udpListener *sdk.UDPListener,
238
return fmt.Errorf("accept datagram: %w", err)
239
}
240
241
+ logger.Debug().
242
+ Uint32("flow_id", dg.FlowID).
243
+ Int("bytes", len(dg.Payload)).
244
+ Str("target", targetAddr).
245
+ Msg("datagram received from relay, forwarding to local")
246
+
247
localConn, err := getOrCreateFlow(dg.FlowID)
248
if err != nil {
249
logger.Warn().Err(err).Uint32("flow_id", dg.FlowID).Msg("dial local udp failed")
@@ -244,7 +251,7 @@ func proxyUDPRelayConnections(ctx context.Context, udpListener *sdk.UDPListener,
251
}
252
253
if _, err := localConn.Write(dg.Payload); err != nil {
247
- logger.Debug().Err(err).Uint32("flow_id", dg.FlowID).Msg("write to local udp failed")
254
+ logger.Warn().Err(err).Uint32("flow_id", dg.FlowID).Msg("write to local udp failed")
255
}
256
}
257
}
portal/api_server.go
+2
-11
@@ -321,7 +321,7 @@ func (s *Server) registerLease(req types.RegisterRequest, clientIP string) (type
321
if s.ports == nil {
322
return types.RegisterResponse{}, errors.New("udp port allocation not available")
323
}
324
- udpPort, portErr := s.ports.Allocate()
324
+ udpPort, portErr := s.ports.Allocate(name)
325
if portErr != nil {
326
return types.RegisterResponse{}, fmt.Errorf("allocate udp port: %w", portErr)
327
}
@@ -382,16 +382,7 @@ func (s *Server) unregisterLease(req types.UnregisterRequest) error {
382
if err != nil {
383
return err
384
}
385
- record.Broker.Close()
386
- if record.QUICBroker != nil {
387
- record.QUICBroker.Stop()
388
- }
389
- if record.UDPRelay != nil {
390
- record.UDPRelay.Stop()
391
- }
392
- if record.UDPPort > 0 && s.ports != nil {
393
- s.ports.Release(record.UDPPort)
394
- }
385
+ s.closeLease(record)
386
return nil
387
}
388
portal/lease.go
+3
-6
@@ -17,6 +17,7 @@ type leaseRegistry struct {
17
routes *routeTable
18
leaseByID map[string]*leaseRecord
19
policy *policy.Runtime
20
+ onExpired func(*leaseRecord) // called for each expired lease during cleanup
21
mu sync.RWMutex
22
}
23
@@ -206,12 +207,8 @@ func (r *leaseRegistry) Touch(leaseID, clientIP string, now time.Time) *leaseRec
207
208
func (r *leaseRegistry) cleanupExpired(now time.Time) {
209
for _, lease := range r.removeExpired(now) {
209
- lease.Broker.Close()
210
- if lease.QUICBroker != nil {
211
- lease.QUICBroker.Stop()
212
- }
213
- if lease.UDPRelay != nil {
214
- lease.UDPRelay.Stop()
210
+ if r.onExpired != nil {
211
+ r.onExpired(lease)
212
}
213
}
214
}
portal/lease_test.go
+3
@@ -139,6 +139,9 @@ func TestLeaseRegistryCleanupExpiredClosesBroker(t *testing.T) {
139
t.Parallel()
140
141
registry := newLeaseRegistry(policy.NewRuntime())
142
+ registry.onExpired = func(r *leaseRecord) {
143
+ r.Broker.Close()
144
+ }
145
record := &leaseRecord{
146
ID: "lease_expired",
147
Hostname: "expired.example.com",
portal/quic.go
+7
-34
@@ -2,7 +2,6 @@ package portal
2
3
import (
4
"context"
5
- "encoding/binary"
5
"encoding/json"
6
"errors"
7
"net"
@@ -11,41 +10,15 @@ import (
10
11
"github.com/quic-go/quic-go"
12
"github.com/rs/zerolog/log"
13
+
14
+ "github.com/gosuda/portal/v2/types"
15
)
16
17
var (
18
errQUICNoConnection = errors.New("no quic connection registered")
19
errQUICAlreadyClosed = errors.New("quic broker closed")
19
- errDatagramTooSmall = errors.New("datagram too small to decode")
20
)
21
22
-// datagramFrame is the wire format for QUIC DATAGRAM payloads.
23
-// Layout: [flowID varint][payload bytes]
24
-type datagramFrame struct {
25
- FlowID uint32
26
- Payload []byte
27
-}
28
-
29
-func encodeDatagram(flowID uint32, payload []byte) []byte {
30
- var buf [binary.MaxVarintLen32]byte
31
- n := binary.PutUvarint(buf[:], uint64(flowID))
32
- out := make([]byte, n+len(payload))
33
- copy(out, buf[:n])
34
- copy(out[n:], payload)
35
- return out
36
-}
37
-
38
-func decodeDatagram(data []byte) (datagramFrame, error) {
39
- flowID, n := binary.Uvarint(data)
40
- if n <= 0 {
41
- return datagramFrame{}, errDatagramTooSmall
42
- }
43
- return datagramFrame{
44
- FlowID: uint32(flowID),
45
- Payload: data[n:],
46
- }, nil
47
-}
48
-
22
// quicBroker manages a single QUIC connection from a tunnel for one lease.
23
// All UDP traffic for the lease is multiplexed over DATAGRAM frames on this
24
// connection, identified by flow IDs.
@@ -57,7 +30,7 @@ type quicBroker struct {
30
addrIndex map[string]uint32 // "ip:port" → flowID
31
nextFlow uint32
32
60
- incoming chan datagramFrame // frames received from tunnel
33
+ incoming chan types.DatagramFrame // frames received from tunnel
34
done chan struct{}
35
36
mu sync.Mutex
@@ -71,7 +44,7 @@ func newQUICBroker(leaseID string) *quicBroker {
44
flowTable: make(map[uint32]*net.UDPAddr),
45
addrIndex: make(map[string]uint32),
46
nextFlow: 1,
74
- incoming: make(chan datagramFrame, 256),
47
+ incoming: make(chan types.DatagramFrame, 256),
48
done: make(chan struct{}),
49
}
50
}
@@ -113,11 +86,11 @@ func (b *quicBroker) SendDatagram(flowID uint32, payload []byte) error {
86
if conn == nil {
87
return errQUICNoConnection
88
}
116
- return conn.SendDatagram(encodeDatagram(flowID, payload))
89
+ return conn.SendDatagram(types.EncodeDatagram(flowID, payload))
90
}
91
92
// Incoming returns the channel that delivers datagrams received from the tunnel.
120
-func (b *quicBroker) Incoming() <-chan datagramFrame {
93
+func (b *quicBroker) Incoming() <-chan types.DatagramFrame {
94
return b.incoming
95
}
96
@@ -189,7 +162,7 @@ func (b *quicBroker) receiveLoop(conn *quic.Conn) {
162
return
163
}
164
192
- frame, err := decodeDatagram(data)
165
+ frame, err := types.DecodeDatagram(data)
166
if err != nil {
167
continue
168
}
portal/server.go
+31
-18
@@ -32,10 +32,8 @@ const (
32
defaultControlBodyLimit = 4 << 20
33
defaultSessionWriteLimit = 5 * time.Second
34
35
- defaultUDPPortMin = 29000
36
- defaultUDPPortMax = 29999
37
- defaultUDPSessionTimeout = 30 * time.Second
38
- defaultMaxDatagramSize = 1350
35
+ defaultUDPPortMin = 29000
36
+ defaultUDPPortMax = 29999
37
)
38
39
type ServerConfig struct {
@@ -99,12 +97,22 @@ func NewServer(cfg ServerConfig) (*Server, error) {
97
cfg.QUICListenAddr = cfg.APIListenAddr
98
}
99
102
- return &Server{
100
+ registry := newLeaseRegistry(policy.NewRuntime())
101
+ ports := newPortAllocator(cfg.UDPPortMin, cfg.UDPPortMax, 5*time.Minute)
102
+
103
+ s := &Server{
104
cfg: cfg,
105
rootHost: rootHost,
105
- registry: newLeaseRegistry(policy.NewRuntime()),
106
- ports: newPortAllocator(cfg.UDPPortMin, cfg.UDPPortMax),
107
- }, nil
106
+ registry: registry,
107
+ ports: ports,
108
+ }
109
+
110
+ // Tear down all lease resources when leases expire via TTL janitor.
111
+ registry.onExpired = func(record *leaseRecord) {
112
+ s.closeLease(record)
113
+ }
114
+
115
+ return s, nil
116
}
117
118
func (s *Server) Start(ctx context.Context, apiMux *http.ServeMux) error {
@@ -182,16 +190,7 @@ func (s *Server) Shutdown(ctx context.Context) error {
190
}
191
192
for _, lease := range s.registry.CloseAll() {
185
- lease.Broker.Close()
186
- if lease.QUICBroker != nil {
187
- lease.QUICBroker.Stop()
188
- }
189
- if lease.UDPRelay != nil {
190
- lease.UDPRelay.Stop()
191
- }
192
- if lease.UDPPort > 0 && s.ports != nil {
193
- s.ports.Release(lease.UDPPort)
194
- }
193
+ s.closeLease(lease)
194
}
195
196
if s.quicTunnel != nil {
@@ -472,6 +471,20 @@ func (s *Server) startUDPRelay(ctx context.Context, leaseID string, relay *udpRe
471
}
472
}
473
474
+// closeLease tears down all resources associated with a single lease record.
475
+func (s *Server) closeLease(record *leaseRecord) {
476
+ record.Broker.Close()
477
+ if record.QUICBroker != nil {
478
+ record.QUICBroker.Stop()
479
+ }
480
+ if record.UDPRelay != nil {
481
+ record.UDPRelay.Stop()
482
+ }
483
+ if record.UDPPort > 0 && s.ports != nil {
484
+ s.ports.Release(record.UDPPort)
485
+ }
486
+}
487
+
488
func (s *Server) quicPublicAddr() string {
489
_, port, err := net.SplitHostPort(s.cfg.QUICListenAddr)
490
if err != nil {
portal/udp_port.go
+77
-10
@@ -2,51 +2,118 @@ package portal
2
3
import (
4
"errors"
5
+ "sort"
6
"sync"
7
+ "time"
8
)
9
10
var errPortExhausted = errors.New("no udp ports available")
11
12
+// portReservation holds a released port for a grace period so the same lease
13
+// name can reclaim it on rapid reconnect.
14
+type portReservation struct {
15
+ port int
16
+ expiresAt time.Time
17
+}
18
+
19
// portAllocator manages a pool of UDP ports for dynamic per-lease allocation.
20
+//
21
+// Features:
22
+// - Sticky allocation: re-registering the same lease name within the grace
23
+// period returns the previously assigned port.
24
+// - Grace period: released ports are held in a reservation map for the
25
+// configured duration before returning to the free pool.
26
+// - Sorted reuse: free ports are kept in ascending order so the lowest
27
+// available port is always allocated first.
28
type portAllocator struct {
12
- available []int
13
- inUse map[int]struct{}
29
+ available []int // sorted ascending
30
+ inUse map[int]string // port → lease name
31
+ reserved map[string]portReservation // lease name → reservation
32
+ grace time.Duration
33
mu sync.Mutex
34
}
35
17
-func newPortAllocator(min, max int) *portAllocator {
36
+func newPortAllocator(min, max int, grace time.Duration) *portAllocator {
37
available := make([]int, 0, max-min+1)
38
for p := min; p <= max; p++ {
39
available = append(available, p)
40
}
41
return &portAllocator{
42
available: available,
24
- inUse: make(map[int]struct{}),
43
+ inUse: make(map[int]string),
44
+ reserved: make(map[string]portReservation),
45
+ grace: grace,
46
}
47
}
48
28
-// Allocate returns the next available port from the pool.
29
-func (a *portAllocator) Allocate() (int, error) {
49
+// Allocate returns a UDP port for the given lease name.
50
+// If the name has a non-expired reservation the same port is returned.
51
+// Otherwise the lowest available port is allocated.
52
+func (a *portAllocator) Allocate(name string) (int, error) {
53
a.mu.Lock()
54
defer a.mu.Unlock()
55
56
+ a.cleanupExpiredLocked(time.Now())
57
+
58
+ // Reclaim reserved port for same name.
59
+ if res, ok := a.reserved[name]; ok {
60
+ delete(a.reserved, name)
61
+ a.inUse[res.port] = name
62
+ return res.port, nil
63
+ }
64
+
65
if len(a.available) == 0 {
66
return 0, errPortExhausted
67
}
68
+
69
port := a.available[0]
70
a.available = a.available[1:]
38
- a.inUse[port] = struct{}{}
71
+ a.inUse[port] = name
72
return port, nil
73
}
74
42
-// Release returns a port back to the available pool.
75
+// Release moves a port from in-use to reserved state. The port is held for
76
+// the grace period so the same lease name can reclaim it.
77
func (a *portAllocator) Release(port int) {
78
a.mu.Lock()
79
defer a.mu.Unlock()
80
47
- if _, ok := a.inUse[port]; !ok {
81
+ name, ok := a.inUse[port]
82
+ if !ok {
83
return
84
}
85
delete(a.inUse, port)
51
- a.available = append(a.available, port)
86
+
87
+ // If the name already has a different reservation (shouldn't happen in
88
+ // normal flow), return that old port to the free pool first.
89
+ if prev, exists := a.reserved[name]; exists {
90
+ a.sortedInsertLocked(prev.port)
91
+ }
92
+
93
+ a.reserved[name] = portReservation{
94
+ port: port,
95
+ expiresAt: time.Now().Add(a.grace),
96
+ }
97
+
98
+ a.cleanupExpiredLocked(time.Now())
99
+}
100
+
101
+// cleanupExpiredLocked moves expired reservations back to the sorted available
102
+// pool. Caller must hold a.mu.
103
+func (a *portAllocator) cleanupExpiredLocked(now time.Time) {
104
+ for name, res := range a.reserved {
105
+ if now.After(res.expiresAt) {
106
+ delete(a.reserved, name)
107
+ a.sortedInsertLocked(res.port)
108
+ }
109
+ }
110
+}
111
+
112
+// sortedInsertLocked inserts port into a.available maintaining ascending order.
113
+// Caller must hold a.mu.
114
+func (a *portAllocator) sortedInsertLocked(port int) {
115
+ i := sort.SearchInts(a.available, port)
116
+ a.available = append(a.available, 0)
117
+ copy(a.available[i+1:], a.available[i:])
118
+ a.available[i] = port
119
}
portal/udp_relay.go
+22
-1
@@ -10,6 +10,11 @@ import (
10
"github.com/rs/zerolog/log"
11
)
12
13
+const (
14
+ defaultUDPSessionTimeout = 30 * time.Second
15
+ defaultMaxDatagramSize = 1350
16
+)
17
+
18
// udpSession tracks one client endpoint sending to a per-lease UDP listener.
19
type udpSession struct {
20
FlowID uint32
@@ -105,6 +110,11 @@ func (r *udpRelay) readLoop(ctx context.Context) {
110
if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
111
continue
112
}
113
+ log.Warn().
114
+ Str("component", "udp-relay").
115
+ Str("lease_id", r.leaseID).
116
+ Err(err).
117
+ Msg("readLoop exiting: unexpected read error")
118
return
119
}
120
@@ -113,7 +123,13 @@ func (r *udpRelay) readLoop(ctx context.Context) {
123
copy(payload, buf[:n])
124
125
if err := r.broker.SendDatagram(flowID, payload); err != nil {
116
- // Tunnel not connected yet — drop silently.
126
+ log.Warn().
127
+ Str("component", "udp-relay").
128
+ Str("lease_id", r.leaseID).
129
+ Err(err).
130
+ Uint32("flow_id", flowID).
131
+ Int("bytes", n).
132
+ Msg("send datagram to tunnel failed, dropping packet")
133
continue
134
}
135
}
@@ -128,6 +144,11 @@ func (r *udpRelay) writeLoop(ctx context.Context) {
144
case frame := <-r.broker.Incoming():
145
addr, ok := r.broker.LookupFlowAddr(frame.FlowID)
146
if !ok {
147
+ log.Debug().
148
+ Str("component", "udp-relay").
149
+ Str("lease_id", r.leaseID).
150
+ Uint32("flow_id", frame.FlowID).
151
+ Msg("write loop: unknown flow id, dropping")
152
continue
153
}
154
_, _ = r.conn.WriteToUDP(frame.Payload, addr)
sdk/api_client.go
+2
-15
@@ -90,13 +90,14 @@ func (a *apiClient) close() {
90
}
91
}
92
93
-func (a *apiClient) registerLease(ctx context.Context, ttl time.Duration) (types.RegisterResponse, error) {
93
+func (a *apiClient) registerLease(ctx context.Context, ttl time.Duration, transport string) (types.RegisterResponse, error) {
94
var resp types.RegisterResponse
95
if err := a.doJSON(ctx, http.MethodPost, types.PathSDKRegister, types.RegisterRequest{
96
Name: a.name,
97
Metadata: a.metadata.Copy(),
98
ReverseToken: a.reverseToken,
99
TTL: int(ttl / time.Second),
100
+ Transport: transport,
101
}, &resp); err != nil {
102
return types.RegisterResponse{}, err
103
}
@@ -294,20 +295,6 @@ func (c *bufferedConn) Read(p []byte) (int, error) {
295
return c.Conn.Read(p)
296
}
297
297
-func (a *apiClient) registerLeaseWithTransport(ctx context.Context, ttl time.Duration, transport string) (types.RegisterResponse, error) {
298
- var resp types.RegisterResponse
299
- if err := a.doJSON(ctx, http.MethodPost, types.PathSDKRegister, types.RegisterRequest{
300
- Name: a.name,
301
- Metadata: a.metadata.Copy(),
302
- ReverseToken: a.reverseToken,
303
- TTL: int(ttl / time.Second),
304
- Transport: transport,
305
- }, &resp); err != nil {
306
- return types.RegisterResponse{}, err
307
- }
308
- return resp, nil
309
-}
310
-
298
// openQUICSession opens a QUIC connection to the relay for datagram transport.
299
func (a *apiClient) openQUICSession(ctx context.Context, quicAddr, leaseID, reverseToken string) (*quic.Conn, error) {
300
tlsConf := a.rawTLSConfig.Clone()
sdk/expose.go
+48
-9
@@ -31,7 +31,7 @@ type Exposure struct {
31
// Expose creates relay listeners for each normalized relay URL and exposes a
32
// merged listener for accepting traffic from all of them. Empty relay input
33
// returns nil, nil so callers can fall back to local-only serving.
34
-func Expose(ctx context.Context, relayUrls []string, name string, metadata types.LeaseMetadata) (*Exposure, error) {
34
+func Expose(ctx context.Context, relayUrls []string, name string, transport string, metadata types.LeaseMetadata) (*Exposure, error) {
35
relayURLs, err := utils.NormalizeRelayURLs(relayUrls)
36
if err != nil {
37
return nil, err
@@ -52,8 +52,9 @@ func Expose(ctx context.Context, relayUrls []string, name string, metadata types
52
53
for _, relayURL := range relayURLs {
54
listener, err := NewListener(ctx, relayURL, ListenerConfig{
55
- Name: name,
56
- Metadata: metadata,
55
+ Name: name,
56
+ Transport: transport,
57
+ Metadata: metadata,
58
})
59
if err != nil {
60
return nil, errors.Join(fmt.Errorf("listen %q: %w", relayURL, err), cleanup())
@@ -190,17 +191,55 @@ func (e *Exposure) AttachUDP(ctx context.Context) ([]*UDPListener, error) {
191
return nil, nil
192
}
193
193
- var out []*UDPListener
194
+ // Attach concurrently so a slow/failing relay does not block the rest.
195
+ // Return collected results as soon as timeout fires or all complete.
196
+ waitCtx, cancel := context.WithTimeout(ctx, 15*time.Second)
197
+ defer cancel()
198
+
199
+ results := make(chan *UDPListener, len(e.relays))
200
+
201
+ count := 0
202
for _, relay := range e.relays {
203
if relay.listener == nil {
204
continue
205
}
198
- udpL, err := relay.listener.AttachUDP(ctx)
199
- if err != nil {
200
- // Skip relays that don't support UDP.
201
- continue
206
+ count++
207
+ go func(l *Listener) {
208
+ udpL, err := l.AttachUDP(waitCtx)
209
+ if err != nil {
210
+ results <- nil
211
+ return
212
+ }
213
+ results <- udpL
214
+ }(relay.listener)
215
+ }
216
+
217
+ var out []*UDPListener
218
+ var graceTimer *time.Timer
219
+ defer func() {
220
+ if graceTimer != nil {
221
+ graceTimer.Stop()
222
+ }
223
+ }()
224
+ var grace <-chan time.Time
225
+ collected := 0
226
+ for collected < count {
227
+ select {
228
+ case l := <-results:
229
+ collected++
230
+ if l != nil {
231
+ out = append(out, l)
232
+ if graceTimer == nil {
233
+ // After first success, give 3 more seconds for remaining relays.
234
+ graceTimer = time.NewTimer(3 * time.Second)
235
+ grace = graceTimer.C
236
+ }
237
+ }
238
+ case <-grace:
239
+ return out, nil
240
+ case <-waitCtx.Done():
241
+ return out, nil
242
}
203
- out = append(out, udpL)
243
}
244
return out, nil
245
}
sdk/listener.go
+19
-9
@@ -65,8 +65,10 @@ type Listener struct {
65
quicAddr string
66
metadata types.LeaseMetadata
67
68
- closeOnce sync.Once
69
- mu sync.Mutex
68
+ registered chan struct{} // closed after first successful registration
69
+ closeOnce sync.Once
70
+ registerOnce sync.Once
71
+ mu sync.Mutex
72
}
73
74
// NewListener creates one relay listener and its dedicated relay transport for one relay URL.
@@ -96,6 +98,7 @@ func NewListener(ctx context.Context, relayURL string, cfg ListenerConfig) (*Lis
98
cancel: cancel,
99
api: api,
100
accepted: make(chan net.Conn, max(readyTarget*2, 1)),
101
+ registered: make(chan struct{}),
102
relayURL: api.baseURL.String(),
103
transport: transport,
104
startupStatus: listenerStatusInactive,
@@ -381,13 +384,7 @@ func (l *Listener) registerAndConfigure(ctx context.Context) error {
384
return err
385
}
386
384
- var resp types.RegisterResponse
385
- var err error
386
- if l.transport != "" {
387
- resp, err = l.api.registerLeaseWithTransport(ctx, l.leaseTTL, l.transport)
388
- } else {
389
- resp, err = l.api.registerLease(ctx, l.leaseTTL)
390
- }
387
+ resp, err := l.api.registerLease(ctx, l.leaseTTL, l.transport)
388
if err != nil {
389
return err
390
}
@@ -424,9 +421,22 @@ func (l *Listener) registerAndConfigure(ctx context.Context) error {
421
if oldCloser != nil {
422
_ = oldCloser.Close()
423
}
424
+ l.registerOnce.Do(func() { close(l.registered) })
425
return nil
426
}
427
428
+// WaitRegistered blocks until the first successful lease registration or context cancellation.
429
+func (l *Listener) WaitRegistered(ctx context.Context) error {
430
+ select {
431
+ case <-l.registered:
432
+ return nil
433
+ case <-l.doneCh:
434
+ return net.ErrClosed
435
+ case <-ctx.Done():
436
+ return ctx.Err()
437
+ }
438
+}
439
+
440
func (l *Listener) reregister(ctx context.Context) error {
441
requestCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
442
defer cancel()
sdk/sdk_test.go
+1
-1
@@ -347,7 +347,7 @@ func TestNewListenerRetriesForeverWhenRetryCountIsNegative(t *testing.T) {
347
}
348
349
func TestExposeNoRelayInputs(t *testing.T) {
350
- exposure, err := Expose(context.Background(), nil, "demo", types.LeaseMetadata{})
350
+ exposure, err := Expose(context.Background(), nil, "demo", types.TransportBoth, types.LeaseMetadata{})
351
if err != nil {
352
t.Fatalf("Expose() error = %v", err)
353
}
sdk/udp_listener.go
+12
-30
@@ -2,7 +2,6 @@ package sdk
2
3
import (
4
"context"
5
- "encoding/binary"
5
"errors"
6
"fmt"
7
"net"
@@ -85,7 +84,7 @@ func (l *UDPListener) SendDatagram(flowID uint32, payload []byte) error {
84
if conn == nil {
85
return errors.New("quic connection not established")
86
}
88
- return conn.SendDatagram(encodeDatagram(flowID, payload))
87
+ return conn.SendDatagram(types.EncodeDatagram(flowID, payload))
88
}
89
90
// UDPAddr returns the public UDP address allocated by the relay.
@@ -170,7 +169,7 @@ func NewUDPListener(ctx context.Context, relayURL string, cfg UDPListenerConfig)
169
return nil, err
170
}
171
173
- registerResp, err := api.registerLeaseWithTransport(ctx, leaseTTL, transport)
172
+ registerResp, err := api.registerLease(ctx, leaseTTL, transport)
173
if err != nil {
174
api.close()
175
return nil, err
@@ -206,6 +205,11 @@ func NewUDPListener(ctx context.Context, relayURL string, cfg UDPListenerConfig)
205
// The caller (the TCP Listener) owns the lease lifecycle.
206
// The Listener must have been registered with transport "udp" or "both".
207
func (l *Listener) AttachUDP(ctx context.Context) (*UDPListener, error) {
208
+ // Wait for lease registration to complete before reading UDP addresses.
209
+ if err := l.WaitRegistered(ctx); err != nil {
210
+ return nil, fmt.Errorf("wait for registration: %w", err)
211
+ }
212
+
213
l.mu.Lock()
214
leaseID := l.leaseID
215
api := l.api
@@ -223,7 +227,10 @@ func (l *Listener) AttachUDP(ctx context.Context) (*UDPListener, error) {
227
return nil, errors.New("lease does not have UDP transport enabled")
228
}
229
226
- listenerCtx, cancel := context.WithCancel(ctx)
230
+ // Use context.Background — the caller's ctx may be a short-lived timeout
231
+ // context (e.g. the 15-second waitCtx from Exposure.AttachUDP). The
232
+ // UDPListener's lifecycle is managed by Close(), not context cancellation.
233
+ listenerCtx, cancel := context.WithCancel(context.Background())
234
udpL := &UDPListener{
235
api: api,
236
baseContext: func() context.Context { return listenerCtx },
@@ -300,7 +307,7 @@ func (l *UDPListener) receiveLoop(conn *quic.Conn) {
307
return
308
}
309
303
- frame, err := decodeDatagram(data)
310
+ frame, err := types.DecodeDatagram(data)
311
if err != nil {
312
continue
313
}
@@ -362,28 +369,3 @@ func (l *UDPListener) isClosed() bool {
369
}
370
}
371
365
-// datagramFrame mirrors portal.datagramFrame for SDK-side decode/encode.
366
-type datagramFrame struct {
367
- FlowID uint32
368
- Payload []byte
369
-}
370
-
371
-func encodeDatagram(flowID uint32, payload []byte) []byte {
372
- var buf [binary.MaxVarintLen32]byte
373
- n := binary.PutUvarint(buf[:], uint64(flowID))
374
- out := make([]byte, n+len(payload))
375
- copy(out, buf[:n])
376
- copy(out[n:], payload)
377
- return out
378
-}
379
-
380
-func decodeDatagram(data []byte) (datagramFrame, error) {
381
- flowID, n := binary.Uvarint(data)
382
- if n <= 0 {
383
- return datagramFrame{}, fmt.Errorf("datagram too small to decode")
384
- }
385
- return datagramFrame{
386
- FlowID: uint32(flowID),
387
- Payload: data[n:],
388
- }, nil
389
-}
types/datagram.go
new
+39
@@ -0,0 +1,39 @@
1
+package types
2
+
3
+import (
4
+ "encoding/binary"
5
+ "errors"
6
+)
7
+
8
+// ErrDatagramTooSmall is returned when a datagram payload is too short to
9
+// contain a valid flow ID varint.
10
+var ErrDatagramTooSmall = errors.New("datagram too small to decode")
11
+
12
+// DatagramFrame is the wire format for QUIC DATAGRAM payloads.
13
+// Layout: [flowID varint][payload bytes]
14
+type DatagramFrame struct {
15
+ FlowID uint32
16
+ Payload []byte
17
+}
18
+
19
+// EncodeDatagram serialises a flow-framed datagram for transmission.
20
+func EncodeDatagram(flowID uint32, payload []byte) []byte {
21
+ var buf [binary.MaxVarintLen32]byte
22
+ n := binary.PutUvarint(buf[:], uint64(flowID))
23
+ out := make([]byte, n+len(payload))
24
+ copy(out, buf[:n])
25
+ copy(out[n:], payload)
26
+ return out
27
+}
28
+
29
+// DecodeDatagram deserialises a flow-framed datagram.
30
+func DecodeDatagram(data []byte) (DatagramFrame, error) {
31
+ flowID, n := binary.Uvarint(data)
32
+ if n <= 0 {
33
+ return DatagramFrame{}, ErrDatagramTooSmall
34
+ }
35
+ return DatagramFrame{
36
+ FlowID: uint32(flowID),
37
+ Payload: data[n:],
38
+ }, nil
39
+}