diff --git a/pkg/service/webtransport.go b/pkg/service/webtransport.go index 0a7431939..2d0ad276d 100644 --- a/pkg/service/webtransport.go +++ b/pkg/service/webtransport.go @@ -27,8 +27,10 @@ import ( "net" "net/http" "strconv" + "sync" "time" + "github.com/quic-go/quic-go" "github.com/quic-go/quic-go/http3" "github.com/quic-go/webtransport-go" "github.com/urfave/negroni/v3" @@ -139,9 +141,26 @@ func ListenWebTransport(wt *webtransport.Server, addrs []string, port uint32) ([ addrs = []string{""} } - conns := make([]*net.UDPConn, 0, len(addrs)) - bound := make([]net.Addr, 0, len(addrs)) + // webtransport.Server.Serve takes a reference on a WaitGroup that + // Server.Close waits on, so running the two concurrently is a race by the + // WaitGroup's own rules. Accepting here instead leaves that counter at zero: + // Server.ServeQUICConn never touches it. + quicConf := &quic.Config{} + if wt.H3.QUICConfig != nil { + quicConf = wt.H3.QUICConfig.Clone() + } + quicConf.EnableDatagrams = true + quicConf.EnableStreamResetPartialDelivery = true + + var ( + conns []*net.UDPConn + lns []*quic.EarlyListener + bound []net.Addr + ) closeAll := func() { + for _, ln := range lns { + _ = ln.Close() + } for _, c := range conns { _ = c.Close() } @@ -159,19 +178,43 @@ func ListenWebTransport(wt *webtransport.Server, addrs []string, port uint32) ([ } conns = append(conns, udp) bound = append(bound, udp.LocalAddr()) + + ln, err := quic.ListenEarly(udp, wt.H3.TLSConfig, quicConf) + if err != nil { + closeAll() + return nil, nil, err + } + lns = append(lns, ln) } - for _, udp := range conns { - go func() { - if err := wt.Serve(udp); err != nil { - logger.Infow("webtransport listener stopped", "error", err) - } - }() + ctx, cancel := context.WithCancel(context.Background()) + var serving sync.WaitGroup + for _, ln := range lns { + serving.Go(func() { acceptWebTransport(ctx, wt, ln, &serving) }) } logger.Infow("webtransport listener started", "addresses", bound) return bound, func() { + // the sockets stay open until everything has drained, so every + // CONNECTION_CLOSE frame still reaches its peer _ = wt.Close() + cancel() + serving.Wait() closeAll() }, nil } + +func acceptWebTransport(ctx context.Context, wt *webtransport.Server, ln *quic.EarlyListener, serving *sync.WaitGroup) { + for { + conn, err := ln.Accept(ctx) + if err != nil { + logger.Infow("webtransport listener stopped", "error", err) + return + } + serving.Go(func() { + if err := wt.ServeQUICConn(conn); err != nil && !errors.Is(err, http.ErrServerClosed) { + logger.Infow("webtransport connection stopped", "error", err) + } + }) + } +} diff --git a/pkg/service/webtransportlisten_test.go b/pkg/service/webtransportlisten_test.go new file mode 100644 index 000000000..fb0654707 --- /dev/null +++ b/pkg/service/webtransportlisten_test.go @@ -0,0 +1,37 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package service_test + +import ( + "net/http" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/livekit/livekit-server/pkg/service" +) + +// a listener stopped before it has served anything must not race its own accept +// loop: under -race this is what catches webtransport.Server.Serve running +// concurrently with Server.Close. +func TestWebTransportStopBeforeFirstConnection(t *testing.T) { + for range 20 { + wt := service.NewWebTransportServer(selfSignedTLS(t)) + wt.H3.Handler = service.NewWebTransportHandler(nil, wt, http.NewServeMux()) + _, stop, err := service.ListenWebTransport(wt, []string{"127.0.0.1"}, 0) + require.NoError(t, err) + stop() + } +}