From c9b5e7f77b8aac7334d81c552c583af81ba1c400 Mon Sep 17 00:00:00 2001
From: Artem Mikheev <30644072+renbou@users.noreply.github.com>
Date: Wed, 23 Mar 2022 02:47:57 +0300
Subject: Fix http3 servers dying after reload (#4654)

---
 listeners.go | 81 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 76 insertions(+), 5 deletions(-)

(limited to 'listeners.go')

diff --git a/listeners.go b/listeners.go
index c23ff78..181db75 100644
--- a/listeners.go
+++ b/listeners.go
@@ -15,6 +15,9 @@
 package caddy
 
 import (
+	"context"
+	"crypto/tls"
+	"errors"
 	"fmt"
 	"net"
 	"os"
@@ -24,6 +27,9 @@ import (
 	"sync/atomic"
 	"syscall"
 	"time"
+
+	"github.com/lucas-clemente/quic-go"
+	"github.com/lucas-clemente/quic-go/http3"
 )
 
 // Listen is like net.Listen, except Caddy's listeners can overlap
@@ -79,6 +85,27 @@ func ListenPacket(network, addr string) (net.PacketConn, error) {
 	return &fakeClosePacketConn{sharedPacketConn: sharedPc.(*sharedPacketConn)}, nil
 }
 
+// ListenQUIC returns a quic.EarlyListener suitable for use in a Caddy module.
+// Note that the context passed to Accept is currently ignored, so using
+// a context other than context.Background is meaningless.
+func ListenQUIC(addr string, tlsConf *tls.Config) (quic.EarlyListener, error) {
+	lnKey := "quic/" + addr
+
+	sharedEl, _, err := listenerPool.LoadOrNew(lnKey, func() (Destructor, error) {
+		el, err := quic.ListenAddrEarly(addr, http3.ConfigureTLSConfig(tlsConf), &quic.Config{})
+		if err != nil {
+			return nil, err
+		}
+		return &sharedQuicListener{EarlyListener: el, key: lnKey}, nil
+	})
+
+	ctx, cancel := context.WithCancel(context.Background())
+	return &fakeCloseQuicListener{
+		sharedQuicListener: sharedEl.(*sharedQuicListener),
+		context:            ctx, contextCancel: cancel,
+	}, err
+}
+
 // fakeCloseListener is a private wrapper over a listener that
 // is shared. The state of fakeCloseListener is not shared.
 // This allows one user of a socket to "close" the listener
@@ -95,7 +122,7 @@ type fakeCloseListener struct {
 func (fcl *fakeCloseListener) Accept() (net.Conn, error) {
 	// if the listener is already "closed", return error
 	if atomic.LoadInt32(&fcl.closed) == 1 {
-		return nil, fcl.fakeClosedErr()
+		return nil, fakeClosedErr(fcl)
 	}
 
 	// call underlying accept
@@ -119,7 +146,7 @@ func (fcl *fakeCloseListener) Accept() (net.Conn, error) {
 		_ = fcl.sharedListener.clearDeadline()
 
 		if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
-			return nil, fcl.fakeClosedErr()
+			return nil, fakeClosedErr(fcl)
 		}
 	}
 
@@ -145,14 +172,47 @@ func (fcl *fakeCloseListener) Close() error {
 	return nil
 }
 
+type fakeCloseQuicListener struct {
+	closed              int32 // accessed atomically; belongs to this struct only
+	*sharedQuicListener       // embedded, so we also become a quic.EarlyListener
+	context             context.Context
+	contextCancel       context.CancelFunc
+}
+
+// Currently Accept ignores the passed context, however a situation where
+// someone would need a hotswappable QUIC-only (not http3, since it uses context.Background here)
+// server on which Accept would be called with non-empty contexts
+// (mind that the default net listeners' Accept doesn't take a context argument)
+// sounds way too rare for us to sacrifice efficiency here.
+func (fcql *fakeCloseQuicListener) Accept(_ context.Context) (quic.EarlySession, error) {
+	conn, err := fcql.sharedQuicListener.Accept(fcql.context)
+	if err == nil {
+		return conn, nil
+	}
+
+	// if the listener is "closed", return a fake closed error instead
+	if atomic.LoadInt32(&fcql.closed) == 1 && errors.Is(err, context.Canceled) {
+		return nil, fakeClosedErr(fcql)
+	}
+	return nil, err
+}
+
+func (fcql *fakeCloseQuicListener) Close() error {
+	if atomic.CompareAndSwapInt32(&fcql.closed, 0, 1) {
+		fcql.contextCancel()
+		_, _ = listenerPool.Delete(fcql.sharedQuicListener.key)
+	}
+	return nil
+}
+
 // fakeClosedErr returns an error value that is not temporary
 // nor a timeout, suitable for making the caller think the
 // listener is actually closed
-func (fcl *fakeCloseListener) fakeClosedErr() error {
+func fakeClosedErr(l interface{ Addr() net.Addr }) error {
 	return &net.OpError{
 		Op:   "accept",
-		Net:  fcl.Addr().Network(),
-		Addr: fcl.Addr(),
+		Net:  l.Addr().Network(),
+		Addr: l.Addr(),
 		Err:  errFakeClosed,
 	}
 }
@@ -244,6 +304,17 @@ func (sl *sharedListener) Destruct() error {
 	return sl.Listener.Close()
 }
 
+// sharedQuicListener is like sharedListener, but for quic.EarlyListeners.
+type sharedQuicListener struct {
+	quic.EarlyListener
+	key string
+}
+
+// Destruct closes the underlying QUIC listener.
+func (sql *sharedQuicListener) Destruct() error {
+	return sql.EarlyListener.Close()
+}
+
 // sharedPacketConn is like sharedListener, but for net.PacketConns.
 type sharedPacketConn struct {
 	net.PacketConn
-- 
cgit v1.2.3