@samitouri / QOSAMI-X402Tunnel-1 / commits / f28a6868

refactor datagram

Kim committed Mar 18, 2026 at 11:44 UTC f28a68680c4319300c84454f6bc2408982a6de2a
27 files changed +1946 -1477
cmd/portal-tunnel/main.go
+39 -28
@@ -226,6 +226,11 @@ func runTunnel(
226 metadata types.LeaseMetadata,
227 ) error {
228 logger := log.With().Str("component", "portal").Logger()
229 + capabilities, err := types.ParseLeaseCapabilities(transport)
230 + if err != nil {
231 + return err
232 + }
233 + transport = capabilities.Transport()
234
235 exposure, err := sdk.Expose(ctx, relayURLs, name, transport, metadata)
236 if err != nil {
@@ -237,8 +242,21 @@ func runTunnel(
242 defer exposure.Close()
243
244 // UDP is best-effort — attach to existing lease, log and continue if it fails.
240 - if transport == types.TransportUDP || transport == types.TransportBoth {
241 - go runUDPBestEffort(ctx, exposure, target)
245 + if capabilities.SupportsDatagram() && !capabilities.SupportsStream() {
246 + runErr := runUDPBestEffort(ctx, exposure, target)
247 + closeErr := exposure.Close()
248 + if runErr != nil && stop != nil {
249 + stop()
250 + }
251 + return errors.Join(runErr, closeErr)
252 + }
253 +
254 + if capabilities.SupportsDatagram() {
255 + go func() {
256 + if err := runUDPBestEffort(ctx, exposure, target); err != nil && ctx.Err() == nil {
257 + logger.Warn().Err(err).Msg("udp transport disabled")
258 + }
259 + }()
260 }
261
262 logger.Info().
@@ -288,39 +306,32 @@ func runTunnel(
306 return errors.Join(waitErr, closeErr)
307 }
308
291 -// runUDPBestEffort attaches UDP listeners to the existing TCP exposure lease.
292 -// Failures are logged but do not bring down the tunnel.
293 -func runUDPBestEffort(ctx context.Context, exposure *sdk.Exposure, target string) {
309 +// runUDPBestEffort waits for the exposure datagram plane and proxies it to the
310 +// local UDP target.
311 +func runUDPBestEffort(ctx context.Context, exposure *sdk.Exposure, target string) error {
312 logger := log.With().Str("component", "portal-tunnel-udp").Logger()
313
296 - udpListeners, err := exposure.AttachUDP(ctx)
314 + udpAddrs, err := exposure.WaitDatagramReady(ctx)
315 if err != nil {
298 - logger.Warn().Err(err).Msg("udp transport disabled: attach failed")
299 - return
316 + if ctx.Err() != nil || errors.Is(err, context.Canceled) {
317 + return ctx.Err()
318 + }
319 + return fmt.Errorf("wait for udp readiness: %w", err)
320 }
301 - if len(udpListeners) == 0 {
302 - logger.Info().Msg("udp transport: no UDP addresses from relay")
303 - return
321 + if len(udpAddrs) == 0 {
322 + if ctx.Err() != nil {
323 + return ctx.Err()
324 + }
325 + return errors.New("relay did not expose any UDP listeners")
326 }
327
306 - var wg sync.WaitGroup
307 - for _, ul := range udpListeners {
308 - wg.Add(1)
309 - go func(l *sdk.UDPListener) {
310 - defer wg.Done()
311 - defer l.Close()
312 -
313 - logger.Info().
314 - Str("udp_addr", l.UDPAddr()).
315 - Str("lease_id", l.LeaseID()).
316 - Msg("UDP tunnel ready")
317 -
318 - if err := proxyUDPRelayConnections(ctx, l, target); err != nil {
319 - logger.Warn().Err(err).Msg("udp proxy ended")
320 - }
321 - }(ul)
328 + for _, udpAddr := range udpAddrs {
329 + logger.Info().
330 + Str("udp_addr", udpAddr).
331 + Msg("UDP tunnel ready")
332 }
323 - wg.Wait()
333 +
334 + return proxyExposureDatagrams(ctx, exposure, target)
335 }
336
337 func resolveRelayURLs(ctx context.Context, registryURL string, inputs []string, includeDefaultRelays bool) ([]string, error) {
cmd/portal-tunnel/relays.go
+72 -24
@@ -134,9 +134,9 @@ func writeEmptyHTTPResponse(conn net.Conn) error {
134 return err
135 }
136
137 -// proxyUDPRelayConnections receives datagrams from the relay via the UDPListener
137 +// proxyExposureDatagrams receives datagrams from the exposure datagram plane
138 // and forwards them to the local UDP service, relaying responses back.
139 -func proxyUDPRelayConnections(ctx context.Context, udpListener *sdk.UDPListener, localAddr string) error {
139 +func proxyExposureDatagrams(ctx context.Context, exposure *sdk.Exposure, localAddr string) error {
140 logger := log.With().Str("component", "portal-tunnel-udp").Logger()
141
142 targetAddr, err := utils.NormalizeTargetAddr(localAddr)
@@ -149,15 +149,20 @@ func proxyUDPRelayConnections(ctx context.Context, udpListener *sdk.UDPListener,
149 return fmt.Errorf("resolve udp addr %q: %w", targetAddr, err)
150 }
151
152 - // Per-flow local UDP connections: flowID → *net.UDPConn
152 + type flowKey struct {
153 + flowID uint32
154 + leaseID string
155 + relayURL string
156 + }
157 type flowEntry struct {
158 conn *net.UDPConn
159 lastSeen time.Time
160 + reply func([]byte) error
161 }
162 +
163 var mu sync.Mutex
158 - flows := make(map[uint32]*flowEntry)
164 + flows := make(map[flowKey]*flowEntry)
165
160 - // Cleanup idle flow connections.
166 go func() {
167 ticker := time.NewTicker(15 * time.Second)
168 defer ticker.Stop()
@@ -168,10 +173,10 @@ func proxyUDPRelayConnections(ctx context.Context, udpListener *sdk.UDPListener,
173 case <-ticker.C:
174 mu.Lock()
175 now := time.Now()
171 - for id, f := range flows {
176 + for key, f := range flows {
177 if now.Sub(f.lastSeen) > 30*time.Second {
178 _ = f.conn.Close()
174 - delete(flows, id)
179 + delete(flows, key)
180 }
181 }
182 mu.Unlock()
@@ -179,11 +184,11 @@ func proxyUDPRelayConnections(ctx context.Context, udpListener *sdk.UDPListener,
184 }
185 }()
186
182 - // getOrCreateFlow returns (or creates) a local UDP conn for a flow.
183 - getOrCreateFlow := func(flowID uint32) (*net.UDPConn, error) {
187 + getOrCreateFlow := func(key flowKey, reply func([]byte) error) (*net.UDPConn, error) {
188 mu.Lock()
185 - if f, ok := flows[flowID]; ok {
189 + if f, ok := flows[key]; ok {
190 f.lastSeen = time.Now()
191 + f.reply = reply
192 mu.Unlock()
193 return f.conn, nil
194 }
@@ -195,17 +200,16 @@ func proxyUDPRelayConnections(ctx context.Context, udpListener *sdk.UDPListener,
200 }
201
202 mu.Lock()
198 - // Double-check after acquiring lock.
199 - if f, ok := flows[flowID]; ok {
203 + if f, ok := flows[key]; ok {
204 mu.Unlock()
205 _ = localConn.Close()
206 f.lastSeen = time.Now()
207 + f.reply = reply
208 return f.conn, nil
209 }
205 - flows[flowID] = &flowEntry{conn: localConn, lastSeen: time.Now()}
210 + flows[key] = &flowEntry{conn: localConn, lastSeen: time.Now(), reply: reply}
211 mu.Unlock()
212
208 - // Start reverse read loop: local service → relay.
213 go func() {
214 buf := make([]byte, 65535)
215 for {
@@ -214,11 +218,33 @@ func proxyUDPRelayConnections(ctx context.Context, udpListener *sdk.UDPListener,
218 if ctx.Err() != nil {
219 return
220 }
217 - logger.Debug().Err(err).Uint32("flow_id", flowID).Msg("local read ended")
221 + logger.Debug().
222 + Err(err).
223 + Uint32("flow_id", key.flowID).
224 + Str("lease_id", key.leaseID).
225 + Str("relay_url", key.relayURL).
226 + Msg("local read ended")
227 return
228 }
220 - if sendErr := udpListener.SendDatagram(flowID, buf[:n]); sendErr != nil {
221 - logger.Debug().Err(sendErr).Uint32("flow_id", flowID).Msg("send datagram to relay failed")
229 +
230 + mu.Lock()
231 + entry := flows[key]
232 + if entry != nil {
233 + entry.lastSeen = time.Now()
234 + }
235 + replyFn := func([]byte) error { return net.ErrClosed }
236 + if entry != nil && entry.reply != nil {
237 + replyFn = entry.reply
238 + }
239 + mu.Unlock()
240 +
241 + if sendErr := replyFn(buf[:n]); sendErr != nil {
242 + logger.Debug().
243 + Err(sendErr).
244 + Uint32("flow_id", key.flowID).
245 + Str("lease_id", key.leaseID).
246 + Str("relay_url", key.relayURL).
247 + Msg("send datagram to relay failed")
248 return
249 }
250 }
@@ -227,13 +253,15 @@ func proxyUDPRelayConnections(ctx context.Context, udpListener *sdk.UDPListener,
253 return localConn, nil
254 }
255
230 - // Main loop: relay → local service.
256 logger.Info().Str("target", targetAddr).Msg("udp proxy loop started, waiting for datagrams")
257 for {
233 - dg, err := udpListener.AcceptDatagram()
258 + dg, err := exposure.AcceptDatagram()
259 if err != nil {
235 - if ctx.Err() != nil || errors.Is(err, net.ErrClosed) {
236 - return nil
260 + if ctx.Err() != nil {
261 + return ctx.Err()
262 + }
263 + if errors.Is(err, net.ErrClosed) {
264 + break
265 }
266 return fmt.Errorf("accept datagram: %w", err)
267 }
@@ -241,17 +269,37 @@ func proxyUDPRelayConnections(ctx context.Context, udpListener *sdk.UDPListener,
269 logger.Debug().
270 Uint32("flow_id", dg.FlowID).
271 Int("bytes", len(dg.Payload)).
272 + Str("lease_id", dg.LeaseID).
273 + Str("relay_url", dg.RelayURL).
274 + Str("udp_addr", dg.UDPAddr).
275 Str("target", targetAddr).
276 Msg("datagram received from relay, forwarding to local")
277
247 - localConn, err := getOrCreateFlow(dg.FlowID)
278 + key := flowKey{
279 + flowID: dg.FlowID,
280 + leaseID: dg.LeaseID,
281 + relayURL: dg.RelayURL,
282 + }
283 + localConn, err := getOrCreateFlow(key, dg.Reply)
284 if err != nil {
249 - logger.Warn().Err(err).Uint32("flow_id", dg.FlowID).Msg("dial local udp failed")
285 + logger.Warn().
286 + Err(err).
287 + Uint32("flow_id", dg.FlowID).
288 + Str("lease_id", dg.LeaseID).
289 + Str("relay_url", dg.RelayURL).
290 + Msg("dial local udp failed")
291 continue
292 }
293
294 if _, err := localConn.Write(dg.Payload); err != nil {
254 - logger.Warn().Err(err).Uint32("flow_id", dg.FlowID).Msg("write to local udp failed")
295 + logger.Warn().
296 + Err(err).
297 + Uint32("flow_id", dg.FlowID).
298 + Str("lease_id", dg.LeaseID).
299 + Str("relay_url", dg.RelayURL).
300 + Msg("write to local udp failed")
301 }
302 }
303 +
304 + return nil
305 }
go.mod
+2 -1
@@ -9,7 +9,9 @@ require (
9 github.com/aws/aws-sdk-go-v2/service/route53 v1.62.1
10 github.com/go-acme/lego/v4 v4.32.0
11 github.com/gosuda/keyless_tls v0.0.1-0.20260304212324-7733f8366abc
12 + github.com/quic-go/quic-go v0.59.0
13 github.com/rs/zerolog v1.34.0
14 + golang.org/x/crypto v0.48.0
15 golang.org/x/net v0.50.0
16 golang.org/x/sync v0.19.0
17 )
@@ -31,7 +33,6 @@ require (
33 github.com/mattn/go-colorable v0.1.13 // indirect
34 github.com/mattn/go-isatty v0.0.20 // indirect
35 github.com/miekg/dns v1.1.72 // indirect
34 - golang.org/x/crypto v0.48.0 // indirect
36 golang.org/x/mod v0.32.0 // indirect
37 golang.org/x/sys v0.41.0 // indirect
38 golang.org/x/text v0.34.0 // indirect
go.sum
+4
@@ -53,11 +53,15 @@ github.com/miekg/dns v1.1.72/go.mod h1:+EuEPhdHOsfk6Wk5TT2CzssZdqkmFhf8r+aVyDETo
53 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
54 github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U=
55 github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
56 +github.com/quic-go/quic-go v0.59.0 h1:OLJkp1Mlm/aS7dpKgTc6cnpynnD2Xg7C1pwL6vy/SAw=
57 +github.com/quic-go/quic-go v0.59.0/go.mod h1:upnsH4Ju1YkqpLXC305eW3yDZ4NfnNbmQRCMWS58IKU=
58 github.com/rs/xid v1.6.0/go.mod h1:7XoLgs4eV+QndskICGsho+ADou8ySMSjJKDIan90Nz0=
59 github.com/rs/zerolog v1.34.0 h1:k43nTLIwcTVQAncfCw4KZ2VY6ukYoZaBPNOE8txlOeY=
60 github.com/rs/zerolog v1.34.0/go.mod h1:bJsvje4Z08ROH4Nhs5iH600c3IkWhwp44iRc54W6wYQ=
61 github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
62 github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
63 +go.uber.org/mock v0.5.2 h1:LbtPTcP8A5k9WPXj54PPPbjcI4Y6lhyOZXn+VS7wNko=
64 +go.uber.org/mock v0.5.2/go.mod h1:wLlUxC2vVTPTaE3UD51E0BGOAElKrILxhVSDYQLld5o=
65 golang.org/x/crypto v0.48.0 h1:/VRzVqiRSggnhY7gNRxPauEQ5Drw9haKdM0jqfcCFts=
66 golang.org/x/crypto v0.48.0/go.mod h1:r0kV5h3qnFPlQnBSrULhlsRfryS2pmewsg+XfMgkVos=
67 golang.org/x/mod v0.32.0 h1:9F4d3PHLljb6x//jOyokMv3eX+YDeepZSEo3mFJy93c=
portal/api_server.go
+40 -30
@@ -1,7 +1,6 @@
1 package portal
2
3 import (
4 - "context"
4 "crypto/tls"
5 "errors"
6 "fmt"
@@ -13,6 +12,7 @@ import (
12
13 "github.com/rs/zerolog/log"
14
15 + "github.com/gosuda/portal/v2/portal/datagram"
16 "github.com/gosuda/portal/v2/portal/keyless"
17 "github.com/gosuda/portal/v2/portal/policy"
18 "github.com/gosuda/portal/v2/types"
@@ -127,6 +127,9 @@ func (s *Server) handleRegister(w http.ResponseWriter, r *http.Request) {
127 if errors.Is(err, errIPBanned) {
128 status, code = http.StatusForbidden, types.APIErrorCodeIPBanned
129 }
130 + if errors.Is(err, datagram.ErrPortExhausted) {
131 + status, code = http.StatusServiceUnavailable, types.APIErrorCodeUDPPortExhausted
132 + }
133 utils.WriteAPIError(w, status, code, err.Error())
134 return
135 }
@@ -229,6 +232,10 @@ func (s *Server) handleConnect(w http.ResponseWriter, r *http.Request) {
232 utils.WriteAPIError(w, http.StatusForbidden, types.APIErrorCodeUnauthorized, authErr.Error())
233 return
234 }
235 + if !lease.SupportsStream() {
236 + utils.WriteAPIError(w, http.StatusConflict, types.APIErrorCodeTransportMismatch, "lease does not support stream transport")
237 + return
238 + }
239
240 hijacker, ok := w.(http.Hijacker)
241 if !ok {
@@ -252,7 +259,13 @@ func (s *Server) handleConnect(w http.ResponseWriter, r *http.Request) {
259 }
260
261 session := newReverseSession(conn, s.cfg.IdleKeepaliveInterval)
255 - if err := lease.Broker.Offer(session); err != nil {
262 + streamBroker := lease.StreamBroker()
263 + if streamBroker == nil {
264 + _ = session.Close()
265 + return
266 + }
267 +
268 + if err := streamBroker.Offer(session); err != nil {
269 log.Warn().
270 Err(err).
271 Str("component", "relay-server").
@@ -270,7 +283,7 @@ func (s *Server) handleConnect(w http.ResponseWriter, r *http.Request) {
283 Str("lease_id", lease.ID).
284 Str("lease_name", lease.Name).
285 Str("remote_addr", session.RemoteAddr()).
273 - Int("ready", lease.Broker.ReadyCount()).
286 + Int("ready", streamBroker.ReadyCount()).
287 Msg("sdk reverse connected")
288 }
289
@@ -295,14 +308,27 @@ func (s *Server) registerLease(req types.RegisterRequest, clientIP string) (type
308 ttl = time.Duration(req.TTL) * time.Second
309 }
310
298 - transport := strings.ToLower(strings.TrimSpace(req.Transport))
299 - if transport == "" {
300 - transport = types.TransportTCP
311 + capabilities, err := types.ParseLeaseCapabilities(req.Transport)
312 + if err != nil {
313 + return types.RegisterResponse{}, err
314 }
315 + transport := capabilities.Transport()
316
317 leaseID := utils.RandomID("lease_")
318 now := time.Now()
319 expiresAt := now.Add(ttl)
320 + runtime, err := newLeaseRuntime(leaseRuntimeConfig{
321 + Capabilities: capabilities,
322 + IdleInterval: s.cfg.IdleKeepaliveInterval,
323 + LeaseID: leaseID,
324 + LeaseName: name,
325 + PortAllocator: s.ports,
326 + ReadyLimit: s.cfg.ReadyQueueLimit,
327 + })
328 + if err != nil {
329 + return types.RegisterResponse{}, err
330 + }
331 +
332 record := &leaseRecord{
333 Lease: types.Lease{
334 ID: leaseID,
@@ -314,38 +340,22 @@ func (s *Server) registerLease(req types.RegisterRequest, clientIP string) (type
340 LastSeenAt: now,
341 ClientIP: clientIP,
342 Transport: transport,
343 + UDPPort: runtime.UDPPort(),
344 },
345 ReverseToken: req.ReverseToken,
319 - Broker: newLeaseBroker(leaseID, s.cfg.IdleKeepaliveInterval, s.cfg.ReadyQueueLimit),
346 + Runtime: runtime,
347 }
348
322 - if transport == types.TransportUDP || transport == types.TransportBoth {
323 - if s.ports == nil {
324 - return types.RegisterResponse{}, errors.New("udp port allocation not available")
325 - }
326 - udpPort, portErr := s.ports.Allocate(name)
327 - if portErr != nil {
328 - return types.RegisterResponse{}, fmt.Errorf("allocate udp port: %w", portErr)
329 - }
330 - record.UDPPort = udpPort
331 - record.QUICBroker = newQUICBroker(leaseID)
332 - record.UDPRelay = newUDPRelay(leaseID, udpPort, record.QUICBroker)
349 + if err := runtime.Start(); err != nil {
350 + runtime.Close(s.ports)
351 + return types.RegisterResponse{}, err
352 }
353
354 if err := s.registry.Register(record); err != nil {
336 - if record.QUICBroker != nil {
337 - record.QUICBroker.Stop()
338 - }
339 - if record.UDPPort > 0 && s.ports != nil {
340 - s.ports.Release(record.UDPPort)
341 - }
355 + runtime.Close(s.ports)
356 return types.RegisterResponse{}, err
357 }
358
345 - if record.UDPRelay != nil {
346 - go s.startUDPRelay(context.Background(), leaseID, record.UDPRelay)
347 - }
348 -
359 resp := types.RegisterResponse{
360 LeaseID: leaseID,
361 Hostname: hostname,
@@ -354,8 +364,8 @@ func (s *Server) registerLease(req types.RegisterRequest, clientIP string) (type
364 ConnectURL: strings.TrimRight(s.cfg.PortalURL, "/") + types.PathSDKConnect,
365 Transport: transport,
366 }
357 - if record.UDPPort > 0 {
358 - resp.UDPAddr = fmt.Sprintf("%s:%d", s.rootHost, record.UDPPort)
367 + if record.SupportsDatagram() {
368 + resp.UDPAddr = fmt.Sprintf("%s:%d", s.rootHost, record.UDPPort())
369 resp.QUICAddr = s.quicPublicAddr()
370 }
371
portal/broker.go
+9 -9
@@ -18,7 +18,7 @@ var (
18 errBrokerFull = errors.New("broker ready queue full")
19 )
20
21 -type leaseBroker struct {
21 +type streamBroker struct {
22 notify chan struct{}
23 leaseID string
24 ready []*reverseSession
@@ -28,8 +28,8 @@ type leaseBroker struct {
28 mu sync.Mutex
29 }
30
31 -func newLeaseBroker(leaseID string, idleInterval time.Duration, readyLimit int) *leaseBroker {
32 - return &leaseBroker{
31 +func newStreamBroker(leaseID string, idleInterval time.Duration, readyLimit int) *streamBroker {
32 + return &streamBroker{
33 leaseID: leaseID,
34 idleInterval: idleInterval,
35 readyLimit: readyLimit,
@@ -37,7 +37,7 @@ func newLeaseBroker(leaseID string, idleInterval time.Duration, readyLimit int)
37 }
38 }
39
40 -func (b *leaseBroker) Offer(session *reverseSession) error {
40 +func (b *streamBroker) Offer(session *reverseSession) error {
41 if session == nil {
42 return errors.New("reverse session is required")
43 }
@@ -63,7 +63,7 @@ func (b *leaseBroker) Offer(session *reverseSession) error {
63 return nil
64 }
65
66 -func (b *leaseBroker) Claim(ctx context.Context) (*reverseSession, error) {
66 +func (b *streamBroker) Claim(ctx context.Context) (*reverseSession, error) {
67 for {
68 b.mu.Lock()
69 if b.closedErr != nil {
@@ -96,7 +96,7 @@ func (b *leaseBroker) Claim(ctx context.Context) (*reverseSession, error) {
96 }
97 }
98
99 -func (b *leaseBroker) Close() {
99 +func (b *streamBroker) Close() {
100 b.mu.Lock()
101 sessions := b.ready
102 b.ready = nil
@@ -111,13 +111,13 @@ func (b *leaseBroker) Close() {
111 }
112 }
113
114 -func (b *leaseBroker) ReadyCount() int {
114 +func (b *streamBroker) ReadyCount() int {
115 b.mu.Lock()
116 defer b.mu.Unlock()
117 return len(b.ready)
118 }
119
120 -func (b *leaseBroker) watchSession(session *reverseSession) {
120 +func (b *streamBroker) watchSession(session *reverseSession) {
121 <-session.Done()
122
123 var readyCount int
@@ -140,7 +140,7 @@ func (b *leaseBroker) watchSession(session *reverseSession) {
140 b.mu.Unlock()
141 }
142
143 -func (b *leaseBroker) signalLocked() {
143 +func (b *streamBroker) signalLocked() {
144 select {
145 case b.notify <- struct{}{}:
146 default:
portal/broker_test.go
+5 -5
@@ -22,7 +22,7 @@ func TestLeaseBrokerClaimActivatesTLSMarker(t *testing.T) {
22 _ = clientConn.Close()
23 })
24
25 - broker := newLeaseBroker("lease-test", time.Hour, 2)
25 + broker := newStreamBroker("lease-test", time.Hour, 2)
26 session := newReverseSession(serverConn, time.Hour)
27 if err := broker.Offer(session); err != nil {
28 t.Fatalf("Offer() error = %v", err)
@@ -66,7 +66,7 @@ func TestLeaseBrokerCloseClosesIdleSessions(t *testing.T) {
66 t.Parallel()
67
68 serverConn, clientConn := net.Pipe()
69 - broker := newLeaseBroker("lease-test", time.Hour, 2)
69 + broker := newStreamBroker("lease-test", time.Hour, 2)
70 session := newReverseSession(serverConn, time.Hour)
71 if err := broker.Offer(session); err != nil {
72 t.Fatalf("Offer() error = %v", err)
@@ -84,7 +84,7 @@ func TestLeaseBrokerCloseClosesIdleSessions(t *testing.T) {
84 func TestLeaseBrokerCloseUnblocksClaim(t *testing.T) {
85 t.Parallel()
86
87 - broker := newLeaseBroker("lease-test", time.Hour, 2)
87 + broker := newStreamBroker("lease-test", time.Hour, 2)
88 claimCtx, cancel := context.WithTimeout(context.Background(), brokerAsyncTestTimeout)
89 defer cancel()
90
@@ -118,7 +118,7 @@ func TestLeaseBrokerCloseUnblocksClaim(t *testing.T) {
118 func TestLeaseBrokerClaimWaitsForLateOffer(t *testing.T) {
119 t.Parallel()
120
121 - broker := newLeaseBroker("lease-test", time.Hour, 2)
121 + broker := newStreamBroker("lease-test", time.Hour, 2)
122
123 serverConn, clientConn := net.Pipe()
124 t.Cleanup(func() {
@@ -190,7 +190,7 @@ func TestLeaseBrokerClaimWaitsForLateOffer(t *testing.T) {
190 func TestLeaseBrokerClaimTimesOutWithoutSessions(t *testing.T) {
191 t.Parallel()
192
193 - broker := newLeaseBroker("lease-test", time.Hour, 2)
193 + broker := newStreamBroker("lease-test", time.Hour, 2)
194
195 claimCtx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
196 defer cancel()
portal/datagram/flow.go new
+193
@@ -0,0 +1,193 @@
1 +package datagram
2 +
3 +import (
4 + "sync"
5 + "time"
6 +
7 + "github.com/quic-go/quic-go"
8 + "github.com/rs/zerolog/log"
9 +
10 + "github.com/gosuda/portal/v2/types"
11 +)
12 +
13 +const (
14 + defaultFlowIdleTimeout = 30 * time.Second
15 + defaultFlowCleanupInterval = 30 * time.Second
16 +)
17 +
18 +type flowReplyFunc func([]byte) error
19 +
20 +type flowState struct {
21 + key string
22 + lastSeen time.Time
23 + reply flowReplyFunc
24 +}
25 +
26 +// FlowMux manages a single QUIC connection from a tunnel for one lease.
27 +// All UDP traffic for the lease is multiplexed over DATAGRAM frames on this
28 +// connection, identified by flow IDs.
29 +type FlowMux struct {
30 + leaseID string
31 +
32 + session *Session
33 + flowTable map[uint32]*flowState // flowID -> client addr + liveness + reply path
34 + addrIndex map[string]uint32 // ingress key -> flowID
35 + nextFlow uint32
36 + mu sync.Mutex
37 +}
38 +
39 +func NewFlowMux(leaseID string) *FlowMux {
40 + mux := &FlowMux{
41 + leaseID: leaseID,
42 + session: NewSession(256, true, func(err error) {
43 + log.Warn().
44 + Err(err).
45 + Str("component", "quic-flow-mux").
46 + Str("lease_id", leaseID).
47 + Msg("quic receive loop ended")
48 + }),
49 + flowTable: make(map[uint32]*flowState),
50 + addrIndex: make(map[string]uint32),
51 + nextFlow: 1,
52 + }
53 + go mux.runDispatchLoop()
54 + go mux.runCleanupLoop()
55 + return mux
56 +}
57 +
58 +// Register stores the QUIC connection from the tunnel for this lease.
59 +// Replaces any existing connection.
60 +func (b *FlowMux) Register(conn *quic.Conn) error {
61 + if _, err := b.session.Bind(conn); err != nil {
62 + return err
63 + }
64 +
65 + log.Info().
66 + Str("component", "quic-flow-mux").
67 + Str("lease_id", b.leaseID).
68 + Str("remote_addr", conn.RemoteAddr().String()).
69 + Msg("quic tunnel connection registered")
70 + return nil
71 +}
72 +
73 +// HasConnection reports whether a tunnel QUIC connection is active.
74 +func (b *FlowMux) HasConnection() bool {
75 + return b.session.HasConnection()
76 +}
77 +
78 +// SendDatagram encodes a flow-framed datagram and sends it to the tunnel.
79 +func (b *FlowMux) SendDatagram(flowID uint32, payload []byte) error {
80 + return b.session.Send(flowID, payload)
81 +}
82 +
83 +// TouchFlow assigns a flow ID for an ingress key and updates its liveness and reply path.
84 +// If the key already has a flow, the existing ID is returned.
85 +func (b *FlowMux) TouchFlow(key string, reply flowReplyFunc) uint32 {
86 + now := time.Now()
87 +
88 + b.mu.Lock()
89 + defer b.mu.Unlock()
90 +
91 + if id, ok := b.addrIndex[key]; ok {
92 + if flow, exists := b.flowTable[id]; exists && flow != nil {
93 + flow.lastSeen = now
94 + if reply != nil {
95 + flow.reply = reply
96 + }
97 + return id
98 + }
99 + delete(b.addrIndex, key)
100 + }
101 +
102 + id := b.nextFlow
103 + b.nextFlow++
104 + b.flowTable[id] = &flowState{
105 + key: key,
106 + lastSeen: now,
107 + reply: reply,
108 + }
109 + b.addrIndex[key] = id
110 + return id
111 +}
112 +
113 +func (b *FlowMux) runDispatchLoop() {
114 + incoming := b.session.Incoming()
115 + for {
116 + select {
117 + case <-b.session.Done():
118 + return
119 + case frame := <-incoming:
120 + b.dispatch(frame)
121 + }
122 + }
123 +}
124 +
125 +func (b *FlowMux) dispatch(frame types.DatagramFrame) {
126 + b.mu.Lock()
127 + flow, ok := b.flowTable[frame.FlowID]
128 + if !ok || flow == nil || flow.reply == nil {
129 + b.mu.Unlock()
130 + return
131 + }
132 +
133 + flow.lastSeen = time.Now()
134 + reply := flow.reply
135 + b.mu.Unlock()
136 +
137 + if err := reply(frame.Payload); err != nil {
138 + log.Warn().
139 + Err(err).
140 + Str("component", "quic-flow-mux").
141 + Str("lease_id", b.leaseID).
142 + Uint32("flow_id", frame.FlowID).
143 + Msg("flow writeback failed")
144 + b.forgetFlow(frame.FlowID)
145 + }
146 +}
147 +
148 +func (b *FlowMux) runCleanupLoop() {
149 + ticker := time.NewTicker(defaultFlowCleanupInterval)
150 + defer ticker.Stop()
151 +
152 + for {
153 + select {
154 + case <-b.session.Done():
155 + return
156 + case now := <-ticker.C:
157 + b.expireIdleFlows(now)
158 + }
159 + }
160 +}
161 +
162 +func (b *FlowMux) expireIdleFlows(now time.Time) {
163 + b.mu.Lock()
164 + defer b.mu.Unlock()
165 +
166 + for flowID, flow := range b.flowTable {
167 + if flow == nil || now.Sub(flow.lastSeen) > defaultFlowIdleTimeout {
168 + if flow != nil {
169 + delete(b.addrIndex, flow.key)
170 + }
171 + delete(b.flowTable, flowID)
172 + }
173 + }
174 +}
175 +
176 +func (b *FlowMux) forgetFlow(flowID uint32) {
177 + b.mu.Lock()
178 + defer b.mu.Unlock()
179 +
180 + flow, ok := b.flowTable[flowID]
181 + if !ok {
182 + return
183 + }
184 + if flow != nil {
185 + delete(b.addrIndex, flow.key)
186 + }
187 + delete(b.flowTable, flowID)
188 +}
189 +
190 +// Stop tears down the QUIC connection and signals done.
191 +func (b *FlowMux) Stop() {
192 + b.session.Stop("lease stopped")
193 +}
portal/datagram/port_allocator.go new
+97
@@ -0,0 +1,97 @@
1 +package datagram
2 +
3 +import (
4 + "errors"
5 + "sort"
6 + "sync"
7 + "time"
8 +)
9 +
10 +var ErrPortExhausted = errors.New("no udp ports available")
11 +
12 +type portReservation struct {
13 + port int
14 + expiresAt time.Time
15 +}
16 +
17 +// PortAllocator manages a pool of UDP ports for dynamic per-lease allocation.
18 +type PortAllocator struct {
19 + available []int
20 + inUse map[int]string
21 + reserved map[string]portReservation
22 + grace time.Duration
23 + mu sync.Mutex
24 +}
25 +
26 +func NewPortAllocator(min, max int, grace time.Duration) *PortAllocator {
27 + available := make([]int, 0, max-min+1)
28 + for p := min; p <= max; p++ {
29 + available = append(available, p)
30 + }
31 + return &PortAllocator{
32 + available: available,
33 + inUse: make(map[int]string),
34 + reserved: make(map[string]portReservation),
35 + grace: grace,
36 + }
37 +}
38 +
39 +func (a *PortAllocator) Allocate(name string) (int, error) {
40 + a.mu.Lock()
41 + defer a.mu.Unlock()
42 +
43 + a.cleanupExpiredLocked(time.Now())
44 +
45 + if res, ok := a.reserved[name]; ok {
46 + delete(a.reserved, name)
47 + a.inUse[res.port] = name
48 + return res.port, nil
49 + }
50 +
51 + if len(a.available) == 0 {
52 + return 0, ErrPortExhausted
53 + }
54 +
55 + port := a.available[0]
56 + a.available = a.available[1:]
57 + a.inUse[port] = name
58 + return port, nil
59 +}
60 +
61 +func (a *PortAllocator) Release(port int) {
62 + a.mu.Lock()
63 + defer a.mu.Unlock()
64 +
65 + name, ok := a.inUse[port]
66 + if !ok {
67 + return
68 + }
69 + delete(a.inUse, port)
70 +
71 + if prev, exists := a.reserved[name]; exists {
72 + a.sortedInsertLocked(prev.port)
73 + }
74 +
75 + a.reserved[name] = portReservation{
76 + port: port,
77 + expiresAt: time.Now().Add(a.grace),
78 + }
79 +
80 + a.cleanupExpiredLocked(time.Now())
81 +}
82 +
83 +func (a *PortAllocator) cleanupExpiredLocked(now time.Time) {
84 + for name, res := range a.reserved {
85 + if now.After(res.expiresAt) {
86 + delete(a.reserved, name)
87 + a.sortedInsertLocked(res.port)
88 + }
89 + }
90 +}
91 +
92 +func (a *PortAllocator) sortedInsertLocked(port int) {
93 + i := sort.SearchInts(a.available, port)
94 + a.available = append(a.available, 0)
95 + copy(a.available[i+1:], a.available[i:])
96 + a.available[i] = port
97 +}
portal/datagram/relay.go new
+119
@@ -0,0 +1,119 @@
1 +package datagram
2 +
3 +import (
4 + "context"
5 + "errors"
6 + "fmt"
7 + "net"
8 + "sync"
9 + "time"
10 +
11 + "github.com/rs/zerolog/log"
12 +)
13 +
14 +const DefaultMaxPacketSize = 1350
15 +
16 +// Relay binds a UDP port for a lease and relays datagrams bidirectionally
17 +// between raw UDP clients and the tunnel's QUIC connection via the flow mux.
18 +type Relay struct {
19 + leaseID string
20 + port int
21 + flowMux *FlowMux
22 + conn *net.UDPConn
23 +
24 + cancel context.CancelFunc
25 + closeOnce sync.Once
26 +}
27 +
28 +func NewRelay(leaseID string, port int, flowMux *FlowMux) *Relay {
29 + return &Relay{
30 + leaseID: leaseID,
31 + port: port,
32 + flowMux: flowMux,
33 + }
34 +}
35 +
36 +func (r *Relay) Start(ctx context.Context) error {
37 + addr := &net.UDPAddr{Port: r.port}
38 + conn, err := net.ListenUDP("udp", addr)
39 + if err != nil {
40 + return fmt.Errorf("listen udp :%d: %w", r.port, err)
41 + }
42 + r.conn = conn
43 +
44 + relayCtx, cancel := context.WithCancel(ctx)
45 + r.cancel = cancel
46 +
47 + go r.readLoop(relayCtx)
48 +
49 + log.Info().
50 + Str("component", "udp-relay").
51 + Str("lease_id", r.leaseID).
52 + Int("port", r.port).
53 + Msg("udp relay started")
54 +
55 + return nil
56 +}
57 +
58 +func (r *Relay) Stop() {
59 + r.closeOnce.Do(func() {
60 + if r.cancel != nil {
61 + r.cancel()
62 + }
63 + if r.conn != nil {
64 + _ = r.conn.Close()
65 + }
66 + log.Info().
67 + Str("component", "udp-relay").
68 + Str("lease_id", r.leaseID).
69 + Int("port", r.port).
70 + Msg("udp relay stopped")
71 + })
72 +}
73 +
74 +func (r *Relay) readLoop(ctx context.Context) {
75 + buf := make([]byte, DefaultMaxPacketSize)
76 + for {
77 + select {
78 + case <-ctx.Done():
79 + return
80 + default:
81 + }
82 +
83 + _ = r.conn.SetReadDeadline(time.Now().Add(5 * time.Second))
84 + n, clientAddr, err := r.conn.ReadFromUDP(buf)
85 + if err != nil {
86 + if ctx.Err() != nil {
87 + return
88 + }
89 + var netErr net.Error
90 + if errors.As(err, &netErr) && netErr.Timeout() {
91 + continue
92 + }
93 + log.Warn().
94 + Str("component", "udp-relay").
95 + Str("lease_id", r.leaseID).
96 + Err(err).
97 + Msg("readLoop exiting: unexpected read error")
98 + return
99 + }
100 +
101 + flowID := r.flowMux.TouchFlow("udp:"+clientAddr.String(), func(payload []byte) error {
102 + _, err := r.conn.WriteToUDP(payload, clientAddr)
103 + return err
104 + })
105 + payload := make([]byte, n)
106 + copy(payload, buf[:n])
107 +
108 + if err := r.flowMux.SendDatagram(flowID, payload); err != nil {
109 + log.Warn().
110 + Str("component", "udp-relay").
111 + Str("lease_id", r.leaseID).
112 + Err(err).
113 + Uint32("flow_id", flowID).
114 + Int("bytes", n).
115 + Msg("send datagram to tunnel failed, dropping packet")
116 + continue
117 + }
118 + }
119 +}
portal/datagram/session.go new
+168
@@ -0,0 +1,168 @@
1 +package datagram
2 +
3 +import (
4 + "context"
5 + "errors"
6 + "sync"
7 +
8 + "github.com/quic-go/quic-go"
9 +
10 + "github.com/gosuda/portal/v2/types"
11 +)
12 +
13 +var (
14 + ErrNoConnection = errors.New("no quic connection registered")
15 + ErrSessionClosed = errors.New("quic datagram session closed")
16 +)
17 +
18 +// Session owns one active QUIC DATAGRAM connection and exposes decoded frames.
19 +type Session struct {
20 + incoming chan types.DatagramFrame
21 + dropIncoming bool
22 + onReceiveError func(error)
23 + done chan struct{}
24 +
25 + mu sync.Mutex
26 + conn *quic.Conn
27 + closed bool
28 +}
29 +
30 +func NewSession(bufferSize int, dropIncoming bool, onReceiveError func(error)) *Session {
31 + if bufferSize <= 0 {
32 + bufferSize = 256
33 + }
34 +
35 + return &Session{
36 + incoming: make(chan types.DatagramFrame, bufferSize),
37 + dropIncoming: dropIncoming,
38 + onReceiveError: onReceiveError,
39 + done: make(chan struct{}),
40 + }
41 +}
42 +
43 +// Bind installs a new active QUIC connection and starts the receive loop.
44 +// Any previously active connection is replaced and closed.
45 +func (s *Session) Bind(conn *quic.Conn) (<-chan struct{}, error) {
46 + if conn == nil {
47 + return nil, errors.New("quic connection is required")
48 + }
49 +
50 + s.mu.Lock()
51 + if s.closed {
52 + s.mu.Unlock()
53 + _ = conn.CloseWithError(0, "session closed")
54 + return nil, ErrSessionClosed
55 + }
56 + old := s.conn
57 + s.conn = conn
58 + s.mu.Unlock()
59 +
60 + if old != nil {
61 + _ = old.CloseWithError(0, "replaced")
62 + }
63 +
64 + recvDone := make(chan struct{})
65 + go s.receiveLoop(conn, recvDone)
66 + return recvDone, nil
67 +}
68 +
69 +func (s *Session) Incoming() <-chan types.DatagramFrame {
70 + return s.incoming
71 +}
72 +
73 +func (s *Session) Done() <-chan struct{} {
74 + return s.done
75 +}
76 +
77 +func (s *Session) HasConnection() bool {
78 + s.mu.Lock()
79 + defer s.mu.Unlock()
80 + return s.conn != nil && !s.closed
81 +}
82 +
83 +func (s *Session) Send(flowID uint32, payload []byte) error {
84 + s.mu.Lock()
85 + conn := s.conn
86 + closed := s.closed
87 + s.mu.Unlock()
88 +
89 + if closed {
90 + return ErrSessionClosed
91 + }
92 + if conn == nil {
93 + return ErrNoConnection
94 + }
95 + return conn.SendDatagram(types.EncodeDatagram(flowID, payload))
96 +}
97 +
98 +// Clear closes the active connection but keeps the session reusable.
99 +func (s *Session) Clear(reason string) {
100 + s.mu.Lock()
101 + conn := s.conn
102 + s.conn = nil
103 + s.mu.Unlock()
104 +
105 + if conn != nil {
106 + _ = conn.CloseWithError(0, reason)
107 + }
108 +}
109 +
110 +// Stop permanently closes the session and any active connection.
111 +func (s *Session) Stop(reason string) {
112 + s.mu.Lock()
113 + if s.closed {
114 + s.mu.Unlock()
115 + return
116 + }
117 + s.closed = true
118 + conn := s.conn
119 + s.conn = nil
120 + close(s.done)
121 + s.mu.Unlock()
122 +
123 + if conn != nil {
124 + _ = conn.CloseWithError(0, reason)
125 + }
126 +}
127 +
128 +func (s *Session) receiveLoop(conn *quic.Conn, recvDone chan struct{}) {
129 + defer close(recvDone)
130 +
131 + for {
132 + data, err := conn.ReceiveDatagram(context.Background())
133 + if err != nil {
134 + s.mu.Lock()
135 + isActive := s.conn == conn
136 + if isActive {
137 + s.conn = nil
138 + }
139 + closed := s.closed
140 + onReceiveError := s.onReceiveError
141 + s.mu.Unlock()
142 +
143 + if isActive && !closed && onReceiveError != nil {
144 + onReceiveError(err)
145 + }
146 + return
147 + }
148 +
149 + frame, err := types.DecodeDatagram(data)
150 + if err != nil {
151 + continue
152 + }
153 +
154 + if s.dropIncoming {
155 + select {
156 + case s.incoming <- frame:
157 + default:
158 + }
159 + continue
160 + }
161 +
162 + select {
163 + case s.incoming <- frame:
164 + case <-s.done:
165 + return
166 + }
167 + }
168 +}
portal/datagram/sni_parse.go renamed
+26 -208
@@ -1,4 +1,4 @@
1 -package portal
1 +package datagram
2
3 import (
4 "crypto"
@@ -7,15 +7,10 @@ import (
7 "encoding/binary"
8 "errors"
9 "fmt"
10 - "net"
11 - "sync"
10
11 "golang.org/x/crypto/hkdf"
14 -
15 - "github.com/gosuda/portal/v2/utils"
12 )
13
18 -// QUIC v1 constants (RFC 9001).
14 var quicV1InitialSalt = []byte{
15 0x38, 0x76, 0x2c, 0xf7, 0xf5, 0x59, 0x34, 0xb3,
16 0x4d, 0x17, 0x9a, 0xe6, 0xa4, 0xc8, 0x0c, 0xad,
@@ -25,36 +20,27 @@ var quicV1InitialSalt = []byte{
20 var errNotQUICInitial = errors.New("not a quic initial packet")
21 var errSNINotFound = errors.New("sni not found in quic initial")
22
28 -// parseQUICInitialSNI extracts the TLS SNI from a QUIC Initial packet.
29 -// It decrypts the Initial packet header and payload using keys derived from
30 -// the Destination Connection ID per RFC 9001 Section 5.2, then parses the
31 -// CRYPTO frame to find the TLS ClientHello SNI extension.
32 -func parseQUICInitialSNI(packet []byte) (string, error) {
23 +func ParseQUICInitialSNI(packet []byte) (string, error) {
24 if len(packet) < 5 {
25 return "", errNotQUICInitial
26 }
27
37 - // Long header: first bit is 1, second bit is 1 (fixed), bits 4-5 are packet type.
28 firstByte := packet[0]
29 if firstByte&0x80 == 0 {
40 - return "", errNotQUICInitial // short header
30 + return "", errNotQUICInitial
31 }
32
43 - // Packet type: bits 4-5 of first byte. Initial = 0.
33 packetType := (firstByte & 0x30) >> 4
34 if packetType != 0 {
35 return "", errNotQUICInitial
36 }
37
49 - // Version (4 bytes).
38 version := binary.BigEndian.Uint32(packet[1:5])
39 if version == 0 {
52 - return "", errNotQUICInitial // version negotiation
40 + return "", errNotQUICInitial
41 }
42
43 offset := 5
56 -
57 - // Destination Connection ID length + DCID.
44 if offset >= len(packet) {
45 return "", errNotQUICInitial
46 }
@@ -66,7 +52,6 @@ func parseQUICInitialSNI(packet []byte) (string, error) {
52 dcid := packet[offset : offset+dcidLen]
53 offset += dcidLen
54
69 - // Source Connection ID length + SCID.
55 if offset >= len(packet) {
56 return "", errNotQUICInitial
57 }
@@ -77,7 +62,6 @@ func parseQUICInitialSNI(packet []byte) (string, error) {
62 return "", errNotQUICInitial
63 }
64
80 - // Token length (varint) + token.
65 tokenLen, n := readVarint(packet[offset:])
66 if n <= 0 {
67 return "", errNotQUICInitial
@@ -87,7 +71,6 @@ func parseQUICInitialSNI(packet []byte) (string, error) {
71 return "", errNotQUICInitial
72 }
73
90 - // Payload length (varint).
74 payloadLen, n := readVarint(packet[offset:])
75 if n <= 0 {
76 return "", errNotQUICInitial
@@ -95,8 +78,6 @@ func parseQUICInitialSNI(packet []byte) (string, error) {
78 offset += n
79 _ = payloadLen
80
98 - // The rest from offset is: packet number (1-4 bytes, encrypted) + encrypted payload.
99 - // We need to decrypt the header first to determine packet number length.
81 clientSecret, err := deriveInitialClientSecret(dcid, version)
82 if err != nil {
83 return "", fmt.Errorf("derive initial secret: %w", err)
@@ -117,7 +98,6 @@ func parseQUICInitialSNI(packet []byte) (string, error) {
98 return "", fmt.Errorf("derive iv: %w", err)
99 }
100
120 - // Header protection: sample 16 bytes starting 4 bytes after packet number offset.
101 pnOffset := offset
102 sampleOffset := pnOffset + 4
103 if sampleOffset+16 > len(packet) {
@@ -125,7 +105,6 @@ func parseQUICInitialSNI(packet []byte) (string, error) {
105 }
106 sample := packet[sampleOffset : sampleOffset+16]
107
128 - // Create AES-ECB cipher for HP mask.
108 block, err := aes.NewCipher(hp)
109 if err != nil {
110 return "", fmt.Errorf("aes cipher: %w", err)
@@ -133,11 +112,9 @@ func parseQUICInitialSNI(packet []byte) (string, error) {
112 mask := make([]byte, aes.BlockSize)
113 block.Encrypt(mask, sample)
114
136 - // Unmask first byte.
137 - unmaskedFirst := packet[0] ^ (mask[0] & 0x0f) // long header: lower 4 bits
115 + unmaskedFirst := packet[0] ^ (mask[0] & 0x0f)
116 pnLength := int(unmaskedFirst&0x03) + 1
117
140 - // Unmask packet number.
118 pnBytes := make([]byte, pnLength)
119 for i := range pnLength {
120 pnBytes[i] = packet[pnOffset+i] ^ mask[1+i]
@@ -148,85 +125,57 @@ func parseQUICInitialSNI(packet []byte) (string, error) {
125 pn = (pn << 8) | uint32(b)
126 }
127
151 - // Build nonce for AEAD.
128 nonce := make([]byte, len(iv))
129 copy(nonce, iv)
154 - for i := range len(nonce) {
130 + for i := range nonce {
131 if i >= len(nonce)-4 {
132 nonce[i] ^= byte(pn >> (8 * (len(nonce) - 1 - i)))
133 }
134 }
135
160 - // Decrypt payload.
136 payloadOffset := pnOffset + pnLength
137 if payloadOffset >= len(packet) {
138 return "", errNotQUICInitial
139 }
140
166 - // AAD = entire header with unmasked first byte and unmasked PN.
141 aad := make([]byte, payloadOffset)
142 copy(aad, packet[:payloadOffset])
143 aad[0] = unmaskedFirst
144 copy(aad[pnOffset:], pnBytes)
145
172 - aead, err := cipher.NewGCM(block)
146 + aeadBlock, err := aes.NewCipher(key)
147 if err != nil {
174 - // Use the key for AEAD, not HP block.
175 - aeadBlock, err2 := aes.NewCipher(key)
176 - if err2 != nil {
177 - return "", fmt.Errorf("aead cipher: %w", err2)
178 - }
179 - aead, err = cipher.NewGCM(aeadBlock)
180 - if err != nil {
181 - return "", fmt.Errorf("gcm: %w", err)
182 - }
183 - } else {
184 - // We used the HP block for AEAD by mistake. Redo with key.
185 - aeadBlock, err2 := aes.NewCipher(key)
186 - if err2 != nil {
187 - return "", fmt.Errorf("aead cipher: %w", err2)
188 - }
189 - aead, err = cipher.NewGCM(aeadBlock)
190 - if err != nil {
191 - return "", fmt.Errorf("gcm: %w", err)
192 - }
148 + return "", fmt.Errorf("aead cipher: %w", err)
149 + }
150 + aead, err := cipher.NewGCM(aeadBlock)
151 + if err != nil {
152 + return "", fmt.Errorf("gcm: %w", err)
153 }
154
195 - ciphertext := packet[payloadOffset:]
196 - plaintext, err := aead.Open(nil, nonce, ciphertext, aad)
155 + plaintext, err := aead.Open(nil, nonce, packet[payloadOffset:], aad)
156 if err != nil {
157 return "", fmt.Errorf("decrypt initial payload: %w", err)
158 }
159
201 - // Parse CRYPTO frames to find ClientHello.
160 return extractSNIFromCryptoFrames(plaintext)
161 }
162
205 -// extractSNIFromCryptoFrames parses QUIC frames looking for CRYPTO frames
206 -// containing a TLS ClientHello, and extracts the SNI server_name extension.
163 func extractSNIFromCryptoFrames(frames []byte) (string, error) {
164 offset := 0
165 for offset < len(frames) {
166 frameType := frames[offset]
167 offset++
168
213 - switch {
214 - case frameType == 0x00:
215 - // PADDING frame — skip.
169 + switch frameType {
170 + case 0x00, 0x01:
171 continue
217 - case frameType == 0x01:
218 - // PING frame — skip.
219 - continue
220 - case frameType == 0x06:
221 - // CRYPTO frame.
222 - // Offset field (varint).
172 + case 0x06:
173 _, n := readVarint(frames[offset:])
174 if n <= 0 {
175 return "", errSNINotFound
176 }
177 offset += n
178
229 - // Length field (varint).
179 dataLen, n := readVarint(frames[offset:])
180 if n <= 0 {
181 return "", errSNINotFound
@@ -244,20 +193,14 @@ func extractSNIFromCryptoFrames(frames []byte) (string, error) {
193 return sni, nil
194 }
195 default:
247 - // Unknown frame — can't continue parsing reliably.
196 return "", errSNINotFound
197 }
198 }
199 return "", errSNINotFound
200 }
201
254 -// parseTLSClientHelloSNI parses a raw TLS ClientHello message and extracts SNI.
202 func parseTLSClientHelloSNI(data []byte) (string, error) {
256 - // TLS handshake: type(1) + length(3) + ...
257 - if len(data) < 4 {
258 - return "", errSNINotFound
259 - }
260 - if data[0] != 0x01 { // ClientHello
203 + if len(data) < 4 || data[0] != 0x01 {
204 return "", errSNINotFound
205 }
206 msgLen := int(data[1])<<16 | int(data[2])<<8 | int(data[3])
@@ -266,51 +209,46 @@ func parseTLSClientHelloSNI(data []byte) (string, error) {
209 }
210 body := data[4 : 4+msgLen]
211
269 - // ClientHello: version(2) + random(32) + session_id_len(1) + session_id + ...
212 if len(body) < 34 {
213 return "", errSNINotFound
214 }
273 - offset := 2 + 32 // skip version + random
215 + offset := 34
216
275 - // Session ID.
217 if offset >= len(body) {
218 return "", errSNINotFound
219 }
220 sessionIDLen := int(body[offset])
221 offset += 1 + sessionIDLen
222
282 - // Cipher suites.
223 if offset+2 > len(body) {
224 return "", errSNINotFound
225 }
226 cipherSuitesLen := int(body[offset])<<8 | int(body[offset+1])
227 offset += 2 + cipherSuitesLen
228
289 - // Compression methods.
229 if offset >= len(body) {
230 return "", errSNINotFound
231 }
232 compMethodsLen := int(body[offset])
233 offset += 1 + compMethodsLen
234
296 - // Extensions.
235 if offset+2 > len(body) {
236 return "", errSNINotFound
237 }
238 extensionsLen := int(body[offset])<<8 | int(body[offset+1])
239 offset += 2
240
303 - extEnd := offset + extensionsLen
304 - if extEnd > len(body) {
305 - extEnd = len(body)
306 - }
241 + extEnd := min(offset+extensionsLen, len(body))
242
243 for offset+4 <= extEnd {
244 extType := int(body[offset])<<8 | int(body[offset+1])
245 extLen := int(body[offset+2])<<8 | int(body[offset+3])
246 offset += 4
247
313 - if extType == 0x0000 { // server_name
248 + if offset+extLen > extEnd {
249 + return "", errSNINotFound
250 + }
251 + if extType == 0x0000 {
252 return parseSNIExtension(body[offset : offset+extLen])
253 }
254 offset += extLen
@@ -323,19 +261,15 @@ func parseSNIExtension(data []byte) (string, error) {
261 if len(data) < 2 {
262 return "", errSNINotFound
263 }
326 - // Server name list length.
264 listLen := int(data[0])<<8 | int(data[1])
265 offset := 2
329 - end := offset + listLen
330 - if end > len(data) {
331 - end = len(data)
332 - }
266 + end := min(offset+listLen, len(data))
267
268 for offset+3 <= end {
269 nameType := data[offset]
270 nameLen := int(data[offset+1])<<8 | int(data[offset+2])
271 offset += 3
338 - if nameType == 0x00 { // host_name
272 + if nameType == 0x00 {
273 if offset+nameLen > end {
274 return "", errSNINotFound
275 }
@@ -346,7 +280,6 @@ func parseSNIExtension(data []byte) (string, error) {
280 return "", errSNINotFound
281 }
282
349 -// QUIC Initial secret derivation (RFC 9001 Section 5.2).
283 func deriveInitialClientSecret(dcid []byte, version uint32) ([]byte, error) {
284 salt := quicV1InitialSalt
285
@@ -387,7 +320,6 @@ func deriveIV(secret []byte) ([]byte, error) {
320 return iv, nil
321 }
322
390 -// hkdfLabel builds a TLS 1.3 HkdfLabel structure for HKDF-Expand-Label.
323 func hkdfLabel(label []byte, length int) []byte {
324 fullLabel := append([]byte("tls13 "), label...)
325 out := make([]byte, 2+1+len(fullLabel)+1)
@@ -395,11 +327,10 @@ func hkdfLabel(label []byte, length int) []byte {
327 out[1] = byte(length)
328 out[2] = byte(len(fullLabel))
329 copy(out[3:], fullLabel)
398 - out[3+len(fullLabel)] = 0 // empty context
330 + out[3+len(fullLabel)] = 0
331 return out
332 }
333
402 -// readVarint reads a QUIC variable-length integer (RFC 9000 Section 16).
334 func readVarint(data []byte) (uint64, int) {
335 if len(data) == 0 {
336 return 0, -1
@@ -417,116 +348,3 @@ func readVarint(data []byte) (uint64, int) {
348 }
349 return val, length
350 }
420 -
421 -// quicSNIRouter listens on a raw UDP socket and routes QUIC connections
422 -// based on SNI extracted from Initial packets.
423 -type quicSNIRouter struct {
424 - conn net.PacketConn
425 - server *Server
426 - connTable map[string]string // "src_ip:port" → leaseID
427 - mu sync.RWMutex
428 - done chan struct{}
429 - closeOnce sync.Once
430 -}
431 -
432 -func newQUICSNIRouter(conn net.PacketConn, server *Server) *quicSNIRouter {
433 - return &quicSNIRouter{
434 - conn: conn,
435 - server: server,
436 - connTable: make(map[string]string),
437 - done: make(chan struct{}),
438 - }
439 -}
440 -
441 -func (r *quicSNIRouter) run() error {
442 - buf := make([]byte, 65535)
443 - for {
444 - n, addr, err := r.conn.ReadFrom(buf)
445 - if err != nil {
446 - select {
447 - case <-r.done:
448 - return nil
449 - default:
450 - }
451 - if errors.Is(err, net.ErrClosed) {
452 - return nil
453 - }
454 - return err
455 - }
456 -
457 - packet := make([]byte, n)
458 - copy(packet, buf[:n])
459 - go r.handlePacket(packet, addr)
460 - }
461 -}
462 -
463 -func (r *quicSNIRouter) handlePacket(packet []byte, srcAddr net.Addr) {
464 - key := srcAddr.String()
465 -
466 - // Check if we already have a mapping for this source.
467 - r.mu.RLock()
468 - leaseID, found := r.connTable[key]
469 - r.mu.RUnlock()
470 -
471 - if !found {
472 - // Try to parse SNI from Initial packet.
473 - sni, err := parseQUICInitialSNI(packet)
474 - if err != nil || sni == "" {
475 - return // drop non-Initial or unparseable packets from unknown sources
476 - }
477 -
478 - serverName := utils.NormalizeHostname(sni)
479 - record, ok := r.server.registry.Lookup(serverName)
480 - if !ok || record == nil {
481 - return // no route
482 - }
483 - leaseID = record.ID
484 -
485 - r.mu.Lock()
486 - r.connTable[key] = leaseID
487 - r.mu.Unlock()
488 - }
489 -
490 - // Forward packet to the tunnel via QUIC DATAGRAM.
491 - record, ok := r.server.registry.Get(leaseID)
492 -
493 - if !ok || record == nil || record.QUICBroker == nil || !record.QUICBroker.HasConnection() {
494 - return
495 - }
496 -
497 - udpAddr, ok := srcAddr.(*net.UDPAddr)
498 - if !ok {
499 - return
500 - }
501 -
502 - flowID := record.QUICBroker.AllocateFlow(udpAddr)
503 - _ = record.QUICBroker.SendDatagram(flowID, packet)
504 -}
505 -
506 -// writeBackLoop reads datagrams from each QUIC broker and writes raw UDP back
507 -// to the public QUIC clients via the SNI router's PacketConn.
508 -func (r *quicSNIRouter) writeBackLoop(leaseID string, broker *quicBroker) {
509 - for {
510 - select {
511 - case <-r.done:
512 - return
513 - case <-broker.Done():
514 - return
515 - case frame := <-broker.Incoming():
516 - addr, ok := broker.LookupFlowAddr(frame.FlowID)
517 - if !ok {
518 - continue
519 - }
520 - _, _ = r.conn.WriteTo(frame.Payload, addr)
521 - }
522 - }
523 -}
524 -
525 -func (r *quicSNIRouter) close() error {
526 - var closeErr error
527 - r.closeOnce.Do(func() {
528 - close(r.done)
529 - closeErr = r.conn.Close()
530 - })
531 - return closeErr
532 -}
portal/lease.go
+164 -4
@@ -8,6 +8,7 @@ import (
8 "sync"
9 "time"
10
11 + portaldatagram "github.com/gosuda/portal/v2/portal/datagram"
12 "github.com/gosuda/portal/v2/portal/policy"
13 "github.com/gosuda/portal/v2/types"
14 "github.com/gosuda/portal/v2/utils"
@@ -220,7 +221,8 @@ func (r *leaseRegistry) Snapshot(record *leaseRecord) types.Lease {
221 snapshot.Metadata = snapshot.Metadata.Copy()
222 clientIP := record.ClientIP
223 snapshot.BPS = r.policy.BPSManager().LeaseBPS(record.ID)
223 - snapshot.Ready = record.Broker.ReadyCount()
224 + snapshot.Ready = record.ReadyCount()
225 + snapshot.UDPPort = record.UDPPort()
226 snapshot.IsApproved = r.policy.EffectiveApproval(record.ID)
227 snapshot.IsBanned = r.policy.IsLeaseBanned(record.ID)
228 snapshot.IsDenied = r.policy.IsLeaseDenied(record.ID)
@@ -230,10 +232,168 @@ func (r *leaseRegistry) Snapshot(record *leaseRecord) types.Lease {
232
233 type leaseRecord struct {
234 types.Lease
233 - Broker *leaseBroker
234 - QUICBroker *quicBroker
235 - UDPRelay *udpRelay
235 ReverseToken string
236 + Runtime *leaseRuntime
237 +}
238 +
239 +func (r *leaseRecord) SupportsDatagram() bool {
240 + return r != nil && r.Runtime != nil && r.Runtime.SupportsDatagram()
241 +}
242 +
243 +func (r *leaseRecord) SupportsStream() bool {
244 + return r != nil && r.Runtime != nil && r.Runtime.SupportsStream()
245 +}
246 +
247 +func (r *leaseRecord) ReadyCount() int {
248 + if r == nil || r.Runtime == nil {
249 + return 0
250 + }
251 + return r.Runtime.ReadyCount()
252 +}
253 +
254 +func (r *leaseRecord) StreamBroker() *streamBroker {
255 + if r == nil || r.Runtime == nil {
256 + return nil
257 + }
258 + return r.Runtime.StreamBroker()
259 +}
260 +
261 +func (r *leaseRecord) DatagramFlowMux() *portaldatagram.FlowMux {
262 + if r == nil || r.Runtime == nil {
263 + return nil
264 + }
265 + return r.Runtime.DatagramFlowMux()
266 +}
267 +
268 +func (r *leaseRecord) UDPPort() int {
269 + if r == nil || r.Runtime == nil {
270 + return 0
271 + }
272 + return r.Runtime.UDPPort()
273 +}
274 +
275 +type leaseRuntime struct {
276 + capabilities types.LeaseCapabilities
277 + datagram *leaseDatagramRuntime
278 + startErr error
279 + startOnce sync.Once
280 + stream *leaseStreamRuntime
281 +}
282 +
283 +type leaseRuntimeConfig struct {
284 + Capabilities types.LeaseCapabilities
285 + IdleInterval time.Duration
286 + LeaseID string
287 + LeaseName string
288 + PortAllocator *portaldatagram.PortAllocator
289 + ReadyLimit int
290 +}
291 +
292 +type leaseStreamRuntime struct {
293 + broker *streamBroker
294 +}
295 +
296 +type leaseDatagramRuntime struct {
297 + flowMux *portaldatagram.FlowMux
298 + port int
299 + relay *portaldatagram.Relay
300 +}
301 +
302 +func newLeaseRuntime(cfg leaseRuntimeConfig) (*leaseRuntime, error) {
303 + runtime := &leaseRuntime{capabilities: cfg.Capabilities}
304 +
305 + if cfg.Capabilities.SupportsStream() {
306 + runtime.stream = &leaseStreamRuntime{
307 + broker: newStreamBroker(cfg.LeaseID, cfg.IdleInterval, cfg.ReadyLimit),
308 + }
309 + }
310 +
311 + if cfg.Capabilities.SupportsDatagram() {
312 + if cfg.PortAllocator == nil {
313 + return nil, errors.New("udp port allocation not available")
314 + }
315 + port, err := cfg.PortAllocator.Allocate(cfg.LeaseName)
316 + if err != nil {
317 + return nil, fmt.Errorf("allocate udp port: %w", err)
318 + }
319 +
320 + flowMux := portaldatagram.NewFlowMux(cfg.LeaseID)
321 + runtime.datagram = &leaseDatagramRuntime{
322 + flowMux: flowMux,
323 + port: port,
324 + relay: portaldatagram.NewRelay(cfg.LeaseID, port, flowMux),
325 + }
326 + }
327 +
328 + return runtime, nil
329 +}
330 +
331 +func (r *leaseRuntime) Start() error {
332 + if r == nil || r.datagram == nil || r.datagram.relay == nil {
333 + return nil
334 + }
335 +
336 + r.startOnce.Do(func() {
337 + r.startErr = r.datagram.relay.Start(context.Background())
338 + })
339 +
340 + return r.startErr
341 +}
342 +
343 +func (r *leaseRuntime) Close(ports *portaldatagram.PortAllocator) {
344 + if r == nil {
345 + return
346 + }
347 + if r.stream != nil && r.stream.broker != nil {
348 + r.stream.broker.Close()
349 + }
350 + if r.datagram != nil {
351 + if r.datagram.flowMux != nil {
352 + r.datagram.flowMux.Stop()
353 + }
354 + if r.datagram.relay != nil {
355 + r.datagram.relay.Stop()
356 + }
357 + if r.datagram.port > 0 && ports != nil {
358 + ports.Release(r.datagram.port)
359 + }
360 + }
361 +}
362 +
363 +func (r *leaseRuntime) SupportsDatagram() bool {
364 + return r != nil && r.capabilities.SupportsDatagram()
365 +}
366 +
367 +func (r *leaseRuntime) SupportsStream() bool {
368 + return r != nil && r.capabilities.SupportsStream()
369 +}
370 +
371 +func (r *leaseRuntime) ReadyCount() int {
372 + if r == nil || r.stream == nil || r.stream.broker == nil {
373 + return 0
374 + }
375 + return r.stream.broker.ReadyCount()
376 +}
377 +
378 +func (r *leaseRuntime) StreamBroker() *streamBroker {
379 + if r == nil || r.stream == nil {
380 + return nil
381 + }
382 + return r.stream.broker
383 +}
384 +
385 +func (r *leaseRuntime) DatagramFlowMux() *portaldatagram.FlowMux {
386 + if r == nil || r.datagram == nil {
387 + return nil
388 + }
389 + return r.datagram.flowMux
390 +}
391 +
392 +func (r *leaseRuntime) UDPPort() int {
393 + if r == nil || r.datagram == nil {
394 + return 0
395 + }
396 + return r.datagram.port
397 }
398
399 type routeTable struct {
portal/lease_test.go
+17 -5
@@ -10,6 +10,15 @@ import (
10 "github.com/gosuda/portal/v2/types"
11 )
12
13 +func newTestStreamLeaseRuntime(leaseID string) *leaseRuntime {
14 + return &leaseRuntime{
15 + capabilities: types.LeaseCapabilities{Stream: true},
16 + stream: &leaseStreamRuntime{
17 + broker: newStreamBroker(leaseID, time.Minute, 1),
18 + },
19 + }
20 +}
21 +
22 func TestLeaseRegistryLifecycle(t *testing.T) {
23 t.Parallel()
24
@@ -22,6 +31,7 @@ func TestLeaseRegistryLifecycle(t *testing.T) {
31 ExpiresAt: time.Now().Add(30 * time.Second),
32 },
33 ReverseToken: "tok_1",
34 + Runtime: newTestStreamLeaseRuntime("lease_1"),
35 }
36
37 if err := registry.Register(record); err != nil {
@@ -71,6 +81,7 @@ func TestLeaseRegistryWildcardAndConflict(t *testing.T) {
81 ExpiresAt: time.Now().Add(30 * time.Second),
82 },
83 ReverseToken: "tok_wildcard",
84 + Runtime: newTestStreamLeaseRuntime("lease_wildcard"),
85 }
86 if err := registry.Register(wildcardLease); err != nil {
87 t.Fatalf("Register(wildcard) error = %v", err)
@@ -90,6 +101,7 @@ func TestLeaseRegistryWildcardAndConflict(t *testing.T) {
101 ExpiresAt: time.Now().Add(30 * time.Second),
102 },
103 ReverseToken: "tok_conflict",
104 + Runtime: newTestStreamLeaseRuntime("lease_conflict"),
105 }
106 err := registry.Register(conflict)
107 if !errors.Is(err, errHostnameConflict) {
@@ -115,7 +127,7 @@ func TestLeaseRegistrySnapshotAndRoutableUsePolicy(t *testing.T) {
127 ClientIP: "203.0.113.20",
128 },
129 ReverseToken: "tok_policy",
118 - Broker: newLeaseBroker("lease_policy", time.Minute, 1),
130 + Runtime: newTestStreamLeaseRuntime("lease_policy"),
131 }
132 if err := registry.Register(record); err != nil {
133 t.Fatalf("Register() error = %v", err)
@@ -149,7 +161,7 @@ func TestLeaseRegistryCleanupExpiredClosesBroker(t *testing.T) {
161
162 registry := newLeaseRegistry(policy.NewRuntime())
163 registry.onExpired = func(r *leaseRecord) {
152 - r.Broker.Close()
164 + r.Runtime.Close(nil)
165 }
166 record := &leaseRecord{
167 Lease: types.Lease{
@@ -158,20 +170,20 @@ func TestLeaseRegistryCleanupExpiredClosesBroker(t *testing.T) {
170 ExpiresAt: time.Now().Add(-time.Second),
171 },
172 ReverseToken: "tok_expired",
161 - Broker: newLeaseBroker("lease_expired", time.Minute, 1),
173 + Runtime: newTestStreamLeaseRuntime("lease_expired"),
174 }
175 if err := registry.Register(record); err != nil {
176 t.Fatalf("Register() error = %v", err)
177 }
178
179 for _, lease := range registry.removeExpired(time.Now()) {
168 - lease.Broker.Close()
180 + lease.Runtime.Close(nil)
181 }
182
183 if _, ok := registry.Lookup("expired.example.com"); ok {
184 t.Fatal("Lookup() after removeExpired() = true, want false")
185 }
174 - if _, err := record.Broker.Claim(context.Background()); !errors.Is(err, errBrokerClosed) {
186 + if _, err := record.StreamBroker().Claim(context.Background()); !errors.Is(err, errBrokerClosed) {
187 t.Fatalf("Claim() after removeExpired() error = %v, want %v", err, errBrokerClosed)
188 }
189 }
portal/quic.go deleted
-296
@@ -1,296 +0,0 @@
1 -package portal
2 -
3 -import (
4 - "context"
5 - "encoding/json"
6 - "errors"
7 - "net"
8 - "sync"
9 - "time"
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")
20 -)
21 -
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.
25 -type quicBroker struct {
26 - leaseID string
27 -
28 - conn *quic.Conn
29 - flowTable map[uint32]*net.UDPAddr // flowID → client addr
30 - addrIndex map[string]uint32 // "ip:port" → flowID
31 - nextFlow uint32
32 -
33 - incoming chan types.DatagramFrame // frames received from tunnel
34 - done chan struct{}
35 -
36 - mu sync.Mutex
37 - closeOnce sync.Once
38 - closed bool
39 -}
40 -
41 -func newQUICBroker(leaseID string) *quicBroker {
42 - return &quicBroker{
43 - leaseID: leaseID,
44 - flowTable: make(map[uint32]*net.UDPAddr),
45 - addrIndex: make(map[string]uint32),
46 - nextFlow: 1,
47 - incoming: make(chan types.DatagramFrame, 256),
48 - done: make(chan struct{}),
49 - }
50 -}
51 -
52 -// Register stores the QUIC connection from the tunnel for this lease.
53 -// Replaces any existing connection.
54 -func (b *quicBroker) Register(conn *quic.Conn) {
55 - b.mu.Lock()
56 - old := b.conn
57 - b.conn = conn
58 - b.mu.Unlock()
59 -
60 - if old != nil {
61 - _ = old.CloseWithError(0, "replaced")
62 - }
63 -
64 - go b.receiveLoop(conn)
65 -
66 - log.Info().
67 - Str("component", "quic-broker").
68 - Str("lease_id", b.leaseID).
69 - Str("remote_addr", conn.RemoteAddr().String()).
70 - Msg("quic tunnel connection registered")
71 -}
72 -
73 -// HasConnection reports whether a tunnel QUIC connection is active.
74 -func (b *quicBroker) HasConnection() bool {
75 - b.mu.Lock()
76 - defer b.mu.Unlock()
77 - return b.conn != nil && !b.closed
78 -}
79 -
80 -// SendDatagram encodes a flow-framed datagram and sends it to the tunnel.
81 -func (b *quicBroker) SendDatagram(flowID uint32, payload []byte) error {
82 - b.mu.Lock()
83 - conn := b.conn
84 - b.mu.Unlock()
85 -
86 - if conn == nil {
87 - return errQUICNoConnection
88 - }
89 - return conn.SendDatagram(types.EncodeDatagram(flowID, payload))
90 -}
91 -
92 -// Incoming returns the channel that delivers datagrams received from the tunnel.
93 -func (b *quicBroker) Incoming() <-chan types.DatagramFrame {
94 - return b.incoming
95 -}
96 -
97 -// Done returns a channel closed when the broker shuts down.
98 -func (b *quicBroker) Done() <-chan struct{} {
99 - return b.done
100 -}
101 -
102 -// AllocateFlow assigns a flow ID for a client address. If the address already
103 -// has a flow, the existing ID is returned.
104 -func (b *quicBroker) AllocateFlow(addr *net.UDPAddr) uint32 {
105 - key := addr.String()
106 -
107 - b.mu.Lock()
108 - defer b.mu.Unlock()
109 -
110 - if id, ok := b.addrIndex[key]; ok {
111 - return id
112 - }
113 - id := b.nextFlow
114 - b.nextFlow++
115 - b.flowTable[id] = addr
116 - b.addrIndex[key] = id
117 - return id
118 -}
119 -
120 -// LookupFlowAddr returns the client address for a flow ID.
121 -func (b *quicBroker) LookupFlowAddr(flowID uint32) (*net.UDPAddr, bool) {
122 - b.mu.Lock()
123 - defer b.mu.Unlock()
124 - addr, ok := b.flowTable[flowID]
125 - return addr, ok
126 -}
127 -
128 -// Stop tears down the QUIC connection and signals done.
129 -func (b *quicBroker) Stop() {
130 - b.closeOnce.Do(func() {
131 - b.mu.Lock()
132 - b.closed = true
133 - conn := b.conn
134 - b.conn = nil
135 - b.mu.Unlock()
136 -
137 - if conn != nil {
138 - _ = conn.CloseWithError(0, "lease stopped")
139 - }
140 - close(b.done)
141 - })
142 -}
143 -
144 -func (b *quicBroker) receiveLoop(conn *quic.Conn) {
145 - for {
146 - data, err := conn.ReceiveDatagram(context.Background())
147 - if err != nil {
148 - b.mu.Lock()
149 - // Only clear if this is still the active connection.
150 - if b.conn == conn {
151 - b.conn = nil
152 - }
153 - b.mu.Unlock()
154 -
155 - if !b.isClosed() {
156 - log.Warn().
157 - Err(err).
158 - Str("component", "quic-broker").
159 - Str("lease_id", b.leaseID).
160 - Msg("quic receive loop ended")
161 - }
162 - return
163 - }
164 -
165 - frame, err := types.DecodeDatagram(data)
166 - if err != nil {
167 - continue
168 - }
169 -
170 - select {
171 - case b.incoming <- frame:
172 - default:
173 - // Drop if channel full — back-pressure on tunnel.
174 - }
175 - }
176 -}
177 -
178 -func (b *quicBroker) isClosed() bool {
179 - select {
180 - case <-b.done:
181 - return true
182 - default:
183 - return false
184 - }
185 -}
186 -
187 -// quicControlMessage is sent by the tunnel on the first QUIC stream after
188 -// connecting. The relay reads it to associate the connection with a lease.
189 -type quicControlMessage struct {
190 - LeaseID string `json:"lease_id"`
191 - ReverseToken string `json:"reverse_token"`
192 -}
193 -
194 -// quicTunnelListener manages the QUIC listener that accepts tunnel connections
195 -// on the relay API UDP port.
196 -type quicTunnelListener struct {
197 - listener *quic.Listener
198 - server *Server
199 -
200 - done chan struct{}
201 - closeOnce sync.Once
202 -}
203 -
204 -func newQUICTunnelListener(listener *quic.Listener, server *Server) *quicTunnelListener {
205 - return &quicTunnelListener{
206 - listener: listener,
207 - server: server,
208 - done: make(chan struct{}),
209 - }
210 -}
211 -
212 -func (l *quicTunnelListener) run() error {
213 - for {
214 - conn, err := l.listener.Accept(context.Background())
215 - if err != nil {
216 - select {
217 - case <-l.done:
218 - return nil
219 - default:
220 - }
221 - if errors.Is(err, quic.ErrServerClosed) {
222 - return nil
223 - }
224 - return err
225 - }
226 - go l.handleConnection(conn)
227 - }
228 -}
229 -
230 -func (l *quicTunnelListener) handleConnection(conn *quic.Conn) {
231 - stream, err := conn.AcceptStream(context.Background())
232 - if err != nil {
233 - _ = conn.CloseWithError(1, "stream accept failed")
234 - return
235 - }
236 -
237 - // Read control message with timeout.
238 - _ = stream.SetReadDeadline(time.Now().Add(10 * time.Second))
239 - var msg quicControlMessage
240 - buf := make([]byte, 4096)
241 - n, err := stream.Read(buf)
242 - if err != nil {
243 - _ = conn.CloseWithError(1, "control read failed")
244 - return
245 - }
246 -
247 - // Simple JSON decode.
248 - if decErr := json.Unmarshal(buf[:n], &msg); decErr != nil {
249 - _ = conn.CloseWithError(1, "invalid control message")
250 - return
251 - }
252 - _ = stream.SetReadDeadline(time.Time{})
253 -
254 - lease, err := l.server.findLeaseByID(msg.LeaseID)
255 - if err != nil {
256 - _, _ = stream.Write([]byte(`{"ok":false,"error":"lease_not_found"}`))
257 - _ = conn.CloseWithError(1, "lease not found")
258 - return
259 - }
260 -
261 - if authErr := l.server.authorizeLeaseToken(lease, msg.ReverseToken); authErr != nil {
262 - _, _ = stream.Write([]byte(`{"ok":false,"error":"unauthorized"}`))
263 - _ = conn.CloseWithError(1, "unauthorized")
264 - return
265 - }
266 -
267 - if lease.QUICBroker == nil {
268 - _, _ = stream.Write([]byte(`{"ok":false,"error":"transport_mismatch"}`))
269 - _ = conn.CloseWithError(1, "lease does not support QUIC transport")
270 - return
271 - }
272 -
273 - // Success — register the QUIC connection with the broker.
274 - lease.QUICBroker.Register(conn)
275 -
276 - // Confirm registration to the tunnel.
277 - _, _ = stream.Write([]byte(`{"ok":true}`))
278 -
279 - l.server.registry.Touch(lease.ID, conn.RemoteAddr().String(), time.Now())
280 -
281 - log.Info().
282 - Str("component", "quic-tunnel-listener").
283 - Str("lease_id", lease.ID).
284 - Str("lease_name", lease.Name).
285 - Str("remote_addr", conn.RemoteAddr().String()).
286 - Msg("quic tunnel connected")
287 -}
288 -
289 -func (l *quicTunnelListener) close() error {
290 - var closeErr error
291 - l.closeOnce.Do(func() {
292 - close(l.done)
293 - closeErr = l.listener.Close()
294 - })
295 - return closeErr
296 -}
portal/server.go
+287 -60
@@ -2,6 +2,7 @@ package portal
2
3 import (
4 "context"
5 + "encoding/json"
6 "errors"
7 "fmt"
8 "io"
@@ -18,6 +19,7 @@ import (
19 "golang.org/x/sync/errgroup"
20
21 "github.com/gosuda/portal/v2/portal/acme"
22 + portaldatagram "github.com/gosuda/portal/v2/portal/datagram"
23 "github.com/gosuda/portal/v2/portal/keyless"
24 "github.com/gosuda/portal/v2/portal/policy"
25 "github.com/gosuda/portal/v2/types"
@@ -32,11 +34,18 @@ const (
34 defaultClientHelloWait = 2 * time.Second
35 defaultControlBodyLimit = 4 << 20
36 defaultSessionWriteLimit = 5 * time.Second
37 + defaultQUICSNIRouteIdle = 30 * time.Second
38 + defaultQUICSNICleanup = 5 * time.Second
39
40 defaultUDPPortMin = 29000
41 defaultUDPPortMax = 29999
42 )
43
44 +type quicSNIRoute struct {
45 + flowMux *portaldatagram.FlowMux
46 + lastSeen time.Time
47 +}
48 +
49 type ServerConfig struct {
50 PortalURL string
51 ACME acme.Config
@@ -55,20 +64,22 @@ type ServerConfig struct {
64 }
65
66 type Server struct {
58 - sniListener net.Listener
59 - apiListener net.Listener
60 - apiServer *http.Server
61 - apiTLSClose io.Closer
62 - acmeManager *acme.Manager
63 - quicTunnel *quicTunnelListener
64 - quicSNIConn net.PacketConn
65 - cancel context.CancelFunc
66 - group *errgroup.Group
67 - registry *leaseRegistry
68 - ports *portAllocator
69 - cfg ServerConfig
70 - rootHost string
71 - shutdownOnce sync.Once
67 + sniListener net.Listener
68 + apiListener net.Listener
69 + apiServer *http.Server
70 + apiTLSClose io.Closer
71 + acmeManager *acme.Manager
72 + quicTunnel *quic.Listener
73 + quicSNI net.PacketConn
74 + cancel context.CancelFunc
75 + group *errgroup.Group
76 + registry *leaseRegistry
77 + ports *portaldatagram.PortAllocator
78 + cfg ServerConfig
79 + rootHost string
80 + shutdownOnce sync.Once
81 + quicSNIRoutes map[string]quicSNIRoute
82 + quicSNIMu sync.RWMutex
83 }
84
85 func NewServer(cfg ServerConfig) (*Server, error) {
@@ -94,13 +105,14 @@ func NewServer(cfg ServerConfig) (*Server, error) {
105 }
106
107 registry := newLeaseRegistry(policy.NewRuntime())
97 - ports := newPortAllocator(cfg.UDPPortMin, cfg.UDPPortMax, 5*time.Minute)
108 + ports := portaldatagram.NewPortAllocator(cfg.UDPPortMin, cfg.UDPPortMax, 5*time.Minute)
109
110 s := &Server{
100 - cfg: cfg,
101 - rootHost: rootHost,
102 - registry: registry,
103 - ports: ports,
111 + cfg: cfg,
112 + rootHost: rootHost,
113 + registry: registry,
114 + ports: ports,
115 + quicSNIRoutes: make(map[string]quicSNIRoute),
116 }
117
118 // Tear down all lease resources when leases expire via TTL janitor.
@@ -161,7 +173,7 @@ func (s *Server) Start(ctx context.Context, apiMux *http.ServeMux) error {
173 group.Go(func() error { return s.watchContext(groupCtx) })
174 s.acmeManager.Start(serverCtx)
175
164 - if err := s.startQUICTunnelListener(serverCtx, apiTLS); err != nil {
176 + if err := s.startQUICTunnelListener(apiTLS); err != nil {
177 log.Warn().Err(err).Msg("quic tunnel listener disabled")
178 }
179 if err := s.startQUICSNIRouter(); err != nil {
@@ -190,10 +202,10 @@ func (s *Server) Shutdown(ctx context.Context) error {
202 }
203
204 if s.quicTunnel != nil {
193 - _ = s.quicTunnel.close()
205 + _ = s.quicTunnel.Close()
206 }
195 - if s.quicSNIConn != nil {
196 - _ = s.quicSNIConn.Close()
207 + if s.quicSNI != nil {
208 + _ = s.quicSNI.Close()
209 }
210 if s.sniListener != nil {
211 if err := s.sniListener.Close(); err != nil && !errors.Is(err, net.ErrClosed) {
@@ -237,10 +249,10 @@ func (s *Server) SNIAddr() string {
249 }
250
251 func (s *Server) QUICAddr() string {
240 - if s.quicTunnel == nil || s.quicTunnel.listener == nil {
252 + if s.quicTunnel == nil {
253 return ""
254 }
243 - return s.quicTunnel.listener.Addr().String()
255 + return s.quicTunnel.Addr().String()
256 }
257
258 func (s *Server) LeaseSnapshots() []types.Lease {
@@ -330,12 +342,8 @@ func (s *Server) handleSNIConn(ctx context.Context, conn net.Conn) {
342 return
343 }
344
333 - record, ok := s.registry.Lookup(serverName)
334 - if !ok || time.Now().After(record.ExpiresAt) {
335 - _ = wrappedConn.Close()
336 - return
337 - }
338 - if !s.registry.policy.IsLeaseRoutable(record.ID) {
345 + broker, err := s.resolveStreamBroker(serverName)
346 + if err != nil {
347 _ = wrappedConn.Close()
348 return
349 }
@@ -343,7 +351,7 @@ func (s *Server) handleSNIConn(ctx context.Context, conn net.Conn) {
351 claimCtx, cancel := context.WithTimeout(ctx, s.cfg.ClaimTimeout)
352 defer cancel()
353
346 - session, err := record.Broker.Claim(claimCtx)
354 + session, err := broker.Claim(claimCtx)
355 if err != nil {
356 _ = wrappedConn.Close()
357 return
@@ -367,6 +375,59 @@ func (s *Server) bridgeToAPI(ctx context.Context, conn net.Conn) {
375 BridgeConns(conn, upstream)
376 }
377
378 +func (s *Server) lookupRoutableLease(serverName string) (*leaseRecord, error) {
379 + record, ok := s.registry.Lookup(serverName)
380 + if !ok || record == nil {
381 + return nil, errors.New("no route")
382 + }
383 + if time.Now().After(record.ExpiresAt) {
384 + return nil, errors.New("lease expired")
385 + }
386 + if !s.registry.policy.IsLeaseRoutable(record.ID) {
387 + return nil, errors.New("not routable")
388 + }
389 + return record, nil
390 +}
391 +
392 +func (s *Server) resolveStreamBroker(serverName string) (*streamBroker, error) {
393 + record, err := s.lookupRoutableLease(serverName)
394 + if err != nil {
395 + return nil, err
396 + }
397 + if !record.SupportsStream() {
398 + return nil, errors.New("transport mismatch")
399 + }
400 + streamBroker := record.StreamBroker()
401 + if streamBroker == nil {
402 + return nil, errors.New("stream broker unavailable")
403 + }
404 + return streamBroker, nil
405 +}
406 +
407 +func (s *Server) resolveDatagramFlowMux(serverName string) (*portaldatagram.FlowMux, error) {
408 + if serverName == s.rootHost {
409 + return nil, errors.New("root host does not accept datagram routes")
410 + }
411 +
412 + record, err := s.lookupRoutableLease(serverName)
413 + if err != nil {
414 + return nil, err
415 + }
416 + if !record.SupportsDatagram() {
417 + return nil, errors.New("transport mismatch")
418 + }
419 + flowMux := record.DatagramFlowMux()
420 + if flowMux == nil {
421 + return nil, errors.New("flow mux unavailable")
422 + }
423 + return flowMux, nil
424 +}
425 +
426 +type quicControlMessage struct {
427 + LeaseID string `json:"lease_id"`
428 + ReverseToken string `json:"reverse_token"`
429 +}
430 +
431 func (s *Server) watchContext(ctx context.Context) error {
432 <-ctx.Done()
433 shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
@@ -401,7 +462,7 @@ func closeWrite(conn net.Conn) {
462 }
463 }
464
404 -func (s *Server) startQUICTunnelListener(ctx context.Context, apiTLS keyless.TLSMaterialConfig) error {
465 +func (s *Server) startQUICTunnelListener(apiTLS keyless.TLSMaterialConfig) error {
466 if len(apiTLS.KeyPEM) == 0 {
467 return fmt.Errorf("quic tunnel requires api tls key")
468 }
@@ -416,9 +477,9 @@ func (s *Server) startQUICTunnelListener(ctx context.Context, apiTLS keyless.TLS
477 MinVersion: tls.VersionTLS13,
478 }
479 quicConf := &quic.Config{
419 - EnableDatagrams: true,
420 - KeepAlivePeriod: 15 * time.Second,
421 - MaxIdleTimeout: 60 * time.Second,
480 + EnableDatagrams: true,
481 + KeepAlivePeriod: 15 * time.Second,
482 + MaxIdleTimeout: 60 * time.Second,
483 MaxIncomingStreams: 16,
484 }
485
@@ -427,9 +488,8 @@ func (s *Server) startQUICTunnelListener(ctx context.Context, apiTLS keyless.TLS
488 return fmt.Errorf("listen quic: %w", err)
489 }
490
430 - tunnel := newQUICTunnelListener(listener, s)
431 - s.quicTunnel = tunnel
432 - s.group.Go(tunnel.run)
491 + s.quicTunnel = listener
492 + s.group.Go(func() error { return s.runQUICTunnelListener(listener) })
493
494 log.Info().
495 Str("component", "relay-server").
@@ -439,14 +499,14 @@ func (s *Server) startQUICTunnelListener(ctx context.Context, apiTLS keyless.TLS
499 }
500
501 func (s *Server) startQUICSNIRouter() error {
442 - conn, err := net.ListenPacket("udp", s.cfg.SNIListenAddr)
502 + var listenConfig net.ListenConfig
503 + conn, err := listenConfig.ListenPacket(context.Background(), "udp", s.cfg.SNIListenAddr)
504 if err != nil {
505 return fmt.Errorf("listen quic sni udp: %w", err)
506 }
507
447 - router := newQUICSNIRouter(conn, s)
448 - s.quicSNIConn = conn
449 - s.group.Go(router.run)
508 + s.quicSNI = conn
509 + s.group.Go(func() error { return s.runQUICSNIRouter(conn) })
510
511 log.Info().
512 Str("component", "relay-server").
@@ -455,31 +515,198 @@ func (s *Server) startQUICSNIRouter() error {
515 return nil
516 }
517
458 -func (s *Server) startUDPRelay(ctx context.Context, leaseID string, relay *udpRelay) {
459 - if relay == nil {
518 +func (s *Server) runQUICTunnelListener(listener *quic.Listener) error {
519 + for {
520 + conn, err := listener.Accept(context.Background())
521 + if err != nil {
522 + if errors.Is(err, quic.ErrServerClosed) || errors.Is(err, net.ErrClosed) {
523 + return nil
524 + }
525 + return err
526 + }
527 + go s.handleQUICTunnelConn(conn)
528 + }
529 +}
530 +
531 +func (s *Server) handleQUICTunnelConn(conn *quic.Conn) {
532 + stream, err := conn.AcceptStream(context.Background())
533 + if err != nil {
534 + _ = conn.CloseWithError(1, "stream accept failed")
535 + return
536 + }
537 +
538 + _ = stream.SetReadDeadline(time.Now().Add(10 * time.Second))
539 + var msg quicControlMessage
540 + buf := make([]byte, 4096)
541 + n, err := stream.Read(buf)
542 + if err != nil {
543 + _ = conn.CloseWithError(1, "control read failed")
544 + return
545 + }
546 + if err := json.Unmarshal(buf[:n], &msg); err != nil {
547 + _, _ = stream.Write([]byte(`{"ok":false,"error":"invalid_control_message"}`))
548 + _ = conn.CloseWithError(1, "invalid control message")
549 + return
550 + }
551 + _ = stream.SetReadDeadline(time.Time{})
552 +
553 + lease, err := s.findLeaseByID(msg.LeaseID)
554 + switch {
555 + case err != nil:
556 + _, _ = stream.Write([]byte(`{"ok":false,"error":"lease_not_found"}`))
557 + _ = conn.CloseWithError(1, "lease not found")
558 + return
559 + case s.authorizeLeaseToken(lease, msg.ReverseToken) != nil:
560 + _, _ = stream.Write([]byte(`{"ok":false,"error":"unauthorized"}`))
561 + _ = conn.CloseWithError(1, "unauthorized")
562 + return
563 + }
564 +
565 + flowMux := lease.DatagramFlowMux()
566 + if flowMux == nil {
567 + _, _ = stream.Write([]byte(`{"ok":false,"error":"transport_mismatch"}`))
568 + _ = conn.CloseWithError(1, "transport mismatch")
569 + return
570 + }
571 + if err := flowMux.Register(conn); err != nil {
572 + _, _ = stream.Write([]byte(`{"ok":false,"error":"broker_closed"}`))
573 + _ = conn.CloseWithError(1, "broker closed")
574 + return
575 + }
576 +
577 + _, _ = stream.Write([]byte(`{"ok":true}`))
578 + s.registry.Touch(lease.ID, conn.RemoteAddr().String(), time.Now())
579 + log.Info().
580 + Str("component", "quic-tunnel-listener").
581 + Str("lease_id", lease.ID).
582 + Str("lease_name", lease.Name).
583 + Str("remote_addr", conn.RemoteAddr().String()).
584 + Msg("quic tunnel connected")
585 +}
586 +
587 +func (s *Server) runQUICSNIRouter(conn net.PacketConn) error {
588 + buf := make([]byte, 65535)
589 + for {
590 + _ = conn.SetReadDeadline(time.Now().Add(defaultQUICSNICleanup))
591 + n, addr, err := conn.ReadFrom(buf)
592 + if err != nil {
593 + if errors.Is(err, net.ErrClosed) {
594 + return nil
595 + }
596 + var netErr net.Error
597 + if errors.As(err, &netErr) && netErr.Timeout() {
598 + s.cleanupQUICSNIRoutes(time.Now())
599 + continue
600 + }
601 + return err
602 + }
603 +
604 + packet := make([]byte, n)
605 + copy(packet, buf[:n])
606 + s.handleQUICSNIPacket(packet, addr, time.Now())
607 + }
608 +}
609 +
610 +func (s *Server) handleQUICSNIPacket(packet []byte, srcAddr net.Addr, now time.Time) {
611 + cacheKey := srcAddr.String()
612 + flowMux, ok := s.lookupQUICSNIRoute(cacheKey, now)
613 + if !ok {
614 + serverName, err := portaldatagram.ParseQUICInitialSNI(packet)
615 + if err != nil || serverName == "" {
616 + return
617 + }
618 +
619 + flowMux, err = s.resolveDatagramFlowMux(utils.NormalizeHostname(serverName))
620 + if err != nil || flowMux == nil {
621 + return
622 + }
623 + s.storeQUICSNIRoute(cacheKey, flowMux, now)
624 + }
625 +
626 + udpAddr, ok := srcAddr.(*net.UDPAddr)
627 + if !ok {
628 + return
629 + }
630 + if s.quicSNI == nil {
631 + return
632 + }
633 +
634 + flowID := flowMux.TouchFlow("quic:"+cacheKey, func(payload []byte) error {
635 + _, err := s.quicSNI.WriteTo(payload, udpAddr)
636 + return err
637 + })
638 + if err := flowMux.SendDatagram(flowID, packet); err != nil {
639 + s.deleteQUICSNIRoute(cacheKey)
640 + }
641 +}
642 +
643 +func (s *Server) lookupQUICSNIRoute(key string, now time.Time) (*portaldatagram.FlowMux, bool) {
644 + s.quicSNIMu.Lock()
645 + defer s.quicSNIMu.Unlock()
646 +
647 + route, ok := s.quicSNIRoutes[key]
648 + if !ok || route.flowMux == nil {
649 + delete(s.quicSNIRoutes, key)
650 + return nil, false
651 + }
652 + if now.Sub(route.lastSeen) > defaultQUICSNIRouteIdle || !route.flowMux.HasConnection() {
653 + delete(s.quicSNIRoutes, key)
654 + return nil, false
655 + }
656 +
657 + route.lastSeen = now
658 + s.quicSNIRoutes[key] = route
659 + return route.flowMux, true
660 +}
661 +
662 +func (s *Server) storeQUICSNIRoute(key string, flowMux *portaldatagram.FlowMux, now time.Time) {
663 + s.quicSNIMu.Lock()
664 + s.quicSNIRoutes[key] = quicSNIRoute{
665 + flowMux: flowMux,
666 + lastSeen: now,
667 + }
668 + s.quicSNIMu.Unlock()
669 +}
670 +
671 +func (s *Server) deleteQUICSNIRoute(key string) {
672 + s.quicSNIMu.Lock()
673 + delete(s.quicSNIRoutes, key)
674 + s.quicSNIMu.Unlock()
675 +}
676 +
677 +func (s *Server) cleanupQUICSNIRoutes(now time.Time) {
678 + s.quicSNIMu.Lock()
679 + defer s.quicSNIMu.Unlock()
680 +
681 + for key, route := range s.quicSNIRoutes {
682 + if route.flowMux == nil || now.Sub(route.lastSeen) > defaultQUICSNIRouteIdle || !route.flowMux.HasConnection() {
683 + delete(s.quicSNIRoutes, key)
684 + }
685 + }
686 +}
687 +
688 +func (s *Server) clearQUICSNIRoutesForFlowMux(flowMux *portaldatagram.FlowMux) {
689 + if flowMux == nil {
690 return
691 }
462 - if err := relay.Start(ctx); err != nil {
463 - log.Error().
464 - Err(err).
465 - Str("component", "relay-server").
466 - Str("lease_id", leaseID).
467 - Msg("failed to start udp relay")
692 +
693 + s.quicSNIMu.Lock()
694 + defer s.quicSNIMu.Unlock()
695 +
696 + for key, route := range s.quicSNIRoutes {
697 + if route.flowMux == flowMux {
698 + delete(s.quicSNIRoutes, key)
699 + }
700 }
701 }
702
703 // closeLease tears down all resources associated with a single lease record.
704 func (s *Server) closeLease(record *leaseRecord) {
473 - record.Broker.Close()
474 - if record.QUICBroker != nil {
475 - record.QUICBroker.Stop()
476 - }
477 - if record.UDPRelay != nil {
478 - record.UDPRelay.Stop()
479 - }
480 - if record.UDPPort > 0 && s.ports != nil {
481 - s.ports.Release(record.UDPPort)
705 + if record == nil || record.Runtime == nil {
706 + return
707 }
708 + s.clearQUICSNIRoutesForFlowMux(record.DatagramFlowMux())
709 + record.Runtime.Close(s.ports)
710 }
711
712 func (s *Server) quicPublicAddr() string {
portal/server_test.go
+46
@@ -151,3 +151,49 @@ func TestRegisterLeaseRejectsInvalidName(t *testing.T) {
151 t.Fatal("registerLease() error = nil, want invalid name error")
152 }
153 }
154 +
155 +func TestRegisterLeaseBuildsDatagramOnlyRuntime(t *testing.T) {
156 + t.Parallel()
157 +
158 + server, err := NewServer(ServerConfig{
159 + PortalURL: "https://portal.example.com",
160 + })
161 + if err != nil {
162 + t.Fatalf("NewServer() error = %v", err)
163 + }
164 +
165 + resp, err := server.registerLease(types.RegisterRequest{
166 + Name: "demo-udp",
167 + ReverseToken: "tok_udp",
168 + Transport: types.TransportUDP,
169 + }, "203.0.113.10")
170 + if err != nil {
171 + t.Fatalf("registerLease() error = %v", err)
172 + }
173 +
174 + record, ok := server.registry.Get(resp.LeaseID)
175 + if !ok {
176 + t.Fatal("registry.Get() = false, want registered lease")
177 + }
178 + if record.SupportsStream() {
179 + t.Fatal("SupportsStream() = true, want false")
180 + }
181 + if !record.SupportsDatagram() {
182 + t.Fatal("SupportsDatagram() = false, want true")
183 + }
184 + if record.StreamBroker() != nil {
185 + t.Fatal("StreamBroker() != nil, want nil")
186 + }
187 + if record.DatagramFlowMux() == nil {
188 + t.Fatal("DatagramFlowMux() = nil, want flow mux")
189 + }
190 + if got := record.UDPPort(); got == 0 {
191 + t.Fatal("UDPPort() = 0, want allocated port")
192 + }
193 + if resp.UDPAddr == "" {
194 + t.Fatal("RegisterResponse.UDPAddr = empty, want public udp address")
195 + }
196 + if resp.QUICAddr == "" {
197 + t.Fatal("RegisterResponse.QUICAddr = empty, want quic address")
198 + }
199 +}
portal/udp_port.go deleted
-119
@@ -1,119 +0,0 @@
1 -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 {
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 -
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,
43 - inUse: make(map[int]string),
44 - reserved: make(map[string]portReservation),
45 - grace: grace,
46 - }
47 -}
48 -
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:]
71 - a.inUse[port] = name
72 - return port, nil
73 -}
74 -
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 -
81 - name, ok := a.inUse[port]
82 - if !ok {
83 - return
84 - }
85 - delete(a.inUse, 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 deleted
-199
@@ -1,199 +0,0 @@
1 -package portal
2 -
3 -import (
4 - "context"
5 - "fmt"
6 - "net"
7 - "sync"
8 - "time"
9 -
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
21 - Addr *net.UDPAddr
22 - LastSeen time.Time
23 -}
24 -
25 -// udpRelay binds a UDP port for a lease and relays datagrams bidirectionally
26 -// between raw UDP clients and the tunnel's QUIC connection via the quicBroker.
27 -type udpRelay struct {
28 - leaseID string
29 - port int
30 - broker *quicBroker
31 - conn *net.UDPConn
32 -
33 - sessions map[string]*udpSession // "ip:port" → session
34 - mu sync.Mutex
35 -
36 - cancel context.CancelFunc
37 - done chan struct{}
38 - closeOnce sync.Once
39 -}
40 -
41 -func newUDPRelay(leaseID string, port int, broker *quicBroker) *udpRelay {
42 - return &udpRelay{
43 - leaseID: leaseID,
44 - port: port,
45 - broker: broker,
46 - sessions: make(map[string]*udpSession),
47 - done: make(chan struct{}),
48 - }
49 -}
50 -
51 -// Start binds the UDP port and launches read/write relay goroutines.
52 -func (r *udpRelay) Start(ctx context.Context) error {
53 - addr := &net.UDPAddr{Port: r.port}
54 - conn, err := net.ListenUDP("udp", addr)
55 - if err != nil {
56 - return fmt.Errorf("listen udp :%d: %w", r.port, err)
57 - }
58 - r.conn = conn
59 -
60 - relayCtx, cancel := context.WithCancel(ctx)
61 - r.cancel = cancel
62 -
63 - go r.readLoop(relayCtx)
64 - go r.writeLoop(relayCtx)
65 - go r.sessionCleanup(relayCtx)
66 -
67 - log.Info().
68 - Str("component", "udp-relay").
69 - Str("lease_id", r.leaseID).
70 - Int("port", r.port).
71 - Msg("udp relay started")
72 -
73 - return nil
74 -}
75 -
76 -// Stop closes the UDP socket and cancels the relay context.
77 -func (r *udpRelay) Stop() {
78 - r.closeOnce.Do(func() {
79 - if r.cancel != nil {
80 - r.cancel()
81 - }
82 - if r.conn != nil {
83 - _ = r.conn.Close()
84 - }
85 - close(r.done)
86 - log.Info().
87 - Str("component", "udp-relay").
88 - Str("lease_id", r.leaseID).
89 - Int("port", r.port).
90 - Msg("udp relay stopped")
91 - })
92 -}
93 -
94 -// readLoop reads raw UDP from public clients and forwards via QUIC DATAGRAM.
95 -func (r *udpRelay) readLoop(ctx context.Context) {
96 - buf := make([]byte, defaultMaxDatagramSize)
97 - for {
98 - select {
99 - case <-ctx.Done():
100 - return
101 - default:
102 - }
103 -
104 - _ = r.conn.SetReadDeadline(time.Now().Add(5 * time.Second))
105 - n, clientAddr, err := r.conn.ReadFromUDP(buf)
106 - if err != nil {
107 - if ctx.Err() != nil {
108 - return
109 - }
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 -
121 - flowID := r.getOrCreateFlow(clientAddr)
122 - payload := make([]byte, n)
123 - copy(payload, buf[:n])
124 -
125 - if err := r.broker.SendDatagram(flowID, payload); err != nil {
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 - }
136 -}
137 -
138 -// writeLoop receives QUIC DATAGRAM frames from the tunnel and sends raw UDP back.
139 -func (r *udpRelay) writeLoop(ctx context.Context) {
140 - for {
141 - select {
142 - case <-ctx.Done():
143 - return
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)
155 - }
156 - }
157 -}
158 -
159 -func (r *udpRelay) getOrCreateFlow(addr *net.UDPAddr) uint32 {
160 - key := addr.String()
161 -
162 - r.mu.Lock()
163 - defer r.mu.Unlock()
164 -
165 - if s, ok := r.sessions[key]; ok {
166 - s.LastSeen = time.Now()
167 - return s.FlowID
168 - }
169 -
170 - flowID := r.broker.AllocateFlow(addr)
171 - r.sessions[key] = &udpSession{
172 - FlowID: flowID,
173 - Addr: addr,
174 - LastSeen: time.Now(),
175 - }
176 - return flowID
177 -}
178 -
179 -// sessionCleanup periodically removes idle UDP sessions.
180 -func (r *udpRelay) sessionCleanup(ctx context.Context) {
181 - ticker := time.NewTicker(defaultUDPSessionTimeout)
182 - defer ticker.Stop()
183 -
184 - for {
185 - select {
186 - case <-ctx.Done():
187 - return
188 - case <-ticker.C:
189 - now := time.Now()
190 - r.mu.Lock()
191 - for key, s := range r.sessions {
192 - if now.Sub(s.LastSeen) > defaultUDPSessionTimeout {
193 - delete(r.sessions, key)
194 - }
195 - }
196 - r.mu.Unlock()
197 - }
198 - }
199 -}
sdk/expose.go
+206 -60
@@ -20,18 +20,43 @@ import (
20 // Exposure owns the lifecycle of one or more relay listeners and accepts
21 // traffic from all of them through one net.Listener.
22 type Exposure struct {
23 - listener net.Listener
24 - listeners []*Listener
25 - done chan struct{}
23 + capabilities types.LeaseCapabilities
24 + listener net.Listener
25 + listeners []*Listener
26 + datagrams chan ExposureDatagram
27 + done chan struct{}
28
29 closeOnce sync.Once
30 connSeq atomic.Uint64
31 }
32
33 +// ExposureDatagram represents one datagram received from any relay backing an
34 +// exposure. Reply sends a response back through the same relay flow.
35 +type ExposureDatagram struct {
36 + FlowID uint32
37 + LeaseID string
38 + Payload []byte
39 + RelayURL string
40 + UDPAddr string
41 +
42 + reply func([]byte) error
43 +}
44 +
45 +func (d ExposureDatagram) Reply(payload []byte) error {
46 + if d.reply == nil {
47 + return errors.New("reply path is unavailable")
48 + }
49 + return d.reply(payload)
50 +}
51 +
52 // Expose creates relay listeners for each normalized relay URL and exposes a
53 // merged listener for accepting traffic from all of them. Empty relay input
54 // returns nil, nil so callers can fall back to local-only serving.
55 func Expose(ctx context.Context, relayUrls []string, name string, transport string, metadata types.LeaseMetadata) (*Exposure, error) {
56 + if ctx == nil {
57 + ctx = context.Background()
58 + }
59 +
60 relayURLs, err := utils.NormalizeRelayURLs(relayUrls)
61 if err != nil {
62 return nil, err
@@ -39,6 +64,11 @@ func Expose(ctx context.Context, relayUrls []string, name string, transport stri
64 if len(relayURLs) == 0 {
65 return nil, nil
66 }
67 + capabilities, err := types.ParseLeaseCapabilities(transport)
68 + if err != nil {
69 + return nil, err
70 + }
71 +
72 listeners := make([]*Listener, 0, len(relayURLs))
73 cleanup := func() error {
74 var closeErr error
@@ -53,7 +83,7 @@ func Expose(ctx context.Context, relayUrls []string, name string, transport stri
83 for _, relayURL := range relayURLs {
84 listener, err := NewListener(ctx, relayURL, ListenerConfig{
85 Name: name,
56 - Transport: transport,
86 + Transport: capabilities.Transport(),
87 Metadata: metadata,
88 })
89 if err != nil {
@@ -63,22 +93,30 @@ func Expose(ctx context.Context, relayUrls []string, name string, transport stri
93 listeners = append(listeners, listener)
94 }
95
66 - mergedListeners := make([]net.Listener, 0, len(listeners))
67 - for _, listener := range listeners {
68 - mergedListeners = append(mergedListeners, listener)
69 - }
96 + var merged net.Listener
97 + if capabilities.SupportsStream() {
98 + mergedListeners := make([]net.Listener, 0, len(listeners))
99 + for _, listener := range listeners {
100 + mergedListeners = append(mergedListeners, listener)
101 + }
102
71 - merged, err := mergeListeners(mergedListeners...)
72 - if err != nil {
73 - return nil, errors.Join(fmt.Errorf("merge listeners: %w", err), cleanup())
103 + merged, err = mergeListeners(mergedListeners...)
104 + if err != nil {
105 + return nil, errors.Join(fmt.Errorf("merge listeners: %w", err), cleanup())
106 + }
107 }
108
109 exposure := &Exposure{
77 - listener: merged,
78 - listeners: listeners,
79 - done: make(chan struct{}),
110 + capabilities: capabilities,
111 + listener: merged,
112 + listeners: listeners,
113 + datagrams: make(chan ExposureDatagram, max(len(listeners)*32, 1)),
114 + done: make(chan struct{}),
115 }
116 go exposure.monitorStartupCounts(ctx)
117 + if exposure.SupportsDatagram() {
118 + exposure.attachDatagramPlanes(ctx)
119 + }
120
121 log.Info().
122 Str("release_version", types.ReleaseVersion).
@@ -129,6 +167,21 @@ func (e *Exposure) Addr() net.Addr {
167 return e.listener.Addr()
168 }
169
170 +// AcceptDatagram returns datagrams from any relay datagram plane attached to
171 +// the exposure.
172 +func (e *Exposure) AcceptDatagram() (ExposureDatagram, error) {
173 + if e == nil || !e.SupportsDatagram() {
174 + return ExposureDatagram{}, net.ErrClosed
175 + }
176 +
177 + select {
178 + case <-e.done:
179 + return ExposureDatagram{}, net.ErrClosed
180 + case dg := <-e.datagrams:
181 + return dg, nil
182 + }
183 +}
184 +
185 // RelayURLs returns the normalized relay URLs backing the exposure.
186 func (e *Exposure) RelayURLs() []string {
187 if e == nil || len(e.listeners) == 0 {
@@ -145,6 +198,38 @@ func (e *Exposure) RelayURLs() []string {
198 return out
199 }
200
201 +// UDPAddrs returns the current public UDP addresses exposed by the datagram
202 +// plane, deduplicated across all backing relays.
203 +func (e *Exposure) UDPAddrs() []string {
204 + if e == nil || len(e.listeners) == 0 || !e.SupportsDatagram() {
205 + return nil
206 + }
207 +
208 + out := make([]string, 0, len(e.listeners))
209 + seen := make(map[string]struct{})
210 + for _, listener := range e.listeners {
211 + if listener == nil {
212 + continue
213 + }
214 +
215 + listener.mu.Lock()
216 + udpAddr := listener.udpAddr
217 + listener.mu.Unlock()
218 + if udpAddr == "" {
219 + continue
220 + }
221 + if _, ok := seen[udpAddr]; ok {
222 + continue
223 + }
224 + seen[udpAddr] = struct{}{}
225 + out = append(out, udpAddr)
226 + }
227 + if len(out) == 0 {
228 + return nil
229 + }
230 + return out
231 +}
232 +
233 // PublicURLs returns the de-duplicated public URLs exposed by the exposure.
234 func (e *Exposure) PublicURLs() []string {
235 if e == nil || len(e.listeners) == 0 {
@@ -178,70 +263,105 @@ func (e *Exposure) PublicURLs() []string {
263 // local-only serving.
264 func (e *Exposure) RunHTTP(ctx context.Context, handler http.Handler, localAddr string) error {
265 var relayListener net.Listener
181 - if e != nil {
266 + if e != nil && e.listener != nil {
267 relayListener = e
268 }
269 return RunHTTP(ctx, relayListener, handler, localAddr)
270 }
271
187 -// AttachUDP creates UDPListeners for each underlying relay listener that has
188 -// UDP transport enabled. Listeners without UDP support are silently skipped.
189 -func (e *Exposure) AttachUDP(ctx context.Context) ([]*UDPListener, error) {
190 - if e == nil || len(e.listeners) == 0 {
191 - return nil, nil
272 +// WaitDatagramReady blocks until at least one backing relay has published a
273 +// public UDP address for this exposure.
274 +func (e *Exposure) WaitDatagramReady(ctx context.Context) ([]string, error) {
275 + if e == nil || !e.SupportsDatagram() {
276 + return nil, errors.New("exposure does not have datagram transport enabled")
277 + }
278 + if ctx == nil {
279 + ctx = context.Background()
280 }
281
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()
282 + ticker := time.NewTicker(50 * time.Millisecond)
283 + defer ticker.Stop()
284
199 - results := make(chan *UDPListener, len(e.listeners))
285 + for {
286 + if addrs := e.UDPAddrs(); len(addrs) > 0 {
287 + return addrs, nil
288 + }
289
201 - count := 0
202 - for _, relay := range e.listeners {
203 - if relay == nil {
290 + select {
291 + case <-e.done:
292 + return nil, net.ErrClosed
293 + case <-ctx.Done():
294 + return nil, ctx.Err()
295 + case <-ticker.C:
296 + }
297 + }
298 +}
299 +
300 +func (e *Exposure) attachDatagramPlanes(ctx context.Context) {
301 + for _, listener := range e.listeners {
302 + if listener == nil {
303 continue
304 }
206 - count++
207 - go func(l *Listener) {
208 - udpL, err := l.AttachUDP(waitCtx)
209 - if err != nil {
210 - results <- nil
305 +
306 + go e.attachDatagramPlane(ctx, listener)
307 + }
308 +}
309 +
310 +func (e *Exposure) attachDatagramPlane(ctx context.Context, listener *Listener) {
311 + err := listener.WaitDatagramReady(ctx)
312 + if err != nil {
313 + switch {
314 + case e.closed():
315 + return
316 + case ctx != nil && ctx.Err() != nil:
317 + return
318 + case errors.Is(err, net.ErrClosed), errors.Is(err, context.Canceled):
319 + return
320 + default:
321 + log.Warn().
322 + Err(err).
323 + Str("relay_url", listener.relayURL).
324 + Msg("attach datagram plane failed")
325 + return
326 + }
327 + }
328 +
329 + e.forwardDatagrams(listener.relayURL, listener)
330 +}
331 +
332 +func (e *Exposure) forwardDatagrams(relayURL string, listener *Listener) {
333 + for {
334 + dg, err := listener.AcceptDatagram()
335 + if err != nil {
336 + if e.closed() || errors.Is(err, net.ErrClosed) {
337 return
338 }
213 - results <- udpL
214 - }(relay)
215 - }
339 + log.Warn().
340 + Err(err).
341 + Str("relay_url", relayURL).
342 + Str("lease_id", listener.LeaseID()).
343 + Msg("datagram accept failed")
344 + return
345 + }
346
217 - var out []*UDPListener
218 - var graceTimer *time.Timer
219 - defer func() {
220 - if graceTimer != nil {
221 - graceTimer.Stop()
347 + flowID := dg.FlowID
348 + reply := func(payload []byte) error {
349 + return listener.SendDatagram(flowID, payload)
350 }
223 - }()
224 - var grace <-chan time.Time
225 - collected := 0
226 - for collected < count {
351 +
352 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
353 + case <-e.done:
354 + return
355 + case e.datagrams <- ExposureDatagram{
356 + FlowID: flowID,
357 + LeaseID: listener.LeaseID(),
358 + Payload: append([]byte(nil), dg.Payload...),
359 + RelayURL: relayURL,
360 + UDPAddr: listener.UDPAddr(),
361 + reply: reply,
362 + }:
363 }
364 }
244 - return out, nil
365 }
366
367 // Close closes the merged listener and all underlying relay listeners.
@@ -258,6 +378,11 @@ func (e *Exposure) Close() error {
378 if e.listener != nil {
379 closeErr = errors.Join(closeErr, e.listener.Close())
380 }
381 + for _, listener := range e.listeners {
382 + if listener != nil {
383 + closeErr = errors.Join(closeErr, listener.Close())
384 + }
385 + }
386
387 event := log.Info().
388 Int("relay_count", len(e.listeners)).
@@ -273,6 +398,27 @@ func (e *Exposure) Close() error {
398 return closeErr
399 }
400
401 +func (e *Exposure) SupportsDatagram() bool {
402 + return e != nil && e.capabilities.SupportsDatagram()
403 +}
404 +
405 +func (e *Exposure) SupportsStream() bool {
406 + return e != nil && e.capabilities.SupportsStream()
407 +}
408 +
409 +func (e *Exposure) closed() bool {
410 + if e == nil || e.done == nil {
411 + return true
412 + }
413 +
414 + select {
415 + case <-e.done:
416 + return true
417 + default:
418 + return false
419 + }
420 +}
421 +
422 // RunHTTP serves one handler on relayListener and, when localAddr is set, on
423 // the provided local HTTP address for app-local access.
424 func RunHTTP(ctx context.Context, relayListener net.Listener, handler http.Handler, localAddr string) error {
sdk/listener.go
+217 -13
@@ -14,6 +14,7 @@ import (
14
15 "github.com/rs/zerolog/log"
16
17 + "github.com/gosuda/portal/v2/portal/datagram"
18 "github.com/gosuda/portal/v2/portal/keyless"
19 "github.com/gosuda/portal/v2/types"
20 "github.com/gosuda/portal/v2/utils"
@@ -42,6 +43,12 @@ const (
43 listenerStatusReady listenerStatus = "ready"
44 )
45
46 +type datagramState struct {
47 + leaseID string
48 + reverseToken string
49 + quicAddr string
50 +}
51 +
52 type Listener struct {
53 tlsCloser io.Closer
54 tlsConfig *tls.Config
@@ -55,6 +62,7 @@ type Listener struct {
62 cancel context.CancelFunc
63 api *apiClient
64 accepted chan net.Conn
65 + capabilities types.LeaseCapabilities
66 relayURL string
67 transport string
68 startupStatus listenerStatus
@@ -64,6 +72,7 @@ type Listener struct {
72 udpAddr string
73 quicAddr string
74 metadata types.LeaseMetadata
75 + datagram *datagram.Session
76
77 registered chan struct{} // closed after first successful registration
78 closeOnce sync.Once
@@ -91,13 +100,20 @@ func NewListener(ctx context.Context, relayURL string, cfg ListenerConfig) (*Lis
100 return nil, err
101 }
102
94 - transport := strings.ToLower(strings.TrimSpace(cfg.Transport))
103 + capabilities, err := types.ParseLeaseCapabilities(cfg.Transport)
104 + if err != nil {
105 + cancel()
106 + api.close()
107 + return nil, err
108 + }
109 + transport := capabilities.Transport()
110
111 l := &Listener{
112 doneCh: listenerCtx.Done(),
113 cancel: cancel,
114 api: api,
115 accepted: make(chan net.Conn, max(readyTarget*2, 1)),
116 + capabilities: capabilities,
117 registered: make(chan struct{}),
118 relayURL: api.baseURL.String(),
119 transport: transport,
@@ -109,7 +125,19 @@ func NewListener(ctx context.Context, relayURL string, cfg ListenerConfig) (*Lis
125 renewBefore: renewBefore,
126 handshakeTimeout: handshakeTimeout,
127 }
128 + if capabilities.SupportsDatagram() {
129 + l.datagram = datagram.NewSession(256, false, func(err error) {
130 + log.Warn().
131 + Err(err).
132 + Str("component", "sdk-datagram-plane").
133 + Str("lease_id", l.LeaseID()).
134 + Msg("quic receive loop ended")
135 + })
136 + }
137
138 + if l.datagram != nil {
139 + go l.runDatagramLoop(listenerCtx)
140 + }
141 go l.runStartup(listenerCtx)
142 return l, nil
143 }
@@ -121,16 +149,22 @@ func (l *Listener) runStartup(ctx context.Context) {
149 err := l.registerAndConfigure(ctx)
150 switch {
151 case err == nil:
124 - for i := 0; i < l.readyTarget; i++ {
125 - go l.runSessionLoop(ctx)
152 + if l.SupportsStream() {
153 + for i := 0; i < l.readyTarget; i++ {
154 + go l.runSessionLoop(ctx)
155 + }
156 + } else {
157 + l.setStartupStatus(listenerStatusReady)
158 }
159 go l.runRenewLoop(ctx)
160 publicURL := l.PublicURL()
129 - log.Info().
161 + event := log.Info().
162 Str("relay_url", l.relayURL).
131 - Str("lease_id", l.LeaseID()).
132 - Str("public_url", publicURL).
133 - Msg("service is available at this URL")
163 + Str("lease_id", l.LeaseID())
164 + if publicURL != "" {
165 + event = event.Str("public_url", publicURL)
166 + }
167 + event.Msg("relay listener registered")
168 return
169 case errors.Is(err, context.Canceled), errors.Is(err, net.ErrClosed):
170 return
@@ -144,6 +178,9 @@ func (l *Listener) runStartup(ctx context.Context) {
178 }
179
180 func (l *Listener) Accept() (net.Conn, error) {
181 + if !l.SupportsStream() {
182 + return nil, net.ErrClosed
183 + }
184 select {
185 case <-l.doneCh:
186 return nil, net.ErrClosed
@@ -162,6 +199,7 @@ func (l *Listener) Close() error {
199 l.mu.Lock()
200 leaseID := l.leaseID
201 tlsCloser := l.tlsCloser
202 + datagram := l.datagram
203 api := l.api
204 l.leaseID = ""
205 l.hostname = ""
@@ -170,6 +208,9 @@ func (l *Listener) Close() error {
208 l.mu.Unlock()
209
210 l.drainAccepted()
211 + if datagram != nil {
212 + datagram.Stop("listener closed")
213 + }
214
215 if api != nil && leaseID != "" {
216 ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
@@ -214,6 +255,10 @@ func (l *Listener) Metadata() types.LeaseMetadata {
255 }
256
257 func (l *Listener) PublicURL() string {
258 + if !l.SupportsStream() {
259 + return ""
260 + }
261 +
262 l.mu.Lock()
263 hostname := l.hostname
264 relayURL := l.relayURL
@@ -265,6 +310,73 @@ func (l *Listener) runSessionLoop(ctx context.Context) {
310 }
311 }
312
313 +func (l *Listener) runDatagramLoop(ctx context.Context) {
314 + for {
315 + select {
316 + case <-ctx.Done():
317 + l.datagram.Stop("listener context closed")
318 + return
319 + default:
320 + }
321 +
322 + state, ok := l.currentDatagramState()
323 + if !ok {
324 + if !utils.SleepOrDone(ctx, time.Second) {
325 + l.datagram.Stop("listener context closed")
326 + return
327 + }
328 + continue
329 + }
330 +
331 + conn, err := l.api.openQUICSession(ctx, state.quicAddr, state.leaseID, state.reverseToken)
332 + if err != nil {
333 + log.Warn().
334 + Err(err).
335 + Str("component", "sdk-datagram-plane").
336 + Str("lease_id", state.leaseID).
337 + Msg("quic session open failed, retrying")
338 + if !utils.SleepOrDone(ctx, 2*time.Second) {
339 + l.datagram.Stop("listener context closed")
340 + return
341 + }
342 + continue
343 + }
344 +
345 + log.Info().
346 + Str("component", "sdk-datagram-plane").
347 + Str("lease_id", state.leaseID).
348 + Str("remote_addr", conn.RemoteAddr().String()).
349 + Msg("quic tunnel connected")
350 +
351 + recvDone, err := l.datagram.Bind(conn)
352 + if err != nil {
353 + if ctx.Err() != nil {
354 + return
355 + }
356 + log.Warn().
357 + Err(err).
358 + Str("component", "sdk-datagram-plane").
359 + Str("lease_id", state.leaseID).
360 + Msg("quic session bind failed")
361 + if !utils.SleepOrDone(ctx, time.Second) {
362 + return
363 + }
364 + continue
365 + }
366 +
367 + select {
368 + case <-ctx.Done():
369 + l.datagram.Stop("listener context closed")
370 + return
371 + case <-recvDone:
372 + }
373 +
374 + if !utils.SleepOrDone(ctx, time.Second) {
375 + return
376 + }
377 + }
378 +}
379 +
380 func (l *Listener) runRenewLoop(ctx context.Context) {
381 interval := l.leaseTTL / 2
382 if interval <= 0 {
@@ -389,15 +501,23 @@ func (l *Listener) registerAndConfigure(ctx context.Context) error {
501 return err
502 }
503
392 - tlsConf, tlsCloser, err := keyless.BuildClientTLSConfig(l.api.baseURL.String(), []string{resp.Hostname})
393 - if err != nil {
394 - _ = l.api.unregisterLease(context.Background(), resp.LeaseID)
395 - return err
504 + var (
505 + tlsConf *tls.Config
506 + tlsCloser io.Closer
507 + )
508 + if l.SupportsStream() {
509 + tlsConf, tlsCloser, err = keyless.BuildClientTLSConfig(l.api.baseURL.String(), []string{resp.Hostname})
510 + if err != nil {
511 + _ = l.api.unregisterLease(context.Background(), resp.LeaseID)
512 + return err
513 + }
514 }
515
516 if ctx.Err() != nil {
517 _ = l.api.unregisterLease(context.Background(), resp.LeaseID)
400 - _ = tlsCloser.Close()
518 + if tlsCloser != nil {
519 + _ = tlsCloser.Close()
520 + }
521 return ctx.Err()
522 }
523
@@ -405,10 +525,13 @@ func (l *Listener) registerAndConfigure(ctx context.Context) error {
525 if ctx.Err() != nil {
526 l.mu.Unlock()
527 _ = l.api.unregisterLease(context.Background(), resp.LeaseID)
408 - _ = tlsCloser.Close()
528 + if tlsCloser != nil {
529 + _ = tlsCloser.Close()
530 + }
531 return ctx.Err()
532 }
533 oldCloser := l.tlsCloser
534 + datagram := l.datagram
535 l.leaseID = resp.LeaseID
536 l.hostname = resp.Hostname
537 l.udpAddr = resp.UDPAddr
@@ -421,12 +544,78 @@ func (l *Listener) registerAndConfigure(ctx context.Context) error {
544 if oldCloser != nil {
545 _ = oldCloser.Close()
546 }
547 + if datagram != nil {
548 + datagram.Clear("lease updated")
549 + }
550 l.registerOnce.Do(func() { close(l.registered) })
551 return nil
552 }
553
554 +func (l *Listener) SupportsDatagram() bool {
555 + return l != nil && l.capabilities.SupportsDatagram()
556 +}
557 +
558 +func (l *Listener) SupportsStream() bool {
559 + return l != nil && l.capabilities.SupportsStream()
560 +}
561 +
562 +func (l *Listener) AcceptDatagram() (types.DatagramFrame, error) {
563 + if l == nil || !l.SupportsDatagram() || l.datagram == nil {
564 + return types.DatagramFrame{}, net.ErrClosed
565 + }
566 +
567 + select {
568 + case <-l.doneCh:
569 + return types.DatagramFrame{}, net.ErrClosed
570 + case dg := <-l.datagram.Incoming():
571 + return dg, nil
572 + }
573 +}
574 +
575 +func (l *Listener) SendDatagram(flowID uint32, payload []byte) error {
576 + if l == nil || !l.SupportsDatagram() || l.datagram == nil {
577 + return net.ErrClosed
578 + }
579 + return l.datagram.Send(flowID, payload)
580 +}
581 +
582 +func (l *Listener) UDPAddr() string {
583 + l.mu.Lock()
584 + defer l.mu.Unlock()
585 + return l.udpAddr
586 +}
587 +
588 +func (l *Listener) QUICAddr() string {
589 + l.mu.Lock()
590 + defer l.mu.Unlock()
591 + return l.quicAddr
592 +}
593 +
594 +func (l *Listener) currentDatagramState() (datagramState, bool) {
595 + if l == nil || !l.SupportsDatagram() {
596 + return datagramState{}, false
597 + }
598 +
599 + l.mu.Lock()
600 + defer l.mu.Unlock()
601 +
602 + if l.api == nil || l.leaseID == "" {
603 + return datagramState{}, false
604 + }
605 +
606 + return datagramState{
607 + leaseID: l.leaseID,
608 + reverseToken: l.api.reverseToken,
609 + quicAddr: l.quicAddr,
610 + }, true
611 +}
612 +
613 // WaitRegistered blocks until the first successful lease registration or context cancellation.
614 func (l *Listener) WaitRegistered(ctx context.Context) error {
615 + if ctx == nil {
616 + ctx = context.Background()
617 + }
618 +
619 select {
620 case <-l.registered:
621 return nil
@@ -437,6 +626,21 @@ func (l *Listener) WaitRegistered(ctx context.Context) error {
626 }
627 }
628
629 +// WaitDatagramReady blocks until the listener has registered a datagram-capable
630 +// lease and the relay has assigned its UDP/QUIC endpoints.
631 +func (l *Listener) WaitDatagramReady(ctx context.Context) error {
632 + if l == nil || !l.SupportsDatagram() {
633 + return errors.New("lease does not have datagram transport enabled")
634 + }
635 + if err := l.WaitRegistered(ctx); err != nil {
636 + return err
637 + }
638 + if l.UDPAddr() == "" && l.QUICAddr() == "" {
639 + return errors.New("lease registration did not expose datagram addresses")
640 + }
641 + return nil
642 +}
643 +
644 func (l *Listener) reregister(ctx context.Context) error {
645 requestCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
646 defer cancel()
sdk/sdk_test.go
+149
@@ -154,6 +154,9 @@ func TestNewListenerRegistersLeaseWithMainContract(t *testing.T) {
154 if registerReq.TTL != 42 {
155 t.Fatalf("register request TTL = %d, want 42", registerReq.TTL)
156 }
157 + if registerReq.Transport != types.TransportTCP {
158 + t.Fatalf("register request Transport = %q, want %q", registerReq.Transport, types.TransportTCP)
159 + }
160 if registerReq.Name != "demo-app" {
161 t.Fatalf("register request Name = %q, want %q", registerReq.Name, "demo-app")
162 }
@@ -356,6 +359,152 @@ func TestExposeNoRelayInputs(t *testing.T) {
359 }
360 }
361
362 +func TestNewListenerTransportUDPDoesNotOpenReverseSessions(t *testing.T) {
363 + var connectCount atomic.Int32
364 + server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
365 + switch r.URL.Path {
366 + case types.PathSDKDomain:
367 + writeSDKTestEnvelope(w, http.StatusOK, types.APIEnvelope[types.DomainResponse]{
368 + OK: true,
369 + Data: types.DomainResponse{
370 + Version: types.SDKProtocolVersion,
371 + },
372 + })
373 + case types.PathSDKRegister:
374 + writeSDKTestEnvelope(w, http.StatusCreated, types.APIEnvelope[types.RegisterResponse]{
375 + OK: true,
376 + Data: types.RegisterResponse{
377 + LeaseID: "lease-udp",
378 + Hostname: "demo.example.com",
379 + UDPAddr: "demo.example.com:29000",
380 + QUICAddr: "demo.example.com:4017",
381 + Transport: types.TransportUDP,
382 + },
383 + })
384 + case types.PathSDKConnect:
385 + connectCount.Add(1)
386 + writeSDKTestEnvelope(w, http.StatusForbidden, types.APIEnvelope[any]{
387 + OK: false,
388 + Error: &types.APIError{Code: types.APIErrorCodeUnauthorized, Message: "stream should be disabled"},
389 + })
390 + case types.PathSDKRenew:
391 + writeSDKTestEnvelope(w, http.StatusOK, types.APIEnvelope[types.RenewResponse]{
392 + OK: true,
393 + Data: types.RenewResponse{LeaseID: "lease-udp"},
394 + })
395 + case types.PathSDKUnregister:
396 + writeSDKTestEnvelope(w, http.StatusOK, types.APIEnvelope[any]{OK: true})
397 + default:
398 + http.NotFound(w, r)
399 + }
400 + }))
401 + defer server.Close()
402 +
403 + listener, err := NewListener(context.Background(), server.URL, ListenerConfig{
404 + Name: "demo",
405 + Transport: types.TransportUDP,
406 + LeaseTTL: 100 * time.Millisecond,
407 + })
408 + if err != nil {
409 + t.Fatalf("NewListener() error = %v", err)
410 + }
411 + defer listener.Close()
412 +
413 + waitForSDKTest(t, func() bool {
414 + return listener.LeaseID() == "lease-udp"
415 + })
416 + time.Sleep(150 * time.Millisecond)
417 +
418 + if connectCount.Load() != 0 {
419 + t.Fatalf("connect count = %d, want 0", connectCount.Load())
420 + }
421 + if got := listener.PublicURL(); got != "" {
422 + t.Fatalf("PublicURL() = %q, want empty", got)
423 + }
424 + if !listener.SupportsDatagram() {
425 + t.Fatal("SupportsDatagram() = false, want true")
426 + }
427 + if listener.SupportsStream() {
428 + t.Fatal("SupportsStream() = true, want false")
429 + }
430 +}
431 +
432 +func TestListenerWaitDatagramReadyPublishesRelayAddresses(t *testing.T) {
433 + server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
434 + switch r.URL.Path {
435 + case types.PathSDKDomain:
436 + writeSDKTestEnvelope(w, http.StatusOK, types.APIEnvelope[types.DomainResponse]{
437 + OK: true,
438 + Data: types.DomainResponse{
439 + Version: types.SDKProtocolVersion,
440 + },
441 + })
442 + case types.PathSDKRegister:
443 + writeSDKTestEnvelope(w, http.StatusCreated, types.APIEnvelope[types.RegisterResponse]{
444 + OK: true,
445 + Data: types.RegisterResponse{
446 + LeaseID: "lease-udp",
447 + Hostname: "demo.example.com",
448 + UDPAddr: "demo.example.com:29000",
449 + QUICAddr: "demo.example.com:4017",
450 + Transport: types.TransportUDP,
451 + },
452 + })
453 + case types.PathSDKRenew:
454 + writeSDKTestEnvelope(w, http.StatusOK, types.APIEnvelope[types.RenewResponse]{
455 + OK: true,
456 + Data: types.RenewResponse{LeaseID: "lease-udp"},
457 + })
458 + case types.PathSDKUnregister:
459 + writeSDKTestEnvelope(w, http.StatusOK, types.APIEnvelope[any]{OK: true})
460 + default:
461 + http.NotFound(w, r)
462 + }
463 + }))
464 + defer server.Close()
465 +
466 + listener, err := NewListener(context.Background(), server.URL, ListenerConfig{
467 + Name: "demo",
468 + Transport: types.TransportUDP,
469 + })
470 + if err != nil {
471 + t.Fatalf("NewListener() error = %v", err)
472 + }
473 + defer listener.Close()
474 +
475 + if err := listener.WaitDatagramReady(context.Background()); err != nil {
476 + t.Fatalf("WaitDatagramReady() error = %v", err)
477 + }
478 + if got := listener.UDPAddr(); got != "demo.example.com:29000" {
479 + t.Fatalf("UDPAddr() = %q, want %q", got, "demo.example.com:29000")
480 + }
481 + if got := listener.QUICAddr(); got != "demo.example.com:4017" {
482 + t.Fatalf("QUICAddr() = %q, want %q", got, "demo.example.com:4017")
483 + }
484 +}
485 +
486 +func TestExposureDatagramReply(t *testing.T) {
487 + called := false
488 + dg := ExposureDatagram{
489 + FlowID: 7,
490 + Payload: []byte("hello"),
491 + reply: func(payload []byte) error {
492 + called = true
493 + if string(payload) != "world" {
494 + t.Fatalf("reply payload = %q, want %q", payload, "world")
495 + }
496 + return nil
497 + },
498 + }
499 +
500 + if err := dg.Reply([]byte("world")); err != nil {
501 + t.Fatalf("Reply() error = %v", err)
502 + }
503 + if !called {
504 + t.Fatal("Reply() did not invoke reply function")
505 + }
506 +}
507 +
508 func writeSDKTestEnvelope[T any](w http.ResponseWriter, status int, envelope types.APIEnvelope[T]) {
509 w.Header().Set("Content-Type", "application/json")
510 w.WriteHeader(status)
sdk/udp_listener.go deleted
-371
@@ -1,371 +0,0 @@
1 -package sdk
2 -
3 -import (
4 - "context"
5 - "errors"
6 - "fmt"
7 - "net"
8 - "strings"
9 - "sync"
10 - "time"
11 -
12 - "github.com/quic-go/quic-go"
13 - "github.com/rs/zerolog/log"
14 -
15 - "github.com/gosuda/portal/v2/types"
16 - "github.com/gosuda/portal/v2/utils"
17 -)
18 -
19 -// UDPListenerConfig configures a standalone UDP listener that registers its own lease.
20 -type UDPListenerConfig struct {
21 - Name string
22 - ReverseToken string
23 - Metadata types.LeaseMetadata
24 - Transport string // "udp" or "both", defaults to "udp"
25 - LeaseTTL time.Duration
26 - RootCAPEM []byte
27 - DialTimeout time.Duration
28 - RequestTimeout time.Duration
29 -}
30 -
31 -// UDPListener manages a QUIC connection to the relay for a UDP-transport lease.
32 -// It receives DATAGRAM frames from the relay and delivers decoded datagrams via
33 -// the AcceptDatagram method.
34 -type UDPListener struct {
35 - api *apiClient
36 - baseContext func() context.Context
37 - ctxDone <-chan struct{}
38 - cancel context.CancelFunc
39 -
40 - name string
41 - leaseID string
42 - reverseToken string
43 - udpAddr string
44 - quicAddr string
45 - hostname string
46 - metadata types.LeaseMetadata
47 - leaseTTL time.Duration
48 -
49 - conn *quic.Conn
50 - datagrams chan UDPDatagram
51 - done chan struct{}
52 -
53 - ownsLease bool
54 - closeOnce sync.Once
55 - mu sync.Mutex
56 -}
57 -
58 -// UDPDatagram represents a single datagram received from a public client
59 -// through the relay.
60 -type UDPDatagram struct {
61 - FlowID uint32
62 - Payload []byte
63 -}
64 -
65 -// AcceptDatagram blocks until a datagram is available or the listener is closed.
66 -func (l *UDPListener) AcceptDatagram() (UDPDatagram, error) {
67 - select {
68 - case <-l.ctxDone:
69 - return UDPDatagram{}, net.ErrClosed
70 - case dg, ok := <-l.datagrams:
71 - if !ok {
72 - return UDPDatagram{}, net.ErrClosed
73 - }
74 - return dg, nil
75 - }
76 -}
77 -
78 -// SendDatagram sends a response datagram back to a client via the relay.
79 -func (l *UDPListener) SendDatagram(flowID uint32, payload []byte) error {
80 - l.mu.Lock()
81 - conn := l.conn
82 - l.mu.Unlock()
83 -
84 - if conn == nil {
85 - return errors.New("quic connection not established")
86 - }
87 - return conn.SendDatagram(types.EncodeDatagram(flowID, payload))
88 -}
89 -
90 -// UDPAddr returns the public UDP address allocated by the relay.
91 -func (l *UDPListener) UDPAddr() string {
92 - return l.udpAddr
93 -}
94 -
95 -// LeaseID returns the current lease ID.
96 -func (l *UDPListener) LeaseID() string {
97 - l.mu.Lock()
98 - defer l.mu.Unlock()
99 - return l.leaseID
100 -}
101 -
102 -// Hostname returns the hostname registered for this lease.
103 -func (l *UDPListener) Hostname() string {
104 - l.mu.Lock()
105 - defer l.mu.Unlock()
106 - return l.hostname
107 -}
108 -
109 -// Close tears down the QUIC connection. If this listener owns the lease
110 -// (created via NewUDPListener), it also unregisters the lease. Attached listeners
111 -// (created via Listener.AttachUDP) leave lease lifecycle to the TCP Listener.
112 -func (l *UDPListener) Close() error {
113 - var closeErr error
114 - l.closeOnce.Do(func() {
115 - l.cancel()
116 -
117 - l.mu.Lock()
118 - conn := l.conn
119 - leaseID := l.leaseID
120 - l.mu.Unlock()
121 -
122 - if conn != nil {
123 - _ = conn.CloseWithError(0, "listener closed")
124 - }
125 -
126 - if l.ownsLease && l.api != nil && leaseID != "" {
127 - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
128 - defer cancel()
129 - closeErr = l.api.unregisterLease(ctx, leaseID)
130 - }
131 - })
132 - return closeErr
133 -}
134 -
135 -// NewUDPListener registers a UDP-transport lease and returns a UDPListener.
136 -func NewUDPListener(ctx context.Context, relayURL string, cfg UDPListenerConfig) (*UDPListener, error) {
137 - if strings.TrimSpace(cfg.Name) == "" {
138 - return nil, errors.New("listener name is required")
139 - }
140 - if ctx == nil {
141 - ctx = context.Background()
142 - }
143 -
144 - reverseToken := strings.TrimSpace(cfg.ReverseToken)
145 - if reverseToken == "" {
146 - reverseToken = utils.RandomID("tok_")
147 - }
148 - leaseTTL := utils.DurationOrDefault(cfg.LeaseTTL, defaultLeaseTTL)
149 -
150 - transport := strings.ToLower(strings.TrimSpace(cfg.Transport))
151 - if transport == "" {
152 - transport = types.TransportUDP
153 - }
154 -
155 - api, err := newApiClient(relayURL, ListenerConfig{
156 - Name: cfg.Name,
157 - ReverseToken: reverseToken,
158 - Metadata: cfg.Metadata,
159 - RootCAPEM: cfg.RootCAPEM,
160 - DialTimeout: cfg.DialTimeout,
161 - RequestTimeout: cfg.RequestTimeout,
162 - })
163 - if err != nil {
164 - return nil, err
165 - }
166 -
167 - if err := api.ensureReady(ctx); err != nil {
168 - api.close()
169 - return nil, err
170 - }
171 -
172 - registerResp, err := api.registerLease(ctx, leaseTTL, transport)
173 - if err != nil {
174 - api.close()
175 - return nil, err
176 - }
177 -
178 - listenerCtx, cancel := context.WithCancel(ctx)
179 - listener := &UDPListener{
180 - api: api,
181 - baseContext: func() context.Context { return listenerCtx },
182 - ctxDone: listenerCtx.Done(),
183 - cancel: cancel,
184 - name: strings.TrimSpace(cfg.Name),
185 - leaseID: registerResp.LeaseID,
186 - reverseToken: reverseToken,
187 - udpAddr: registerResp.UDPAddr,
188 - quicAddr: registerResp.QUICAddr,
189 - hostname: registerResp.Hostname,
190 - metadata: registerResp.Metadata,
191 - leaseTTL: leaseTTL,
192 - datagrams: make(chan UDPDatagram, 256),
193 - done: make(chan struct{}),
194 - ownsLease: true,
195 - }
196 -
197 - go listener.runSupervisor()
198 - go listener.runRenewLoop()
199 -
200 - return listener, nil
201 -}
202 -
203 -// AttachUDP creates a UDPListener that connects to an existing Listener's
204 -// QUIC broker without registering a new lease or running a renew loop.
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
216 - udpAddr := l.udpAddr
217 - quicAddr := l.quicAddr
218 - l.mu.Unlock()
219 -
220 - if leaseID == "" {
221 - return nil, errors.New("lease not registered yet")
222 - }
223 - if api == nil {
224 - return nil, errors.New("api client not available")
225 - }
226 - if udpAddr == "" && quicAddr == "" {
227 - return nil, errors.New("lease does not have UDP transport enabled")
228 - }
229 -
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 },
237 - ctxDone: listenerCtx.Done(),
238 - cancel: cancel,
239 - leaseID: leaseID,
240 - reverseToken: api.reverseToken,
241 - udpAddr: udpAddr,
242 - quicAddr: quicAddr,
243 - datagrams: make(chan UDPDatagram, 256),
244 - done: make(chan struct{}),
245 - ownsLease: false,
246 - }
247 -
248 - go udpL.runSupervisor()
249 - return udpL, nil
250 -}
251 -
252 -func (l *UDPListener) runSupervisor() {
253 - for {
254 - select {
255 - case <-l.ctxDone:
256 - return
257 - default:
258 - }
259 -
260 - conn, err := l.api.openQUICSession(l.context(), l.quicAddr, l.leaseID, l.reverseToken)
261 - if err != nil {
262 - log.Warn().
263 - Err(err).
264 - Str("component", "sdk-udp-listener").
265 - Str("lease_id", l.leaseID).
266 - Msg("quic session open failed, retrying")
267 - utils.SleepOrDone(l.context(), 2*time.Second)
268 - continue
269 - }
270 -
271 - l.mu.Lock()
272 - l.conn = conn
273 - l.mu.Unlock()
274 -
275 - log.Info().
276 - Str("component", "sdk-udp-listener").
277 - Str("lease_id", l.leaseID).
278 - Str("remote_addr", conn.RemoteAddr().String()).
279 - Msg("quic tunnel connected")
280 -
281 - l.receiveLoop(conn)
282 -
283 - l.mu.Lock()
284 - if l.conn == conn {
285 - l.conn = nil
286 - }
287 - l.mu.Unlock()
288 -
289 - if l.isClosed() {
290 - return
291 - }
292 - utils.SleepOrDone(l.context(), time.Second)
293 - }
294 -}
295 -
296 -func (l *UDPListener) receiveLoop(conn *quic.Conn) {
297 - for {
298 - data, err := conn.ReceiveDatagram(l.context())
299 - if err != nil {
300 - if !l.isClosed() {
301 - log.Warn().
302 - Err(err).
303 - Str("component", "sdk-udp-listener").
304 - Str("lease_id", l.leaseID).
305 - Msg("quic receive loop ended")
306 - }
307 - return
308 - }
309 -
310 - frame, err := types.DecodeDatagram(data)
311 - if err != nil {
312 - continue
313 - }
314 -
315 - select {
316 - case l.datagrams <- UDPDatagram{FlowID: frame.FlowID, Payload: frame.Payload}:
317 - case <-l.ctxDone:
318 - return
319 - }
320 - }
321 -}
322 -
323 -func (l *UDPListener) runRenewLoop() {
324 - interval := l.leaseTTL / 2
325 - if interval <= 0 {
326 - interval = 30 * time.Second
327 - }
328 -
329 - ticker := time.NewTicker(interval)
330 - defer ticker.Stop()
331 -
332 - for {
333 - select {
334 - case <-l.ctxDone:
335 - return
336 - case <-ticker.C:
337 - ctx, cancel := context.WithTimeout(l.context(), 10*time.Second)
338 - err := l.api.renewLease(ctx, l.leaseID, l.leaseTTL)
339 - cancel()
340 - if err != nil {
341 - log.Warn().
342 - Err(err).
343 - Str("component", "sdk-udp-listener").
344 - Str("lease_id", l.leaseID).
345 - Msg("lease renewal failed")
346 - }
347 - }
348 - }
349 -}
350 -
351 -func (l *UDPListener) context() context.Context {
352 - if l.baseContext != nil {
353 - if ctx := l.baseContext(); ctx != nil {
354 - return ctx
355 - }
356 - }
357 - return context.Background()
358 -}
359 -
360 -func (l *UDPListener) isClosed() bool {
361 - if l.ctxDone == nil {
362 - return false
363 - }
364 - select {
365 - case <-l.ctxDone:
366 - return true
367 - default:
368 - return false
369 - }
370 -}
371 -
types/api.go
-4
@@ -6,10 +6,6 @@ import (
6 "time"
7 )
8
9 -const (
10 - MarkerQUICReady = byte(0x03)
11 -)
12 -
9 const (
10 TransportTCP = "tcp"
11 TransportUDP = "udp"
types/datagram.go deleted
-39
@@ -1,39 +0,0 @@
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 -}
types/paths.go
+1 -2
@@ -26,6 +26,5 @@ const (
26 PathSDKRegister = "/sdk/register"
27 PathSDKRenew = "/sdk/renew"
28 PathSDKUnregister = "/sdk/unregister"
29 - PathSDKConnect = "/sdk/connect"
30 - PathSDKQUICConnect = "/sdk/quic-connect"
29 + PathSDKConnect = "/sdk/connect"
30 )
types/transport.go new
+85
@@ -0,0 +1,85 @@
1 +package types
2 +
3 +import (
4 + "encoding/binary"
5 + "errors"
6 + "fmt"
7 + "strings"
8 +)
9 +
10 +// LeaseCapabilities describes which data planes a lease exposes.
11 +// Stream maps to reverse TCP/TLS sessions; Datagram maps to QUIC/UDP.
12 +type LeaseCapabilities struct {
13 + Datagram bool
14 + Stream bool
15 +}
16 +
17 +// ParseLeaseCapabilities normalizes the public transport string into the
18 +// internal capability model shared by relay and SDK.
19 +func ParseLeaseCapabilities(raw string) (LeaseCapabilities, error) {
20 + switch strings.ToLower(strings.TrimSpace(raw)) {
21 + case "", TransportTCP:
22 + return LeaseCapabilities{Stream: true}, nil
23 + case TransportUDP:
24 + return LeaseCapabilities{Datagram: true}, nil
25 + case TransportBoth:
26 + return LeaseCapabilities{Datagram: true, Stream: true}, nil
27 + default:
28 + return LeaseCapabilities{}, fmt.Errorf("unsupported transport %q", strings.TrimSpace(raw))
29 + }
30 +}
31 +
32 +func (c LeaseCapabilities) SupportsDatagram() bool {
33 + return c.Datagram
34 +}
35 +
36 +func (c LeaseCapabilities) SupportsStream() bool {
37 + return c.Stream
38 +}
39 +
40 +// Transport returns the canonical public transport label for the capability set.
41 +func (c LeaseCapabilities) Transport() string {
42 + switch {
43 + case c.Stream && c.Datagram:
44 + return TransportBoth
45 + case c.Datagram:
46 + return TransportUDP
47 + case c.Stream:
48 + return TransportTCP
49 + default:
50 + return ""
51 + }
52 +}
53 +
54 +// ErrDatagramTooSmall is returned when a datagram payload is too short to
55 +// contain a valid flow ID varint.
56 +var ErrDatagramTooSmall = errors.New("datagram too small to decode")
57 +
58 +// DatagramFrame is the wire format for QUIC DATAGRAM payloads.
59 +// Layout: [flowID varint][payload bytes]
60 +type DatagramFrame struct {
61 + FlowID uint32
62 + Payload []byte
63 +}
64 +
65 +// EncodeDatagram serialises a flow-framed datagram for transmission.
66 +func EncodeDatagram(flowID uint32, payload []byte) []byte {
67 + var buf [binary.MaxVarintLen32]byte
68 + n := binary.PutUvarint(buf[:], uint64(flowID))
69 + out := make([]byte, n+len(payload))
70 + copy(out, buf[:n])
71 + copy(out[n:], payload)
72 + return out
73 +}
74 +
75 +// DecodeDatagram deserialises a flow-framed datagram.
76 +func DecodeDatagram(data []byte) (DatagramFrame, error) {
77 + flowID, n := binary.Uvarint(data)
78 + if n <= 0 {
79 + return DatagramFrame{}, ErrDatagramTooSmall
80 + }
81 + return DatagramFrame{
82 + FlowID: uint32(flowID),
83 + Payload: data[n:],
84 + }, nil
85 +}