Skip to content

Commit

Permalink
Move all Wasm related code into ws_js.go
Browse files Browse the repository at this point in the history
This way we don't pollute the directory tree.
  • Loading branch information
nhooyr committed May 18, 2020
1 parent 02861b4 commit b453d3e
Show file tree
Hide file tree
Showing 8 changed files with 783 additions and 677 deletions.
20 changes: 0 additions & 20 deletions accept_js.go

This file was deleted.

205 changes: 205 additions & 0 deletions close.go
Original file line number Diff line number Diff line change
@@ -1,8 +1,16 @@
// +build !js

package websocket

import (
"context"
"encoding/binary"
"errors"
"fmt"
"log"
"time"

"nhooyr.io/websocket/internal/errd"
)

// StatusCode represents a WebSocket status code.
Expand Down Expand Up @@ -74,3 +82,200 @@ func CloseStatus(err error) StatusCode {
}
return -1
}

// Close performs the WebSocket close handshake with the given status code and reason.
//
// It will write a WebSocket close frame with a timeout of 5s and then wait 5s for
// the peer to send a close frame.
// All data messages received from the peer during the close handshake will be discarded.
//
// The connection can only be closed once. Additional calls to Close
// are no-ops.
//
// The maximum length of reason must be 125 bytes. Avoid
// sending a dynamic reason.
//
// Close will unblock all goroutines interacting with the connection once
// complete.
func (c *Conn) Close(code StatusCode, reason string) error {
return c.closeHandshake(code, reason)
}

func (c *Conn) closeHandshake(code StatusCode, reason string) (err error) {
defer errd.Wrap(&err, "failed to close WebSocket")

writeErr := c.writeClose(code, reason)
closeHandshakeErr := c.waitCloseHandshake()

if writeErr != nil {
return writeErr
}

if CloseStatus(closeHandshakeErr) == -1 {
return closeHandshakeErr
}

return nil
}

var errAlreadyWroteClose = errors.New("already wrote close")

func (c *Conn) writeClose(code StatusCode, reason string) error {
c.closeMu.Lock()
wroteClose := c.wroteClose
c.wroteClose = true
c.closeMu.Unlock()
if wroteClose {
return errAlreadyWroteClose
}

ce := CloseError{
Code: code,
Reason: reason,
}

var p []byte
var marshalErr error
if ce.Code != StatusNoStatusRcvd {
p, marshalErr = ce.bytes()
if marshalErr != nil {
log.Printf("websocket: %v", marshalErr)
}
}

writeErr := c.writeControl(context.Background(), opClose, p)
if CloseStatus(writeErr) != -1 {
// Not a real error if it's due to a close frame being received.
writeErr = nil
}

// We do this after in case there was an error writing the close frame.
c.setCloseErr(fmt.Errorf("sent close frame: %w", ce))

if marshalErr != nil {
return marshalErr
}
return writeErr
}

func (c *Conn) waitCloseHandshake() error {
defer c.close(nil)

ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
defer cancel()

err := c.readMu.lock(ctx)
if err != nil {
return err
}
defer c.readMu.unlock()

if c.readCloseFrameErr != nil {
return c.readCloseFrameErr
}

for {
h, err := c.readLoop(ctx)
if err != nil {
return err
}

for i := int64(0); i < h.payloadLength; i++ {
_, err := c.br.ReadByte()
if err != nil {
return err
}
}
}
}

func parseClosePayload(p []byte) (CloseError, error) {
if len(p) == 0 {
return CloseError{
Code: StatusNoStatusRcvd,
}, nil
}

if len(p) < 2 {
return CloseError{}, fmt.Errorf("close payload %q too small, cannot even contain the 2 byte status code", p)
}

ce := CloseError{
Code: StatusCode(binary.BigEndian.Uint16(p)),
Reason: string(p[2:]),
}

if !validWireCloseCode(ce.Code) {
return CloseError{}, fmt.Errorf("invalid status code %v", ce.Code)
}

return ce, nil
}

// See http://www.iana.org/assignments/websocket/websocket.xhtml#close-code-number
// and https://tools.ietf.org/html/rfc6455#section-7.4.1
func validWireCloseCode(code StatusCode) bool {
switch code {
case statusReserved, StatusNoStatusRcvd, StatusAbnormalClosure, StatusTLSHandshake:
return false
}

if code >= StatusNormalClosure && code <= StatusBadGateway {
return true
}
if code >= 3000 && code <= 4999 {
return true
}

return false
}

func (ce CloseError) bytes() ([]byte, error) {
p, err := ce.bytesErr()
if err != nil {
err = fmt.Errorf("failed to marshal close frame: %w", err)
ce = CloseError{
Code: StatusInternalError,
}
p, _ = ce.bytesErr()
}
return p, err
}

const maxCloseReason = maxControlPayload - 2

func (ce CloseError) bytesErr() ([]byte, error) {
if len(ce.Reason) > maxCloseReason {
return nil, fmt.Errorf("reason string max is %v but got %q with length %v", maxCloseReason, ce.Reason, len(ce.Reason))
}

if !validWireCloseCode(ce.Code) {
return nil, fmt.Errorf("status code %v cannot be set", ce.Code)
}

buf := make([]byte, 2+len(ce.Reason))
binary.BigEndian.PutUint16(buf, uint16(ce.Code))
copy(buf[2:], ce.Reason)
return buf, nil
}

func (c *Conn) setCloseErr(err error) {
c.closeMu.Lock()
c.setCloseErrLocked(err)
c.closeMu.Unlock()
}

func (c *Conn) setCloseErrLocked(err error) {
if c.closeErr == nil {
c.closeErr = fmt.Errorf("WebSocket closed: %w", err)
}
}

func (c *Conn) isClosed() bool {
select {
case <-c.closed:
return true
default:
return false
}
}

0 comments on commit b453d3e

Please sign in to comment.