Use new quic api

matrix-org · Aug 10, 2023 · ea4c337 · ea4c337
1 parent 0b58977
commit ea4c337
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 6 deletions.
diff --git a/sessions/dial.go b/sessions/dial.go
@@ -25,7 +25,6 @@ import (
 	"net"
 
 	"github.com/matrix-org/pinecone/types"
-	"github.com/quic-go/quic-go"
 )
 
 // DialContext dials a given public key using the supplied network.
@@ -87,7 +86,7 @@ retry:
 			},
 		}
 
-		session.Connection, err = quic.Dial(ctx, s.s.r, addr, tlsConfig, s.s.quicConfig)
+		session.Connection, err = s.transport.Dial(ctx, addr, tlsConfig, s.s.quicConfig)
 		session.Unlock()
 		if err != nil {
 			if err == context.DeadlineExceeded {

diff --git a/sessions/sessions.go b/sessions/sessions.go
@@ -48,6 +48,7 @@ type Sessions struct {
 
 type SessionProtocol struct {
 	s         *Sessions
+	transport *quic.Transport
 	proto     string
 	streams   chan net.Conn
 	sessions  sync.Map // types.PublicKey -> *activeSession
@@ -60,6 +61,7 @@ type activeSession struct {
 }
 
 func NewSessions(log types.Logger, r *router.Router, protos []string) *Sessions {
+	transport := &quic.Transport{Conn: r}
 	ctx, cancel := context.WithCancel(context.Background())
 	s := &Sessions{
 		r:         r,
@@ -74,9 +76,10 @@ func NewSessions(log types.Logger, r *router.Router, protos []string) *Sessions
 	}
 	for _, proto := range protos {
 		s.protocols[proto] = &SessionProtocol{
-			s:       s,
-			proto:   proto,
-			streams: make(chan net.Conn, 1),
+			s:         s,
+			transport: transport,
+			proto:     proto,
+			streams:   make(chan net.Conn, 1),
 		}
 	}
 
@@ -87,7 +90,7 @@ func NewSessions(log types.Logger, r *router.Router, protos []string) *Sessions
 		NextProtos:   protos,
 	}
 
-	listener, err := quic.Listen(r, s.tlsServerCfg, s.quicConfig)
+	listener, err := transport.Listen(s.tlsServerCfg, s.quicConfig)
 	if err != nil {
 		panic(fmt.Errorf("quic.NewSocketFromPacketConnNoClose: %w", err))
 	}