147 lines
3.4 KiB
Go
147 lines
3.4 KiB
Go
package qq_obv11
|
|
|
|
import (
|
|
"context"
|
|
"net"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
// TestWSClientServerRoundTrip exercises a full handshake plus bidirectional
|
|
// text frames over a real loopback server, verifying that a client's masked
|
|
// frames are unmasked by the server and the server's unmasked frames are read
|
|
// correctly by the client.
|
|
func TestWSClientServerRoundTrip(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
conn, err := upgradeWS(w, r)
|
|
if err != nil {
|
|
t.Errorf("upgrade error: %v", err)
|
|
return
|
|
}
|
|
defer conn.Close()
|
|
op, payload, err := conn.ReadMessage()
|
|
if err != nil {
|
|
return
|
|
}
|
|
if op != opText {
|
|
t.Errorf("server read opcode = %d, want text", op)
|
|
return
|
|
}
|
|
_ = conn.WriteMessage(opText, payload) // echo
|
|
}))
|
|
defer srv.Close()
|
|
|
|
url := "ws://" + strings.TrimPrefix(srv.URL, "http://")
|
|
client, err := dialWS(context.Background(), nil, url, "")
|
|
if err != nil {
|
|
t.Fatalf("dialWS error: %v", err)
|
|
}
|
|
defer client.Close()
|
|
|
|
if err := client.WriteMessage(opText, []byte("hello")); err != nil {
|
|
t.Fatalf("client write error: %v", err)
|
|
}
|
|
_, echo, err := client.ReadMessage()
|
|
if err != nil {
|
|
t.Fatalf("client read error: %v", err)
|
|
}
|
|
if string(echo) != "hello" {
|
|
t.Errorf("echo = %q, want hello", echo)
|
|
}
|
|
}
|
|
|
|
// rawFrame builds an unmasked frame for tests that drive reassembly directly.
|
|
func rawFrame(fin bool, opcode byte, payload []byte) []byte {
|
|
first := opcode
|
|
if fin {
|
|
first |= 0x80
|
|
}
|
|
out := []byte{first}
|
|
n := len(payload)
|
|
switch {
|
|
case n < 126:
|
|
out = append(out, byte(n))
|
|
case n <= 0xFFFF:
|
|
out = append(out, 126, byte(n>>8), byte(n))
|
|
default:
|
|
out = append(out, 127, 0, 0, 0, 0, byte(n>>24), byte(n>>16), byte(n>>8), byte(n))
|
|
}
|
|
return append(out, payload...)
|
|
}
|
|
|
|
func TestWSFragmentedReadAssembles(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
c1, c2 := net.Pipe()
|
|
defer c1.Close()
|
|
defer c2.Close()
|
|
|
|
client := newWSConn(c2, nil, nil, true)
|
|
go func() {
|
|
_, _ = c1.Write(rawFrame(false, opText, []byte("hel")))
|
|
_, _ = c1.Write(rawFrame(true, opContinuation, []byte("lo")))
|
|
}()
|
|
|
|
op, payload, err := client.ReadMessage()
|
|
if err != nil {
|
|
t.Fatalf("fragmented read error: %v", err)
|
|
}
|
|
if op != opText {
|
|
t.Errorf("opcode = %d, want text", op)
|
|
}
|
|
if string(payload) != "hello" {
|
|
t.Errorf("payload = %q, want hello", payload)
|
|
}
|
|
}
|
|
|
|
func TestWSPingRepliesPong(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
c1, c2 := net.Pipe()
|
|
defer c1.Close()
|
|
defer c2.Close()
|
|
|
|
server := newWSConn(c1, nil, nil, false)
|
|
client := newWSConn(c2, nil, nil, true)
|
|
|
|
pong := make(chan string, 1)
|
|
go func() {
|
|
if err := server.WriteMessage(opPing, []byte("hb")); err != nil {
|
|
pong <- "err"
|
|
return
|
|
}
|
|
for {
|
|
_, op, payload, err := server.readFrame()
|
|
if err != nil {
|
|
pong <- "err"
|
|
return
|
|
}
|
|
if op == opPong {
|
|
_ = server.Close() // unblocks the client's pending read
|
|
pong <- string(payload)
|
|
return
|
|
}
|
|
}
|
|
}()
|
|
|
|
// The client's ReadMessage processes the ping (auto-replying a pong) and
|
|
// then blocks until the server closes, which unblocks it with EOF.
|
|
if _, _, err := client.ReadMessage(); err == nil {
|
|
t.Error("client.ReadMessage returned nil error, want EOF after close")
|
|
}
|
|
|
|
select {
|
|
case got := <-pong:
|
|
if got != "hb" {
|
|
t.Errorf("pong payload = %q, want hb", got)
|
|
}
|
|
case <-time.After(3 * time.Second):
|
|
t.Fatal("server never saw a pong")
|
|
}
|
|
}
|