 f80b4dc586
			
		
	
	f80b4dc586
	
	
	
		
			
			full diff: f3200d17e0...ab34263943
Worth mentioning that there's a comment updated in golang.org/x/net/websocket:
    This package currently lacks some features found in alternative
    and more actively maintained WebSocket packages:
        https://godoc.org/github.com/gorilla/websocket
        https://godoc.org/nhooyr.io/websocket
It's used in k8s.io/apiserver/pkg/util/wsstream/stream.go, so perhaps that should
be reviewed if the alternatives are better for how it's used.
Signed-off-by: Sebastiaan van Stijn <github@gone.nl>
		
	
		
			
				
	
	
		
			452 lines
		
	
	
		
			12 KiB
		
	
	
	
		
			Go
		
	
	
	
	
	
			
		
		
	
	
			452 lines
		
	
	
		
			12 KiB
		
	
	
	
		
			Go
		
	
	
	
	
	
| // Copyright 2009 The Go Authors. All rights reserved.
 | |
| // Use of this source code is governed by a BSD-style
 | |
| // license that can be found in the LICENSE file.
 | |
| 
 | |
| // Package websocket implements a client and server for the WebSocket protocol
 | |
| // as specified in RFC 6455.
 | |
| //
 | |
| // This package currently lacks some features found in alternative
 | |
| // and more actively maintained WebSocket packages:
 | |
| //
 | |
| //     https://godoc.org/github.com/gorilla/websocket
 | |
| //     https://godoc.org/nhooyr.io/websocket
 | |
| package websocket // import "golang.org/x/net/websocket"
 | |
| 
 | |
| import (
 | |
| 	"bufio"
 | |
| 	"crypto/tls"
 | |
| 	"encoding/json"
 | |
| 	"errors"
 | |
| 	"io"
 | |
| 	"io/ioutil"
 | |
| 	"net"
 | |
| 	"net/http"
 | |
| 	"net/url"
 | |
| 	"sync"
 | |
| 	"time"
 | |
| )
 | |
| 
 | |
| const (
 | |
| 	ProtocolVersionHybi13    = 13
 | |
| 	ProtocolVersionHybi      = ProtocolVersionHybi13
 | |
| 	SupportedProtocolVersion = "13"
 | |
| 
 | |
| 	ContinuationFrame = 0
 | |
| 	TextFrame         = 1
 | |
| 	BinaryFrame       = 2
 | |
| 	CloseFrame        = 8
 | |
| 	PingFrame         = 9
 | |
| 	PongFrame         = 10
 | |
| 	UnknownFrame      = 255
 | |
| 
 | |
| 	DefaultMaxPayloadBytes = 32 << 20 // 32MB
 | |
| )
 | |
| 
 | |
| // ProtocolError represents WebSocket protocol errors.
 | |
| type ProtocolError struct {
 | |
| 	ErrorString string
 | |
| }
 | |
| 
 | |
| func (err *ProtocolError) Error() string { return err.ErrorString }
 | |
| 
 | |
| var (
 | |
| 	ErrBadProtocolVersion   = &ProtocolError{"bad protocol version"}
 | |
| 	ErrBadScheme            = &ProtocolError{"bad scheme"}
 | |
| 	ErrBadStatus            = &ProtocolError{"bad status"}
 | |
| 	ErrBadUpgrade           = &ProtocolError{"missing or bad upgrade"}
 | |
| 	ErrBadWebSocketOrigin   = &ProtocolError{"missing or bad WebSocket-Origin"}
 | |
| 	ErrBadWebSocketLocation = &ProtocolError{"missing or bad WebSocket-Location"}
 | |
| 	ErrBadWebSocketProtocol = &ProtocolError{"missing or bad WebSocket-Protocol"}
 | |
| 	ErrBadWebSocketVersion  = &ProtocolError{"missing or bad WebSocket Version"}
 | |
| 	ErrChallengeResponse    = &ProtocolError{"mismatch challenge/response"}
 | |
| 	ErrBadFrame             = &ProtocolError{"bad frame"}
 | |
| 	ErrBadFrameBoundary     = &ProtocolError{"not on frame boundary"}
 | |
| 	ErrNotWebSocket         = &ProtocolError{"not websocket protocol"}
 | |
| 	ErrBadRequestMethod     = &ProtocolError{"bad method"}
 | |
| 	ErrNotSupported         = &ProtocolError{"not supported"}
 | |
| )
 | |
| 
 | |
| // ErrFrameTooLarge is returned by Codec's Receive method if payload size
 | |
| // exceeds limit set by Conn.MaxPayloadBytes
 | |
| var ErrFrameTooLarge = errors.New("websocket: frame payload size exceeds limit")
 | |
| 
 | |
| // Addr is an implementation of net.Addr for WebSocket.
 | |
| type Addr struct {
 | |
| 	*url.URL
 | |
| }
 | |
| 
 | |
| // Network returns the network type for a WebSocket, "websocket".
 | |
| func (addr *Addr) Network() string { return "websocket" }
 | |
| 
 | |
| // Config is a WebSocket configuration
 | |
| type Config struct {
 | |
| 	// A WebSocket server address.
 | |
| 	Location *url.URL
 | |
| 
 | |
| 	// A Websocket client origin.
 | |
| 	Origin *url.URL
 | |
| 
 | |
| 	// WebSocket subprotocols.
 | |
| 	Protocol []string
 | |
| 
 | |
| 	// WebSocket protocol version.
 | |
| 	Version int
 | |
| 
 | |
| 	// TLS config for secure WebSocket (wss).
 | |
| 	TlsConfig *tls.Config
 | |
| 
 | |
| 	// Additional header fields to be sent in WebSocket opening handshake.
 | |
| 	Header http.Header
 | |
| 
 | |
| 	// Dialer used when opening websocket connections.
 | |
| 	Dialer *net.Dialer
 | |
| 
 | |
| 	handshakeData map[string]string
 | |
| }
 | |
| 
 | |
| // serverHandshaker is an interface to handle WebSocket server side handshake.
 | |
| type serverHandshaker interface {
 | |
| 	// ReadHandshake reads handshake request message from client.
 | |
| 	// Returns http response code and error if any.
 | |
| 	ReadHandshake(buf *bufio.Reader, req *http.Request) (code int, err error)
 | |
| 
 | |
| 	// AcceptHandshake accepts the client handshake request and sends
 | |
| 	// handshake response back to client.
 | |
| 	AcceptHandshake(buf *bufio.Writer) (err error)
 | |
| 
 | |
| 	// NewServerConn creates a new WebSocket connection.
 | |
| 	NewServerConn(buf *bufio.ReadWriter, rwc io.ReadWriteCloser, request *http.Request) (conn *Conn)
 | |
| }
 | |
| 
 | |
| // frameReader is an interface to read a WebSocket frame.
 | |
| type frameReader interface {
 | |
| 	// Reader is to read payload of the frame.
 | |
| 	io.Reader
 | |
| 
 | |
| 	// PayloadType returns payload type.
 | |
| 	PayloadType() byte
 | |
| 
 | |
| 	// HeaderReader returns a reader to read header of the frame.
 | |
| 	HeaderReader() io.Reader
 | |
| 
 | |
| 	// TrailerReader returns a reader to read trailer of the frame.
 | |
| 	// If it returns nil, there is no trailer in the frame.
 | |
| 	TrailerReader() io.Reader
 | |
| 
 | |
| 	// Len returns total length of the frame, including header and trailer.
 | |
| 	Len() int
 | |
| }
 | |
| 
 | |
| // frameReaderFactory is an interface to creates new frame reader.
 | |
| type frameReaderFactory interface {
 | |
| 	NewFrameReader() (r frameReader, err error)
 | |
| }
 | |
| 
 | |
| // frameWriter is an interface to write a WebSocket frame.
 | |
| type frameWriter interface {
 | |
| 	// Writer is to write payload of the frame.
 | |
| 	io.WriteCloser
 | |
| }
 | |
| 
 | |
| // frameWriterFactory is an interface to create new frame writer.
 | |
| type frameWriterFactory interface {
 | |
| 	NewFrameWriter(payloadType byte) (w frameWriter, err error)
 | |
| }
 | |
| 
 | |
| type frameHandler interface {
 | |
| 	HandleFrame(frame frameReader) (r frameReader, err error)
 | |
| 	WriteClose(status int) (err error)
 | |
| }
 | |
| 
 | |
| // Conn represents a WebSocket connection.
 | |
| //
 | |
| // Multiple goroutines may invoke methods on a Conn simultaneously.
 | |
| type Conn struct {
 | |
| 	config  *Config
 | |
| 	request *http.Request
 | |
| 
 | |
| 	buf *bufio.ReadWriter
 | |
| 	rwc io.ReadWriteCloser
 | |
| 
 | |
| 	rio sync.Mutex
 | |
| 	frameReaderFactory
 | |
| 	frameReader
 | |
| 
 | |
| 	wio sync.Mutex
 | |
| 	frameWriterFactory
 | |
| 
 | |
| 	frameHandler
 | |
| 	PayloadType        byte
 | |
| 	defaultCloseStatus int
 | |
| 
 | |
| 	// MaxPayloadBytes limits the size of frame payload received over Conn
 | |
| 	// by Codec's Receive method. If zero, DefaultMaxPayloadBytes is used.
 | |
| 	MaxPayloadBytes int
 | |
| }
 | |
| 
 | |
| // Read implements the io.Reader interface:
 | |
| // it reads data of a frame from the WebSocket connection.
 | |
| // if msg is not large enough for the frame data, it fills the msg and next Read
 | |
| // will read the rest of the frame data.
 | |
| // it reads Text frame or Binary frame.
 | |
| func (ws *Conn) Read(msg []byte) (n int, err error) {
 | |
| 	ws.rio.Lock()
 | |
| 	defer ws.rio.Unlock()
 | |
| again:
 | |
| 	if ws.frameReader == nil {
 | |
| 		frame, err := ws.frameReaderFactory.NewFrameReader()
 | |
| 		if err != nil {
 | |
| 			return 0, err
 | |
| 		}
 | |
| 		ws.frameReader, err = ws.frameHandler.HandleFrame(frame)
 | |
| 		if err != nil {
 | |
| 			return 0, err
 | |
| 		}
 | |
| 		if ws.frameReader == nil {
 | |
| 			goto again
 | |
| 		}
 | |
| 	}
 | |
| 	n, err = ws.frameReader.Read(msg)
 | |
| 	if err == io.EOF {
 | |
| 		if trailer := ws.frameReader.TrailerReader(); trailer != nil {
 | |
| 			io.Copy(ioutil.Discard, trailer)
 | |
| 		}
 | |
| 		ws.frameReader = nil
 | |
| 		goto again
 | |
| 	}
 | |
| 	return n, err
 | |
| }
 | |
| 
 | |
| // Write implements the io.Writer interface:
 | |
| // it writes data as a frame to the WebSocket connection.
 | |
| func (ws *Conn) Write(msg []byte) (n int, err error) {
 | |
| 	ws.wio.Lock()
 | |
| 	defer ws.wio.Unlock()
 | |
| 	w, err := ws.frameWriterFactory.NewFrameWriter(ws.PayloadType)
 | |
| 	if err != nil {
 | |
| 		return 0, err
 | |
| 	}
 | |
| 	n, err = w.Write(msg)
 | |
| 	w.Close()
 | |
| 	return n, err
 | |
| }
 | |
| 
 | |
| // Close implements the io.Closer interface.
 | |
| func (ws *Conn) Close() error {
 | |
| 	err := ws.frameHandler.WriteClose(ws.defaultCloseStatus)
 | |
| 	err1 := ws.rwc.Close()
 | |
| 	if err != nil {
 | |
| 		return err
 | |
| 	}
 | |
| 	return err1
 | |
| }
 | |
| 
 | |
| // IsClientConn reports whether ws is a client-side connection.
 | |
| func (ws *Conn) IsClientConn() bool { return ws.request == nil }
 | |
| 
 | |
| // IsServerConn reports whether ws is a server-side connection.
 | |
| func (ws *Conn) IsServerConn() bool { return ws.request != nil }
 | |
| 
 | |
| // LocalAddr returns the WebSocket Origin for the connection for client, or
 | |
| // the WebSocket location for server.
 | |
| func (ws *Conn) LocalAddr() net.Addr {
 | |
| 	if ws.IsClientConn() {
 | |
| 		return &Addr{ws.config.Origin}
 | |
| 	}
 | |
| 	return &Addr{ws.config.Location}
 | |
| }
 | |
| 
 | |
| // RemoteAddr returns the WebSocket location for the connection for client, or
 | |
| // the Websocket Origin for server.
 | |
| func (ws *Conn) RemoteAddr() net.Addr {
 | |
| 	if ws.IsClientConn() {
 | |
| 		return &Addr{ws.config.Location}
 | |
| 	}
 | |
| 	return &Addr{ws.config.Origin}
 | |
| }
 | |
| 
 | |
| var errSetDeadline = errors.New("websocket: cannot set deadline: not using a net.Conn")
 | |
| 
 | |
| // SetDeadline sets the connection's network read & write deadlines.
 | |
| func (ws *Conn) SetDeadline(t time.Time) error {
 | |
| 	if conn, ok := ws.rwc.(net.Conn); ok {
 | |
| 		return conn.SetDeadline(t)
 | |
| 	}
 | |
| 	return errSetDeadline
 | |
| }
 | |
| 
 | |
| // SetReadDeadline sets the connection's network read deadline.
 | |
| func (ws *Conn) SetReadDeadline(t time.Time) error {
 | |
| 	if conn, ok := ws.rwc.(net.Conn); ok {
 | |
| 		return conn.SetReadDeadline(t)
 | |
| 	}
 | |
| 	return errSetDeadline
 | |
| }
 | |
| 
 | |
| // SetWriteDeadline sets the connection's network write deadline.
 | |
| func (ws *Conn) SetWriteDeadline(t time.Time) error {
 | |
| 	if conn, ok := ws.rwc.(net.Conn); ok {
 | |
| 		return conn.SetWriteDeadline(t)
 | |
| 	}
 | |
| 	return errSetDeadline
 | |
| }
 | |
| 
 | |
| // Config returns the WebSocket config.
 | |
| func (ws *Conn) Config() *Config { return ws.config }
 | |
| 
 | |
| // Request returns the http request upgraded to the WebSocket.
 | |
| // It is nil for client side.
 | |
| func (ws *Conn) Request() *http.Request { return ws.request }
 | |
| 
 | |
| // Codec represents a symmetric pair of functions that implement a codec.
 | |
| type Codec struct {
 | |
| 	Marshal   func(v interface{}) (data []byte, payloadType byte, err error)
 | |
| 	Unmarshal func(data []byte, payloadType byte, v interface{}) (err error)
 | |
| }
 | |
| 
 | |
| // Send sends v marshaled by cd.Marshal as single frame to ws.
 | |
| func (cd Codec) Send(ws *Conn, v interface{}) (err error) {
 | |
| 	data, payloadType, err := cd.Marshal(v)
 | |
| 	if err != nil {
 | |
| 		return err
 | |
| 	}
 | |
| 	ws.wio.Lock()
 | |
| 	defer ws.wio.Unlock()
 | |
| 	w, err := ws.frameWriterFactory.NewFrameWriter(payloadType)
 | |
| 	if err != nil {
 | |
| 		return err
 | |
| 	}
 | |
| 	_, err = w.Write(data)
 | |
| 	w.Close()
 | |
| 	return err
 | |
| }
 | |
| 
 | |
| // Receive receives single frame from ws, unmarshaled by cd.Unmarshal and stores
 | |
| // in v. The whole frame payload is read to an in-memory buffer; max size of
 | |
| // payload is defined by ws.MaxPayloadBytes. If frame payload size exceeds
 | |
| // limit, ErrFrameTooLarge is returned; in this case frame is not read off wire
 | |
| // completely. The next call to Receive would read and discard leftover data of
 | |
| // previous oversized frame before processing next frame.
 | |
| func (cd Codec) Receive(ws *Conn, v interface{}) (err error) {
 | |
| 	ws.rio.Lock()
 | |
| 	defer ws.rio.Unlock()
 | |
| 	if ws.frameReader != nil {
 | |
| 		_, err = io.Copy(ioutil.Discard, ws.frameReader)
 | |
| 		if err != nil {
 | |
| 			return err
 | |
| 		}
 | |
| 		ws.frameReader = nil
 | |
| 	}
 | |
| again:
 | |
| 	frame, err := ws.frameReaderFactory.NewFrameReader()
 | |
| 	if err != nil {
 | |
| 		return err
 | |
| 	}
 | |
| 	frame, err = ws.frameHandler.HandleFrame(frame)
 | |
| 	if err != nil {
 | |
| 		return err
 | |
| 	}
 | |
| 	if frame == nil {
 | |
| 		goto again
 | |
| 	}
 | |
| 	maxPayloadBytes := ws.MaxPayloadBytes
 | |
| 	if maxPayloadBytes == 0 {
 | |
| 		maxPayloadBytes = DefaultMaxPayloadBytes
 | |
| 	}
 | |
| 	if hf, ok := frame.(*hybiFrameReader); ok && hf.header.Length > int64(maxPayloadBytes) {
 | |
| 		// payload size exceeds limit, no need to call Unmarshal
 | |
| 		//
 | |
| 		// set frameReader to current oversized frame so that
 | |
| 		// the next call to this function can drain leftover
 | |
| 		// data before processing the next frame
 | |
| 		ws.frameReader = frame
 | |
| 		return ErrFrameTooLarge
 | |
| 	}
 | |
| 	payloadType := frame.PayloadType()
 | |
| 	data, err := ioutil.ReadAll(frame)
 | |
| 	if err != nil {
 | |
| 		return err
 | |
| 	}
 | |
| 	return cd.Unmarshal(data, payloadType, v)
 | |
| }
 | |
| 
 | |
| func marshal(v interface{}) (msg []byte, payloadType byte, err error) {
 | |
| 	switch data := v.(type) {
 | |
| 	case string:
 | |
| 		return []byte(data), TextFrame, nil
 | |
| 	case []byte:
 | |
| 		return data, BinaryFrame, nil
 | |
| 	}
 | |
| 	return nil, UnknownFrame, ErrNotSupported
 | |
| }
 | |
| 
 | |
| func unmarshal(msg []byte, payloadType byte, v interface{}) (err error) {
 | |
| 	switch data := v.(type) {
 | |
| 	case *string:
 | |
| 		*data = string(msg)
 | |
| 		return nil
 | |
| 	case *[]byte:
 | |
| 		*data = msg
 | |
| 		return nil
 | |
| 	}
 | |
| 	return ErrNotSupported
 | |
| }
 | |
| 
 | |
| /*
 | |
| Message is a codec to send/receive text/binary data in a frame on WebSocket connection.
 | |
| To send/receive text frame, use string type.
 | |
| To send/receive binary frame, use []byte type.
 | |
| 
 | |
| Trivial usage:
 | |
| 
 | |
| 	import "websocket"
 | |
| 
 | |
| 	// receive text frame
 | |
| 	var message string
 | |
| 	websocket.Message.Receive(ws, &message)
 | |
| 
 | |
| 	// send text frame
 | |
| 	message = "hello"
 | |
| 	websocket.Message.Send(ws, message)
 | |
| 
 | |
| 	// receive binary frame
 | |
| 	var data []byte
 | |
| 	websocket.Message.Receive(ws, &data)
 | |
| 
 | |
| 	// send binary frame
 | |
| 	data = []byte{0, 1, 2}
 | |
| 	websocket.Message.Send(ws, data)
 | |
| 
 | |
| */
 | |
| var Message = Codec{marshal, unmarshal}
 | |
| 
 | |
| func jsonMarshal(v interface{}) (msg []byte, payloadType byte, err error) {
 | |
| 	msg, err = json.Marshal(v)
 | |
| 	return msg, TextFrame, err
 | |
| }
 | |
| 
 | |
| func jsonUnmarshal(msg []byte, payloadType byte, v interface{}) (err error) {
 | |
| 	return json.Unmarshal(msg, v)
 | |
| }
 | |
| 
 | |
| /*
 | |
| JSON is a codec to send/receive JSON data in a frame from a WebSocket connection.
 | |
| 
 | |
| Trivial usage:
 | |
| 
 | |
| 	import "websocket"
 | |
| 
 | |
| 	type T struct {
 | |
| 		Msg string
 | |
| 		Count int
 | |
| 	}
 | |
| 
 | |
| 	// receive JSON type T
 | |
| 	var data T
 | |
| 	websocket.JSON.Receive(ws, &data)
 | |
| 
 | |
| 	// send JSON type T
 | |
| 	websocket.JSON.Send(ws, data)
 | |
| */
 | |
| var JSON = Codec{jsonMarshal, jsonUnmarshal}
 |