Merge pull request #366 from SiaFoundation/chris/cancel-read

Add final Read when closing QUIC stream + reproduction test

Author: n8mgr

.changeset/fix_quic_streams_not_getting_closed_properly.md

@@ -0,0 +1,5 @@
+---
+default: patch
+---
+
+# Fix QUIC streams not getting closed properly.

rhp/v4/quic/quic.go

@@ -5,6 +5,7 @@ import (
 	"crypto/tls"
 	"errors"
 	"fmt"
+	"io"
 	"net"
 	"net/http"
 	"time"
@@ -25,6 +26,10 @@ const (
 
 	// TLSNextProtoRHP4 is the ALPN identifier for the Quic RHP4 protocol.
 	TLSNextProtoRHP4 = "sia/rhp4"
+
+	// maxIncomingStreams is the maximum number of incoming streams allowed per
+	// QUIC connection by the server.
+	maxIncomingStreams = 100000
 )
 
 type (
@@ -60,6 +65,16 @@ func WithTLSConfig(fn func(*tls.Config)) ClientOption {
 	}
 }
 
+func (s *stream) Close() error {
+	err := s.Stream.Close()
+	_, errCopy := io.CopyN(io.Discard, s, 4096)
+	if !errors.Is(errCopy, io.EOF) {
+		// fall back to forcefully canceling the read if we couldn't reach EOF
+		s.CancelRead(1)
+	}
+	return err
+}
+
 // LocalAddr implements net.Conn
 func (s *stream) LocalAddr() net.Addr {
 	return s.localAddr
@@ -192,7 +207,7 @@ func Listen(conn net.PacketConn, certs CertManager) (*quic.Listener, error) {
 		EnableDatagrams:    true,
 		KeepAlivePeriod:    30 * time.Second,
 		MaxIdleTimeout:     30 * time.Minute,
-		MaxIncomingStreams: 1000,
+		MaxIncomingStreams: maxIncomingStreams,
 	})
 }
 

rhp/v4/quic/quic_test.go

@@ -0,0 +1,92 @@
+package quic
+
+import (
+	"context"
+	"crypto/tls"
+	"net"
+	"sync"
+	"testing"
+
+	"github.com/quic-go/quic-go"
+	"go.sia.tech/core/types"
+	"go.sia.tech/coreutils/rhp/v4"
+	"go.sia.tech/coreutils/testutil/certs"
+)
+
+// setupTestPair sets up a QUIC server and client for testing.
+func setupTestPair(tb testing.TB) (*quic.Listener, rhp.TransportClient) {
+	tb.Helper()
+
+	udpAddr, err := net.ResolveUDPAddr("udp", "localhost:0")
+	if err != nil {
+		tb.Fatal(err)
+	}
+
+	conn, err := net.ListenUDP("udp", udpAddr)
+	if err != nil {
+		tb.Fatal(err)
+	}
+	tb.Cleanup(func() { conn.Close() })
+
+	l, err := Listen(conn, &certs.EphemeralCertManager{})
+	if err != nil {
+		tb.Fatal(err)
+	}
+	tb.Cleanup(func() { l.Close() })
+
+	client, err := Dial(context.Background(), conn.LocalAddr().String(), types.PublicKey{}, WithTLSConfig(func(tc *tls.Config) {
+		tc.InsecureSkipVerify = true
+	}))
+	if err != nil {
+		tb.Fatal(err)
+	}
+	tb.Cleanup(func() { client.Close() })
+	return l, client
+}
+
+// TestStreamLimit is a regression test that ensures we can perform more than
+// maxIncomingStreams RPCs on a single QUIC connection. This should be the case
+// if both sides close their streams.
+// Before fixing the issue, this test would hang for the full 10 minute test
+// timeout.
+func TestStreamLimit(t *testing.T) {
+	server, client := setupTestPair(t)
+
+	// the server just accepts streams and cancels them
+	go func() {
+		for {
+			conn, err := server.Accept(context.Background())
+			if err != nil {
+				return
+			}
+
+			transport := &transport{conn}
+
+			for {
+				stream, err := transport.AcceptStream()
+				if err != nil {
+					return
+				}
+				stream.Close()
+			}
+		}
+	}()
+
+	// open the maximum number of streams + 1 which should neither block forever
+	// nor return an error
+	var wg sync.WaitGroup
+	for range maxIncomingStreams + 1 {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+
+			stream, err := client.DialStream(context.Background())
+			if err != nil {
+				t.Error(err)
+				return
+			}
+			stream.Close()
+		}()
+	}
+	wg.Wait()
+}

testutil/certs/certs.go

@@ -0,0 +1,37 @@
+package certs
+
+import (
+	"crypto/rsa"
+	"crypto/tls"
+	"crypto/x509"
+	"encoding/pem"
+	"fmt"
+	"math/big"
+
+	"lukechampine.com/frand"
+)
+
+// An EphemeralCertManager is an in-memory minimal rhp4.CertManager for testing.
+// Calls to GetCertificate will return a new self-signed certificate each time.
+type EphemeralCertManager struct{}
+
+// GetCertificate returns a new self-signed certificate each time it is called.
+func (ec *EphemeralCertManager) GetCertificate(*tls.ClientHelloInfo) (*tls.Certificate, error) {
+	key, err := rsa.GenerateKey(frand.Reader, 2048)
+	if err != nil {
+		return nil, fmt.Errorf("failed to generate key: %w", err)
+	}
+	template := x509.Certificate{SerialNumber: big.NewInt(1)}
+	certDER, err := x509.CreateCertificate(frand.Reader, &template, &template, &key.PublicKey, key)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create cert: %w", err)
+	}
+	keyPEM := pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(key)})
+	certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: certDER})
+
+	cert, err := tls.X509KeyPair(certPEM, keyPEM)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create tls cert: %w", err)
+	}
+	return &cert, nil
+}

testutil/host.go

@@ -1,13 +1,7 @@
 package testutil
 
 import (
-	"crypto/rsa"
-	"crypto/tls"
-	"crypto/x509"
-	"encoding/pem"
 	"errors"
-	"fmt"
-	"math/big"
 	"net"
 	"sync"
 	"testing"
@@ -19,35 +13,10 @@ import (
 	rhp4 "go.sia.tech/coreutils/rhp/v4"
 	"go.sia.tech/coreutils/rhp/v4/quic"
 	"go.sia.tech/coreutils/rhp/v4/siamux"
+	"go.sia.tech/coreutils/testutil/certs"
 	"go.uber.org/zap"
-	"lukechampine.com/frand"
 )
 
-// An EphemeralCertManager is an in-memory minimal rhp4.CertManager for testing.
-// Calls to GetCertificate will return a new self-signed certificate each time.
-type EphemeralCertManager struct{}
-
-// GetCertificate returns a new self-signed certificate each time it is called.
-func (ec *EphemeralCertManager) GetCertificate(*tls.ClientHelloInfo) (*tls.Certificate, error) {
-	key, err := rsa.GenerateKey(frand.Reader, 2048)
-	if err != nil {
-		return nil, fmt.Errorf("failed to generate key: %w", err)
-	}
-	template := x509.Certificate{SerialNumber: big.NewInt(1)}
-	certDER, err := x509.CreateCertificate(frand.Reader, &template, &template, &key.PublicKey, key)
-	if err != nil {
-		return nil, fmt.Errorf("failed to create cert: %w", err)
-	}
-	keyPEM := pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(key)})
-	certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: certDER})
-
-	cert, err := tls.X509KeyPair(certPEM, keyPEM)
-	if err != nil {
-		return nil, fmt.Errorf("failed to create tls cert: %w", err)
-	}
-	return &cert, nil
-}
-
 // An EphemeralSectorStore is an in-memory minimal rhp4.SectorStore for testing.
 type EphemeralSectorStore struct {
 	mu      sync.Mutex
@@ -438,7 +407,7 @@ func ServeQUIC(tb testing.TB, s *rhp4.Server, log *zap.Logger) string {
 	}
 	tb.Cleanup(func() { conn.Close() })
 
-	l, err := quic.Listen(conn, &EphemeralCertManager{})
+	l, err := quic.Listen(conn, &certs.EphemeralCertManager{})
 	if err != nil {
 		tb.Fatal(err)
 	}