Implement handshake

Author: nhooyr

.github/fmt/entrypoint.sh

@@ -8,7 +8,6 @@ gen() {
 	go list ./... > /dev/null
 	go mod tidy
 
-	go install golang.org/x/tools/cmd/stringer
 	go generate ./...
 }
 

README.md

@@ -17,6 +17,7 @@ go get nhooyr.io/websocket@master
 ## Features
 
 - Full support of the WebSocket protocol
+- Only depends on the stdlib
 - Simple to use because of the minimal API
 - Uses the context package for cancellation
 - Uses net/http's Client to do WebSocket dials

accept.go

@@ -1,33 +1,48 @@
 package websocket
 
 import (
-	"fmt"
+	"crypto/sha1"
+	"encoding/base64"
 	"net/http"
+	"net/url"
+	"strings"
+
+	"golang.org/x/net/http/httpguts"
+	"golang.org/x/xerrors"
 )
 
 // AcceptOption is an option that can be passed to Accept.
+// The implementations of this interface are printable.
 type AcceptOption interface {
 	acceptOption()
-	fmt.Stringer
 }
 
+type acceptSubprotocols []string
+
+func (o acceptSubprotocols) acceptOption() {}
+
 // AcceptSubprotocols list the subprotocols that Accept will negotiate with a client.
 // The first protocol that a client supports will be negotiated.
-// Pass "" as a subprotocol if you would like to allow the default protocol.
+// Pass "" as a subprotocol if you would like to allow the default protocol along with
+// specific subprotocols.
 func AcceptSubprotocols(subprotocols ...string) AcceptOption {
-	panic("TODO")
+	return acceptSubprotocols(subprotocols)
 }
 
+type acceptOrigins []string
+
+func (o acceptOrigins) acceptOption() {}
+
 // AcceptOrigins lists the origins that Accept will accept.
 // Accept will always accept r.Host as the origin so you do not need to
 // specify that with this option.
 //
 // Use this option with caution to avoid exposing your WebSocket
 // server to a CSRF attack.
 // See https://stackoverflow.com/a/37837709/4283659
-// You can use a * to specify wildcards.
+// You can use a * for wildcards.
 func AcceptOrigins(origins ...string) AcceptOption {
-	panic("TODO")
+	return AcceptOrigins(origins...)
 }
 
 // Accept accepts a WebSocket handshake from a client and upgrades the
@@ -36,5 +51,121 @@ func AcceptOrigins(origins ...string) AcceptOption {
 // InsecureAcceptOrigin is passed.
 // Accept uses w to write the handshake response so the timeouts on the http.Server apply.
 func Accept(w http.ResponseWriter, r *http.Request, opts ...AcceptOption) (*Conn, error) {
-	panic("TODO")
+	var subprotocols []string
+	origins := []string{r.Host}
+	for _, opt := range opts {
+		switch opt := opt.(type) {
+		case acceptOrigins:
+			origins = []string(opt)
+		case acceptSubprotocols:
+			subprotocols = []string(opt)
+		}
+	}
+
+	if !httpguts.HeaderValuesContainsToken(r.Header["Connection"], "Upgrade") {
+		err := xerrors.Errorf("websocket: protocol violation: Connection header does not contain Upgrade: %q", r.Header.Get("Connection"))
+		http.Error(w, err.Error(), http.StatusBadRequest)
+		return nil, err
+	}
+
+	if !httpguts.HeaderValuesContainsToken(r.Header["Upgrade"], "websocket") {
+		err := xerrors.Errorf("websocket: protocol violation: Upgrade header does not contain websocket: %q", r.Header.Get("Upgrade"))
+		http.Error(w, err.Error(), http.StatusBadRequest)
+		return nil, err
+	}
+
+	if r.Method != "GET" {
+		err := xerrors.Errorf("websocket: protocol violation: handshake request method is not GET: %q", r.Method)
+		http.Error(w, err.Error(), http.StatusBadRequest)
+		return nil, err
+	}
+
+	if r.Header.Get("Sec-WebSocket-Version") != "13" {
+		err := xerrors.Errorf("websocket: unsupported protocol version: %q", r.Header.Get("Sec-WebSocket-Version"))
+		http.Error(w, err.Error(), http.StatusBadRequest)
+		return nil, err
+	}
+
+	if r.Header.Get("Sec-WebSocket-Key") == "" {
+		err := xerrors.New("websocket: protocol violation: missing Sec-WebSocket-Key")
+		http.Error(w, err.Error(), http.StatusBadRequest)
+		return nil, err
+	}
+
+	origins = append(origins, r.Host)
+
+	err := authenticateOrigin(r, origins)
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusForbidden)
+		return nil, err
+	}
+
+	hj, ok := w.(http.Hijacker)
+	if !ok {
+		err = xerrors.New("websocket: response writer does not implement http.Hijacker")
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return nil, err
+	}
+
+	w.Header().Set("Upgrade", "websocket")
+	w.Header().Set("Connection", "Upgrade")
+
+	handleKey(w, r)
+
+	selectSubprotocol(w, r, subprotocols)
+
+	w.WriteHeader(http.StatusSwitchingProtocols)
+
+	c, 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
+
+	return nil, nil
+}
+
+func selectSubprotocol(w http.ResponseWriter, r *http.Request, subprotocols []string) {
+	clientSubprotocols := strings.Split(r.Header.Get("Sec-WebSocket-Protocol"), "\n")
+	for _, sp := range subprotocols {
+		for _, cp := range clientSubprotocols {
+			if sp == strings.TrimSpace(cp) {
+				w.Header().Set("Sec-WebSocket-Protocol", sp)
+				return
+			}
+		}
+	}
+}
+
+var keyGUID = []byte("258EAFA5-E914-47DA-95CA-C5AB0DC85B11")
+
+func handleKey(w http.ResponseWriter, r *http.Request) {
+	key := r.Header.Get("Sec-WebSocket-Key")
+	h := sha1.New()
+	h.Write([]byte(key))
+	h.Write(keyGUID)
+
+	responseKey := base64.StdEncoding.EncodeToString(h.Sum(nil))
+	w.Header().Set("Sec-WebSocket-Accept", responseKey)
+}
+
+func authenticateOrigin(r *http.Request, origins []string) error {
+	origin := r.Header.Get("Origin")
+	if origin == "" {
+		return nil
+	}
+	u, err := url.Parse(origin)
+	if err != nil {
+		return xerrors.Errorf("failed to parse Origin header %q: %v", origin, err)
+	}
+	for _, o := range origins {
+		if u.Host == o {
+			return nil
+		}
+	}
+	return xerrors.New("request origin is not authorized")
 }

datatype.go

@@ -1,7 +1,7 @@
 package websocket
 
 // DataType represents the Opcode of a WebSocket data frame.
-//go:generate stringer -type=DataType
+//go:generate go run golang.org/x/tools/cmd/stringer -type=DataType
 type DataType int
 
 // DataType constants.

dial.go

@@ -1,34 +1,53 @@
 package websocket
 
 import (
+	"bytes"
 	"context"
 	"encoding/base64"
-	"fmt"
+	"io"
+	"io/ioutil"
 	"net/http"
+	"net/url"
+	"strings"
+
+	"golang.org/x/net/http/httpguts"
+	"golang.org/x/xerrors"
 )
 
 // DialOption represents a dial option that can be passed to Dial.
+// The implementations are printable for easy debugging.
 type DialOption interface {
 	dialOption()
-	fmt.Stringer
 }
 
+type dialHTTPClient http.Client
+
+func (o dialHTTPClient) dialOption() {}
+
 // DialHTTPClient is the http client used for the handshake.
 // Its Transport must use HTTP/1.1 and must return writable bodies
 // for WebSocket handshakes.
 // http.Transport does this correctly.
-func DialHTTPClient(h *http.Client) DialOption {
-	panic("TODO")
+func DialHTTPClient(hc *http.Client) DialOption {
+	return (*dialHTTPClient)(hc)
 }
 
+type dialHeader http.Header
+
+func (o dialHeader) dialOption() {}
+
 // DialHeader are the HTTP headers included in the handshake request.
 func DialHeader(h http.Header) DialOption {
-	panic("TODO")
+	return dialHeader(h)
 }
 
+type dialSubprotocols []string
+
+func (o dialSubprotocols) dialOption() {}
+
 // DialSubprotocols accepts a slice of protcols to include in the Sec-WebSocket-Protocol header.
 func DialSubprotocols(subprotocols ...string) DialOption {
-	panic("TODO")
+	return dialSubprotocols(subprotocols)
 }
 
 // We use this key for all client requests as the Sec-WebSocket-Key header is useless.
@@ -37,6 +56,81 @@ func DialSubprotocols(subprotocols ...string) DialOption {
 var secWebSocketKey = base64.StdEncoding.EncodeToString(make([]byte, 16))
 
 // Dial performs a WebSocket handshake on the given url with the given options.
-func Dial(ctx context.Context, u string, opts ...DialOption) (*Conn, *http.Response, error) {
-	panic("TODO")
+func Dial(ctx context.Context, u string, opts ...DialOption) (_ *Conn, _ *http.Response, err error) {
+	httpClient := http.DefaultClient
+	var subprotocols []string
+	header := http.Header{}
+	for _, o := range opts {
+		switch o := o.(type) {
+		case dialSubprotocols:
+			subprotocols = o
+		case dialHeader:
+			header = http.Header(o)
+		case *dialHTTPClient:
+			httpClient = (*http.Client)(o)
+		}
+	}
+
+	parsedURL, err := url.Parse(u)
+	if err != nil {
+		return nil, nil, xerrors.Errorf("failed to parse websocket url: %v", err)
+	}
+
+	switch parsedURL.Scheme {
+	case "ws", "http":
+		parsedURL.Scheme = "http"
+	case "wss", "https":
+		parsedURL.Scheme = "https"
+	default:
+		return nil, nil, xerrors.Errorf("unknown scheme in url: %q", parsedURL.Scheme)
+	}
+
+	req, _ := http.NewRequest("GET", u, nil)
+	req = req.WithContext(ctx)
+	req.Header = header
+	req.Header.Set("Connection", "Upgrade")
+	req.Header.Set("Upgrade", "websocket")
+	req.Header.Set("Sec-WebSocket-Version", "13")
+	req.Header.Set("Sec-WebSocket-Key", secWebSocketKey)
+	if len(subprotocols) > 0 {
+		req.Header.Set("Sec-WebSocket-Protocol", strings.Join(subprotocols, ","))
+	}
+
+	resp, err := httpClient.Do(req)
+	if err != nil {
+		return nil, nil, err
+	}
+	defer func() {
+		respBody := resp.Body
+		if err != nil {
+			// We read a bit of the body for better debugging.
+			r := io.LimitReader(resp.Body, 1024)
+			b, _ := ioutil.ReadAll(r)
+			resp.Body = ioutil.NopCloser(bytes.NewReader(b))
+		}
+		respBody.Close()
+	}()
+
+	if resp.StatusCode != http.StatusSwitchingProtocols {
+		return nil, resp, xerrors.Errorf("websocket: expected status code %v but got %v", http.StatusSwitchingProtocols)
+	}
+
+	if !httpguts.HeaderValuesContainsToken(resp.Header["Connection"], "Upgrade") {
+		return nil, resp, xerrors.Errorf("websocket: protocol violation: Connection header does not contain Upgrade: %q", resp.Header.Get("Connection"))
+	}
+
+	if !httpguts.HeaderValuesContainsToken(resp.Header["Upgrade"], "websocket") {
+		return nil, resp, xerrors.Errorf("websocket: protocol violation: Upgrade header does not contain websocket: %q", resp.Header.Get("Upgrade"))
+
+	}
+
+	// We do not care about Sec-WebSocket-Accept because it does not matter.
+	// See the secWebSocketKey global variable.
+
+	rwc, ok := resp.Body.(io.ReadWriteCloser)
+	if !ok {
+		return nil, resp, xerrors.Errorf("websocket: body is not a read write closer but should be: %T", rwc)
+	}
+
+	return nil, resp, nil
 }

doc.go

@@ -1,4 +1,6 @@
 // Package websocket is a minimal and idiomatic implementation of the WebSocket protocol.
 //
+// See https://tools.ietf.org/html/rfc6455
+//
 // For now the docs are at https://github.com/nhooyr/websocket#websocket. I will move them here later.
 package websocket

go.mod

@@ -7,7 +7,9 @@ require (
 	github.com/kr/pretty v0.1.0 // indirect
 	go.coder.com/go-tools v0.0.0-20190317003359-0c6a35b74a16
 	golang.org/x/lint v0.0.0-20190313153728-d0100b6bd8b3
+	golang.org/x/net v0.0.0-20190311183353-d8887717615a
 	golang.org/x/time v0.0.0-20190308202827-9d24e82272b4
 	golang.org/x/tools v0.0.0-20190329215204-73054e8977d1
+	golang.org/x/xerrors v0.0.0-20190315151331-d61658bd2e18
 	mvdan.cc/sh v2.6.4+incompatible
 )

go.sum

@@ -20,5 +20,7 @@ golang.org/x/time v0.0.0-20190308202827-9d24e82272b4/go.mod h1:tRJNPiyCQ0inRvYxb
 golang.org/x/tools v0.0.0-20190311212946-11955173bddd/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs=
 golang.org/x/tools v0.0.0-20190329215204-73054e8977d1 h1:rLRH2E2wN5JjGJSVlBe1ioUkCKgb6eoL9X8bDmtEpsk=
 golang.org/x/tools v0.0.0-20190329215204-73054e8977d1/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs=
+golang.org/x/xerrors v0.0.0-20190315151331-d61658bd2e18 h1:1AGvnywFL1aB5KLRxyLseWJI6aSYPo3oF7HSpXdWQdU=
+golang.org/x/xerrors v0.0.0-20190315151331-d61658bd2e18/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 mvdan.cc/sh v2.6.4+incompatible h1:eD6tDeh0pw+/TOTI1BBEryZ02rD2nMcFsgcvde7jffM=
 mvdan.cc/sh v2.6.4+incompatible/go.mod h1:IeeQbZq+x2SUGBensq/jge5lLQbS3XT2ktyp3wrt4x8=

opcode.go

@@ -1,7 +1,7 @@
 package websocket
 
 // opcode represents a WebSocket Opcode.
-//go:generate stringer -type=opcode
+//go:generate go run golang.org/x/tools/cmd/stringer -type=opcode
 type opcode int
 
 // opcode constants.

statuscode.go

@@ -8,7 +8,7 @@ import (
 )
 
 // StatusCode represents a WebSocket status code.
-//go:generate stringer -type=StatusCode
+//go:generate go run golang.org/x/tools/cmd/stringer -type=StatusCode
 type StatusCode int
 
 // These codes were retrieved from: