Implement Conn

Author: nhooyr

accept.go

@@ -116,17 +116,22 @@ func Accept(w http.ResponseWriter, r *http.Request, opts ...AcceptOption) (*Conn
 
 	w.WriteHeader(http.StatusSwitchingProtocols)
 
-	c, brw, err := hj.Hijack()
+	netConn, brw, err := hj.Hijack()
 	if err != nil {
 		err = xerrors.Errorf("websocket: failed to hijack connection: %v", err)
 		http.Error(w, err.Error(), http.StatusInternalServerError)
 		return nil, err
 	}
 
-	_ = c
-	_ = brw
+	c := &Conn{
+		subprotocol: w.Header().Get("Sec-WebSocket-Protocol"),
+		br:          brw.Reader,
+		bw:          brw.Writer,
+		closer:      netConn,
+	}
+	c.init()
 
-	return nil, nil
+	return c, nil
 }
 
 func selectSubprotocol(w http.ResponseWriter, r *http.Request, subprotocols []string) {

dial.go

@@ -1,6 +1,7 @@
 package websocket
 
 import (
+	"bufio"
 	"bytes"
 	"context"
 	"encoding/base64"
@@ -132,5 +133,14 @@ func Dial(ctx context.Context, u string, opts ...DialOption) (_ *Conn, _ *http.R
 		return nil, resp, xerrors.Errorf("websocket: body is not a read write closer but should be: %T", rwc)
 	}
 
-	return nil, resp, nil
+	c := &Conn{
+		subprotocol: resp.Header.Get("Sec-WebSocket-Protocol"),
+		br:          bufio.NewReader(rwc),
+		bw:          bufio.NewWriter(rwc),
+		closer:      rwc,
+		client:      true,
+	}
+	c.init()
+
+	return c, resp, nil
 }

header.go

@@ -2,6 +2,7 @@ package websocket
 
 import (
 	"encoding/binary"
+	"fmt"
 	"io"
 
 	"golang.org/x/xerrors"
@@ -16,7 +17,6 @@ const maxHeaderSize = 1 + 1 + 8 + 4
 
 // header represents a WebSocket frame header.
 // See https://tools.ietf.org/html/rfc6455#section-5.2
-// The fields are exported for easy printing for debugging.
 type header struct {
 	fin    bool
 	rsv1   bool
@@ -34,7 +34,7 @@ type header struct {
 
 // bytes returns the bytes of the header.
 // See https://tools.ietf.org/html/rfc6455#section-5.2
-func marshalHeader(h header) ([]byte, error) {
+func marshalHeader(h header) []byte {
 	b := make([]byte, 2, maxHeaderSize)
 
 	if h.fin {
@@ -54,7 +54,7 @@ func marshalHeader(h header) ([]byte, error) {
 
 	switch {
 	case h.payloadLength < 0:
-		return nil, xerrors.Errorf("invalid header: negative length: %v", h.payloadLength)
+		panic(fmt.Sprintf("websocket: invalid header: negative length: %v", h.payloadLength))
 	case h.payloadLength <= 125:
 		b[1] = byte(h.payloadLength)
 	case h.payloadLength <= 1<<16:
@@ -73,7 +73,7 @@ func marshalHeader(h header) ([]byte, error) {
 		copy(b[len(b)-4:], h.maskKey[:])
 	}
 
-	return b, nil
+	return b
 }
 
 // readHeader reads a header from the reader.
@@ -130,6 +130,9 @@ func readHeader(r io.Reader) (header, error) {
 		b = b[2:]
 	case payloadLength == 127:
 		h.payloadLength = int64(binary.BigEndian.Uint64(b))
+		if h.payloadLength < 0 {
+			return header{}, xerrors.Errorf("websocket: header has negative payload length: %v", h.payloadLength)
+		}
 		b = b[8:]
 	}
 

header_test.go

@@ -38,11 +38,7 @@ func TestHeader(t *testing.T) {
 
 		t.Logf("header: %#v", h)
 
-		b, err := marshalHeader(h)
-		if err != nil {
-			t.Fatalf("failed to marshal header: %v", err)
-		}
-
+		b := marshalHeader(h)
 		t.Logf("bytes: %b", b)
 
 		r := bytes.NewReader(b)

opcode.go

@@ -15,3 +15,11 @@ const (
 	opPong
 	// 11-16 are reserved for further control frames.
 )
+
+func (o opcode) controlOp() bool {
+	switch o {
+	case opClose, opPing, opPong:
+		return true
+	}
+	return false
+}

statuscode.go

@@ -5,6 +5,8 @@ import (
 	"errors"
 	"fmt"
 	"math/bits"
+
+	"golang.org/x/xerrors"
 )
 
 // StatusCode represents a WebSocket status code.
@@ -43,18 +45,23 @@ func (e CloseError) Error() string {
 	return fmt.Sprintf("WebSocket closed with status = %v and reason = %q", e.Code, e.Reason)
 }
 
-func parseClosePayload(p []byte) (code StatusCode, reason []byte, err error) {
+func parseClosePayload(p []byte) (code StatusCode, reason string, err error) {
 	if len(p) < 2 {
-		return 0, nil, fmt.Errorf("close payload too small, cannot even contain the 2 byte status code")
+		return 0, "", fmt.Errorf("close payload too small, cannot even contain the 2 byte status code")
 	}
 
 	code = StatusCode(binary.BigEndian.Uint16(p))
-	reason = p[2:]
+	reason = string(p[2:])
 
 	return code, reason, nil
 }
 
-func closePayload(code StatusCode, reason []byte) ([]byte, error) {
+const maxControlFramePayload = 125
+
+func closePayload(code StatusCode, reason string) ([]byte, error) {
+	if len(reason) > maxControlFramePayload-2 {
+		return nil, xerrors.Errorf("reason string max is %v but got %q with length %v", maxControlFramePayload-2, reason, len(reason))
+	}
 	if bits.Len(uint(code)) > 16 {
 		return nil, errors.New("status code is larger than 2 bytes")
 	}