@samitouri / QOSAMI-X402Tunnel / 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

@@ -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

3 import (

4 - "context"

4 "crypto/tls"

5 "errors"

6 "fmt"

@@ -13,6 +12,7 @@ import (

13 "github.com/rs/zerolog/log"

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 )

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 }

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 }

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 }

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 }

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 })

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()

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()

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()

@@ -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

3 import (

4 "crypto"

@@ -7,15 +7,10 @@ import (

7 "encoding/binary"

8 "errors"

9 "fmt"

10 - "net"

11 - "sync"

11 "golang.org/x/crypto/hkdf"

14 -

15 - "github.com/gosuda/portal/v2/utils"

12 )

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")

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 }

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 }

43 - // Packet type: bits 4-5 of first byte. Initial = 0.

33 packetType := (firstByte & 0x30) >> 4

34 if packetType != 0 {

35 return "", errNotQUICInitial

36 }

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 }

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

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 }

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 }

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

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"

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 )

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()

@@ -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 }

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

3 import (

4 "context"

5 + "encoding/json"

6 "errors"

7 "fmt"

8 "io"

@@ -18,6 +19,7 @@ import (

19 "golang.org/x/sync/errgroup"

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

40 defaultUDPPortMin = 29000

41 defaultUDPPortMax = 29999

42 )

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 }

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 }

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{}

29 closeOnce sync.Once

30 connSeq atomic.Uint64

31 }

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 }

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 (

15 "github.com/rs/zerolog/log"

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 )

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

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 )

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 +}