github.com/tumi8/quic-go@v0.37.4-tum/integrationtests/self/datagram_test.go (about)

     1  package self_test
     2  
     3  import (
     4  	"context"
     5  	"encoding/binary"
     6  	"fmt"
     7  	mrand "math/rand"
     8  	"net"
     9  	"sync"
    10  	"sync/atomic"
    11  	"time"
    12  
    13  	"github.com/tumi8/quic-go"
    14  	quicproxy "github.com/tumi8/quic-go/integrationtests/tools/proxy"
    15  	"github.com/tumi8/quic-go/noninternal/wire"
    16  
    17  	. "github.com/onsi/ginkgo/v2"
    18  	. "github.com/onsi/gomega"
    19  )
    20  
    21  var _ = Describe("Datagram test", func() {
    22  	const num = 100
    23  
    24  	var (
    25  		serverConn, clientConn *net.UDPConn
    26  		dropped, total         int32
    27  	)
    28  
    29  	startServerAndProxy := func(enableDatagram, expectDatagramSupport bool) (port int, closeFn func()) {
    30  		addr, err := net.ResolveUDPAddr("udp", "localhost:0")
    31  		Expect(err).ToNot(HaveOccurred())
    32  		serverConn, err = net.ListenUDP("udp", addr)
    33  		Expect(err).ToNot(HaveOccurred())
    34  		ln, err := quic.Listen(
    35  			serverConn,
    36  			getTLSConfig(),
    37  			getQuicConfig(&quic.Config{EnableDatagrams: enableDatagram}),
    38  		)
    39  		Expect(err).ToNot(HaveOccurred())
    40  
    41  		accepted := make(chan struct{})
    42  		go func() {
    43  			defer GinkgoRecover()
    44  			defer close(accepted)
    45  			conn, err := ln.Accept(context.Background())
    46  			Expect(err).ToNot(HaveOccurred())
    47  
    48  			if expectDatagramSupport {
    49  				Expect(conn.ConnectionState().SupportsDatagrams).To(BeTrue())
    50  
    51  				if enableDatagram {
    52  					var wg sync.WaitGroup
    53  					wg.Add(num)
    54  					for i := 0; i < num; i++ {
    55  						go func(i int) {
    56  							defer GinkgoRecover()
    57  							defer wg.Done()
    58  							b := make([]byte, 8)
    59  							binary.BigEndian.PutUint64(b, uint64(i))
    60  							Expect(conn.SendMessage(b)).To(Succeed())
    61  						}(i)
    62  					}
    63  					wg.Wait()
    64  				}
    65  			} else {
    66  				Expect(conn.ConnectionState().SupportsDatagrams).To(BeFalse())
    67  			}
    68  		}()
    69  
    70  		serverPort := ln.Addr().(*net.UDPAddr).Port
    71  		proxy, err := quicproxy.NewQuicProxy("localhost:0", &quicproxy.Opts{
    72  			RemoteAddr: fmt.Sprintf("localhost:%d", serverPort),
    73  			// drop 10% of Short Header packets sent from the server
    74  			DropPacket: func(dir quicproxy.Direction, packet []byte) bool {
    75  				if dir == quicproxy.DirectionIncoming {
    76  					return false
    77  				}
    78  				// don't drop Long Header packets
    79  				if wire.IsLongHeaderPacket(packet[0]) {
    80  					return false
    81  				}
    82  				drop := mrand.Int()%10 == 0
    83  				if drop {
    84  					atomic.AddInt32(&dropped, 1)
    85  				}
    86  				atomic.AddInt32(&total, 1)
    87  				return drop
    88  			},
    89  		})
    90  		Expect(err).ToNot(HaveOccurred())
    91  		return proxy.LocalPort(), func() {
    92  			Eventually(accepted).Should(BeClosed())
    93  			proxy.Close()
    94  			ln.Close()
    95  		}
    96  	}
    97  
    98  	BeforeEach(func() {
    99  		addr, err := net.ResolveUDPAddr("udp", "localhost:0")
   100  		Expect(err).ToNot(HaveOccurred())
   101  		clientConn, err = net.ListenUDP("udp", addr)
   102  		Expect(err).ToNot(HaveOccurred())
   103  	})
   104  
   105  	It("sends datagrams", func() {
   106  		proxyPort, close := startServerAndProxy(true, true)
   107  		defer close()
   108  		raddr, err := net.ResolveUDPAddr("udp", fmt.Sprintf("localhost:%d", proxyPort))
   109  		Expect(err).ToNot(HaveOccurred())
   110  		conn, err := quic.Dial(
   111  			context.Background(),
   112  			clientConn,
   113  			raddr,
   114  			getTLSClientConfig(),
   115  			getQuicConfig(&quic.Config{EnableDatagrams: true}),
   116  		)
   117  		Expect(err).ToNot(HaveOccurred())
   118  		Expect(conn.ConnectionState().SupportsDatagrams).To(BeTrue())
   119  		var counter int
   120  		for {
   121  			// Close the connection if no message is received for 100 ms.
   122  			timer := time.AfterFunc(scaleDuration(100*time.Millisecond), func() { conn.CloseWithError(0, "") })
   123  			if _, err := conn.ReceiveMessage(context.Background()); err != nil {
   124  				break
   125  			}
   126  			timer.Stop()
   127  			counter++
   128  		}
   129  
   130  		numDropped := int(atomic.LoadInt32(&dropped))
   131  		expVal := num - numDropped
   132  		fmt.Fprintf(GinkgoWriter, "Dropped %d out of %d packets.\n", numDropped, atomic.LoadInt32(&total))
   133  		fmt.Fprintf(GinkgoWriter, "Received %d out of %d sent datagrams.\n", counter, num)
   134  		Expect(counter).To(And(
   135  			BeNumerically(">", expVal*9/10),
   136  			BeNumerically("<", num),
   137  		))
   138  		Eventually(conn.Context().Done).Should(BeClosed())
   139  	})
   140  
   141  	It("server can disable datagram", func() {
   142  		proxyPort, close := startServerAndProxy(false, true)
   143  		raddr, err := net.ResolveUDPAddr("udp", fmt.Sprintf("localhost:%d", proxyPort))
   144  		Expect(err).ToNot(HaveOccurred())
   145  		conn, err := quic.Dial(
   146  			context.Background(),
   147  			clientConn,
   148  			raddr,
   149  			getTLSClientConfig(),
   150  			getQuicConfig(&quic.Config{EnableDatagrams: true}),
   151  		)
   152  		Expect(err).ToNot(HaveOccurred())
   153  		Expect(conn.ConnectionState().SupportsDatagrams).To(BeFalse())
   154  
   155  		close()
   156  		conn.CloseWithError(0, "")
   157  	})
   158  
   159  	It("client can disable datagram", func() {
   160  		proxyPort, close := startServerAndProxy(false, true)
   161  		raddr, err := net.ResolveUDPAddr("udp", fmt.Sprintf("localhost:%d", proxyPort))
   162  		Expect(err).ToNot(HaveOccurred())
   163  		conn, err := quic.Dial(
   164  			context.Background(),
   165  			clientConn,
   166  			raddr,
   167  			getTLSClientConfig(),
   168  			getQuicConfig(&quic.Config{EnableDatagrams: true}),
   169  		)
   170  		Expect(err).ToNot(HaveOccurred())
   171  		Expect(conn.ConnectionState().SupportsDatagrams).To(BeFalse())
   172  
   173  		Expect(conn.SendMessage([]byte{0})).To(HaveOccurred())
   174  
   175  		close()
   176  		conn.CloseWithError(0, "")
   177  	})
   178  })