Working connection test

Author: nhooyr

dial.go

@@ -108,8 +108,8 @@ func Dial(ctx context.Context, u string, opts ...DialOption) (_ *Conn, _ *http.R
 			r := io.LimitReader(resp.Body, 1024)
 			b, _ := ioutil.ReadAll(r)
 			resp.Body = ioutil.NopCloser(bytes.NewReader(b))
+			respBody.Close()
 		}
-		respBody.Close()
 	}()
 
 	if resp.StatusCode != http.StatusSwitchingProtocols {

header_test.go

@@ -36,18 +36,18 @@ func TestHeader(t *testing.T) {
 			rand.Read(h.maskKey[:])
 		}
 
-		t.Logf("header: %#v", h)
-
 		b := marshalHeader(h)
-		t.Logf("bytes: %b", b)
-
 		r := bytes.NewReader(b)
 		h2, err := readHeader(r)
 		if err != nil {
+			t.Logf("header: %#v", h)
+			t.Logf("bytes: %b", b)
 			t.Fatalf("failed to read header: %v", err)
 		}
 
 		if !cmp.Equal(h, h2, cmp.AllowUnexported(header{})) {
+			t.Logf("header: %#v", h)
+			t.Logf("bytes: %b", b)
 			t.Fatalf("parsed and read header differ: %v", cmp.Diff(h, h2, cmp.AllowUnexported(header{})))
 		}
 	}

json.go

@@ -2,14 +2,43 @@ package websocket
 
 import (
 	"context"
+	"encoding/json"
+
+	"golang.org/x/net/websocket"
+	"golang.org/x/xerrors"
 )
 
 // ReadJSON reads a json message from c into v.
 func ReadJSON(ctx context.Context, c *Conn, v interface{}) error {
-	panic("TODO")
+	typ, r, err := c.ReadMessage(ctx)
+	if err != nil {
+		return xerrors.Errorf("failed to read json: %v", err)
+	}
+
+	if typ != websocket.TextFrame {
+		return xerrors.Errorf("unexpected frame type for json (expected TextFrame): %v", typ)
+	}
+
+	d := json.NewDecoder(r)
+	err = d.Decode(v)
+	if err != nil {
+		return xerrors.Errorf("failed to read json: %v", err)
+	}
+	return nil
 }
 
 // WriteJSON writes the json message v into c.
 func WriteJSON(ctx context.Context, c *Conn, v interface{}) error {
-	panic("TODO")
+	w := c.MessageWriter(websocket.TextFrame)
+	w.SetContext(ctx)
+	e := json.NewEncoder(w)
+	err := e.Encode(v)
+	if err != nil {
+		return xerrors.Errorf("failed to write json: %v", err)
+	}
+	err = w.Close()
+	if err != nil {
+		return xerrors.Errorf("failed to write json: %v", err)
+	}
+	return nil
 }

websocket.go

@@ -5,6 +5,7 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"log"
 	"sync"
 	"time"
 
@@ -19,6 +20,7 @@ type controlFrame struct {
 // Conn represents a WebSocket connection.
 // Pings will always be automatically responded to with pongs, you do not
 // have to do anything special.
+// TODO set finalizer
 type Conn struct {
 	subprotocol string
 	br          *bufio.Reader
@@ -80,6 +82,9 @@ func (c *Conn) init() {
 	c.write = make(chan opcode)
 	c.read = make(chan opcode)
 	c.readBytes = make(chan []byte)
+
+	go c.writeLoop()
+	go c.readLoop()
 }
 
 func (c *Conn) writeLoop() {
@@ -197,29 +202,49 @@ func (c *Conn) readLoop() {
 		}
 
 		c.readDone = make(chan struct{})
-
-		var maskPos int
-		left := h.payloadLength
+		c.read <- h.opcode
 		for {
-			select {
-			case <-c.closed:
-				return
-			case b := <-c.readBytes:
-				if int64(len(b)) > left {
-					b = b[left:]
-				}
-				_, err = io.ReadFull(c.br, b)
-				if err != nil {
-					c.close(xerrors.Errorf("failed to read from connection: %v", err))
+			var maskPos int
+			left := h.payloadLength
+			for left > 0 {
+				select {
+				case <-c.closed:
 					return
-				}
+				case b := <-c.readBytes:
+					log.Println("readbytes", left)
+
+					if int64(len(b)) > left {
+						b = b[:left]
+					}
+
+					_, err = io.ReadFull(c.br, b)
+					if err != nil {
+						c.close(xerrors.Errorf("failed to read from connection: %v", err))
+						return
+					}
+					left -= int64(len(b))
 
-				if h.masked {
-					maskPos = mask(h.maskKey, maskPos, b)
+					if h.masked {
+						maskPos = mask(h.maskKey, maskPos, b)
+					}
+
+					select {
+					case <-c.closed:
+						return
+					case c.readDone <- struct{}{}:
+					}
 				}
+			}
 
-				c.readDone <- struct{}{}
+			if h.fin {
+				break
+			}
+			h, err = readHeader(c.br)
+			if err != nil {
+				c.close(xerrors.Errorf("failed to read header: %v", err))
+				return
 			}
+			// TODO check opcode.
 		}
 		close(c.readDone)
 	}

websocket_test.go

@@ -16,9 +16,10 @@ var httpclient = &http.Client{
 	Timeout: time.Second * 15,
 }
 
-func TestHandshake(t *testing.T) {
+func TestConnection(t *testing.T) {
 	t.Parallel()
 
+	obj := make(chan interface{})
 	s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		c, err := websocket.Accept(w, r,
 			websocket.AcceptSubprotocols("myproto"),
@@ -27,19 +28,34 @@ func TestHandshake(t *testing.T) {
 			t.Errorf("failed to accept connection: %v", err)
 			return
 		}
-		_ = c
+		defer c.Close(websocket.StatusInternalError, "")
+
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
+		defer cancel()
+
+		var v interface{}
+		err = websocket.ReadJSON(ctx, c, &v)
+		if err != nil {
+			t.Error(err)
+			return
+		}
+
+		t.Log("success", v)
+		obj <- v
+
+		c.Close(websocket.StatusNormalClosure, "")
 	}))
 	defer s.Close()
 
-	ctx, cancel := context.WithTimeout(context.Background(), time.Second*30)
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
 	defer cancel()
 	c, resp, err := websocket.Dial(ctx, s.URL,
 		websocket.DialSubprotocols("myproto"),
 	)
 	if err != nil {
 		t.Fatalf("failed to do handshake request: %v", err)
 	}
-	_ = c
+	defer c.Close(websocket.StatusInternalError, "")
 
 	checkHeader := func(h, exp string) {
 		t.Helper()
@@ -53,4 +69,21 @@ func TestHandshake(t *testing.T) {
 	checkHeader("Upgrade", "websocket")
 	checkHeader("Sec-WebSocket-Accept", "ICX+Yqv66kxgM0FcWaLWlFLwTAI=")
 	checkHeader("Sec-WebSocket-Protocol", "myproto")
+
+	v := map[string]interface{}{
+		"anmol": "wowow",
+	}
+	err = websocket.WriteJSON(ctx, c, v)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	select {
+	case v2 := <-obj:
+		if !cmp.Equal(v, v2) {
+			t.Fatalf("unexpected value read: %v", cmp.Diff(v, v2))
+		}
+	case <-time.After(time.Second * 10):
+		t.Fatalf("test timed out")
+	}
 }