diff --git a/webtransport_test.go b/webtransport_test.go index d819191..79b4bb5 100644 --- a/webtransport_test.go +++ b/webtransport_test.go @@ -2,7 +2,6 @@ package webtransport_test import ( "context" - "crypto/rand" "crypto/tls" "errors" "fmt" @@ -15,6 +14,8 @@ import ( "testing" "time" + "golang.org/x/exp/rand" + "github.com/quic-go/webtransport-go" "github.com/quic-go/quic-go" @@ -595,3 +596,65 @@ func TestWriteCloseRace(t *testing.T) { <-ready close(ch) } + +func TestDatagrams(t *testing.T) { + errChan := make(chan error, 1) + sess, closeServer := establishSession(t, func(sess *webtransport.Session) { + for { + d, err := sess.ReceiveDatagram(context.Background()) + if err != nil { + errChan <- err + return + } + if err := sess.SendDatagram(d); err != nil { + errChan <- err + return + } + } + }) + defer closeServer() + + const num = 10 + m := make(map[string]bool, num) + var mx sync.Mutex + done := make(chan struct{}) + var counter int + go func() { + defer close(done) + for { + b, err := sess.ReceiveDatagram(context.Background()) + fmt.Println("received datagram") + if err != nil { + fmt.Println(err) + return + } + mx.Lock() + if _, ok := m[string(b)]; !ok { + t.Fatal("received unexpected datagram") + } + m[string(b)] = true + mx.Unlock() + counter++ + } + }() + for i := 0; i < num; i++ { + b := make([]byte, 800) + rand.Read(b) + mx.Lock() + m[string(b)] = false + mx.Unlock() + fmt.Println("sending datagram", i) + if err := sess.SendDatagram(b); err != nil { + break + } + } + time.Sleep(time.Second) + sess.CloseWithError(0, "") + select { + case <-done: + t.Logf("sent: %d, received: %d", num, counter) + require.Greater(t, counter, num*4/5) + case <-time.After(5 * time.Second): + t.Fatal("timeout") + } +}