wsclient_test.go raw

   1  package relay
   2  
   3  import (
   4  	"bufio"
   5  	"crypto/rand"
   6  	"encoding/base64"
   7  	"encoding/binary"
   8  	"fmt"
   9  	"io"
  10  	"net"
  11  	"strings"
  12  	"testing"
  13  	"time"
  14  )
  15  
  16  // testWS is a minimal RFC 6455 client used by the integration tests.
  17  type testWS struct {
  18  	conn net.Conn
  19  	br   *bufio.Reader
  20  }
  21  
  22  func dialWS(t *testing.T, addr string) *testWS {
  23  	t.Helper()
  24  	conn, err := net.DialTimeout("tcp", addr, 3*time.Second)
  25  	if err != nil {
  26  		t.Fatalf("dial: %v", err)
  27  	}
  28  	keyBytes := make([]byte, 16)
  29  	rand.Read(keyBytes)
  30  	key := base64.StdEncoding.EncodeToString(keyBytes)
  31  	req := "GET / HTTP/1.1\r\nHost: " + addr +
  32  		"\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Key: " + key +
  33  		"\r\nSec-WebSocket-Version: 13\r\n\r\n"
  34  	if _, err := conn.Write([]byte(req)); err != nil {
  35  		t.Fatalf("handshake write: %v", err)
  36  	}
  37  	br := bufio.NewReader(conn)
  38  	line, err := br.ReadString('\n')
  39  	if err != nil {
  40  		t.Fatalf("handshake read: %v", err)
  41  	}
  42  	if !strings.Contains(line, "101") {
  43  		t.Fatalf("handshake status: %q", line)
  44  	}
  45  	for {
  46  		h, err := br.ReadString('\n')
  47  		if err != nil {
  48  			t.Fatalf("handshake headers: %v", err)
  49  		}
  50  		if h == "\r\n" || h == "\n" {
  51  			break
  52  		}
  53  	}
  54  	return &testWS{conn: conn, br: br}
  55  }
  56  
  57  func (c *testWS) writeText(payload []byte) error {
  58  	mask := make([]byte, 4)
  59  	rand.Read(mask)
  60  	hdr := make([]byte, 0, 14)
  61  	hdr = append(hdr, 0x81) // FIN + text
  62  	switch {
  63  	case len(payload) < 126:
  64  		hdr = append(hdr, 0x80|byte(len(payload)))
  65  	case len(payload) <= 0xffff:
  66  		hdr = append(hdr, 0x80|126, byte(len(payload)>>8), byte(len(payload)))
  67  	default:
  68  		hdr = append(hdr, 0x80|127)
  69  		var ext [8]byte
  70  		binary.BigEndian.PutUint64(ext[:], uint64(len(payload)))
  71  		hdr = append(hdr, ext[:]...)
  72  	}
  73  	hdr = append(hdr, mask...)
  74  	body := make([]byte, len(payload))
  75  	for i, b := range payload {
  76  		body[i] = b ^ mask[i&3]
  77  	}
  78  	if _, err := c.conn.Write(append(hdr, body...)); err != nil {
  79  		return err
  80  	}
  81  	return nil
  82  }
  83  
  84  func (c *testWS) read() ([]byte, error) {
  85  	c.conn.SetReadDeadline(time.Now().Add(5 * time.Second))
  86  	for {
  87  		hdr := make([]byte, 2)
  88  		if _, err := io.ReadFull(c.br, hdr); err != nil {
  89  			return nil, err
  90  		}
  91  		op := hdr[0] & 0x0f
  92  		length := int64(hdr[1] & 0x7f)
  93  		switch length {
  94  		case 126:
  95  			ext := make([]byte, 2)
  96  			if _, err := io.ReadFull(c.br, ext); err != nil {
  97  				return nil, err
  98  			}
  99  			length = int64(binary.BigEndian.Uint16(ext))
 100  		case 127:
 101  			ext := make([]byte, 8)
 102  			if _, err := io.ReadFull(c.br, ext); err != nil {
 103  				return nil, err
 104  			}
 105  			length = int64(binary.BigEndian.Uint64(ext))
 106  		}
 107  		payload := make([]byte, length)
 108  		if _, err := io.ReadFull(c.br, payload); err != nil {
 109  			return nil, err
 110  		}
 111  		switch op {
 112  		case 0x9: // ping
 113  			if err := c.writeFrame(0xa, payload); err != nil {
 114  				return nil, err
 115  			}
 116  			continue
 117  		case 0x8: // close
 118  			return nil, io.EOF
 119  		case 0x1, 0x2:
 120  			return payload, nil
 121  		default:
 122  			continue
 123  		}
 124  	}
 125  }
 126  
 127  func (c *testWS) writeFrame(op byte, payload []byte) error {
 128  	mask := make([]byte, 4)
 129  	rand.Read(mask)
 130  	hdr := []byte{0x80 | op, 0x80 | byte(len(payload))}
 131  	hdr = append(hdr, mask...)
 132  	body := make([]byte, len(payload))
 133  	for i, b := range payload {
 134  		body[i] = b ^ mask[i&3]
 135  	}
 136  	_, err := c.conn.Write(append(hdr, body...))
 137  	return err
 138  }
 139  
 140  // expect reads messages until one starts with prefix and returns it.
 141  func (c *testWS) expect(prefix string) string {
 142  	for i := 0; i < 20; i++ {
 143  		msg, err := c.read()
 144  		if err != nil {
 145  			panic(fmt.Sprintf("read while waiting for %q: %v", prefix, err))
 146  		}
 147  		if strings.HasPrefix(string(msg), prefix) {
 148  			return string(msg)
 149  		}
 150  	}
 151  	panic("expected message not seen: " + prefix)
 152  }
 153  
 154  func (c *testWS) close() { c.conn.Close() }
 155