libgo/go/websocket/websocket.go

   1 // Copyright 2009 The Go Authors. All rights reserved.
   2 // Use of this source code is governed by a BSD-style
   3 // license that can be found in the LICENSE file.
   4
   5 // Package websocket implements a client and server for the WebSocket protocol.
   6 // The protocol is defined at http://tools.ietf.org/html/draft-ietf-hybi-thewebsocketprotocol
   7 package websocket
   8
   9 import (
  10         "bufio"
  11         "crypto/tls"
  12         "http"
  13         "io"
  14         "io/ioutil"
  15         "json"
  16         "net"
  17         "os"
  18         "sync"
  19         "url"
  20 )
  21
  22 const (
  23         ProtocolVersionHixie75   = -75
  24         ProtocolVersionHixie76   = -76
  25         ProtocolVersionHybi00    = 0
  26         ProtocolVersionHybi08    = 8
  27         ProtocolVersionHybi13    = 13
  28         ProtocolVersionHybi      = ProtocolVersionHybi13
  29         SupportedProtocolVersion = "13, 8"
  30
  31         ContinuationFrame = 0
  32         TextFrame         = 1
  33         BinaryFrame       = 2
  34         CloseFrame        = 8
  35         PingFrame         = 9
  36         PongFrame         = 10
  37         UnknownFrame      = 255
  38 )
  39
  40 // WebSocket protocol errors.
  41 type ProtocolError struct {
  42         ErrorString string
  43 }
  44
  45 func (err *ProtocolError) String() string { return err.ErrorString }
  46
  47 var (
  48         ErrBadProtocolVersion   = &ProtocolError{"bad protocol version"}
  49         ErrBadScheme            = &ProtocolError{"bad scheme"}
  50         ErrBadStatus            = &ProtocolError{"bad status"}
  51         ErrBadUpgrade           = &ProtocolError{"missing or bad upgrade"}
  52         ErrBadWebSocketOrigin   = &ProtocolError{"missing or bad WebSocket-Origin"}
  53         ErrBadWebSocketLocation = &ProtocolError{"missing or bad WebSocket-Location"}
  54         ErrBadWebSocketProtocol = &ProtocolError{"missing or bad WebSocket-Protocol"}
  55         ErrBadWebSocketVersion  = &ProtocolError{"missing or bad WebSocket Version"}
  56         ErrChallengeResponse    = &ProtocolError{"mismatch challenge/response"}
  57         ErrBadFrame             = &ProtocolError{"bad frame"}
  58         ErrBadFrameBoundary     = &ProtocolError{"not on frame boundary"}
  59         ErrNotWebSocket         = &ProtocolError{"not websocket protocol"}
  60         ErrBadRequestMethod     = &ProtocolError{"bad method"}
  61         ErrNotSupported         = &ProtocolError{"not supported"}
  62 )
  63
  64 // Addr is an implementation of net.Addr for WebSocket.
  65 type Addr struct {
  66         *url.URL
  67 }
  68
  69 // Network returns the network type for a WebSocket, "websocket".
  70 func (addr *Addr) Network() string { return "websocket" }
  71
  72 // Config is a WebSocket configuration
  73 type Config struct {
  74         // A WebSocket server address.
  75         Location *url.URL
  76
  77         // A Websocket client origin.
  78         Origin *url.URL
  79
  80         // WebSocket subprotocols.
  81         Protocol []string
  82
  83         // WebSocket protocol version.
  84         Version int
  85
  86         // TLS config for secure WebSocket (wss).
  87         TlsConfig *tls.Config
  88
  89         handshakeData map[string]string
  90 }
  91
  92 // serverHandshaker is an interface to handle WebSocket server side handshake.
  93 type serverHandshaker interface {
  94         // ReadHandshake reads handshake request message from client.
  95         // Returns http response code and error if any.
  96         ReadHandshake(buf *bufio.Reader, req *http.Request) (code int, err os.Error)
  97
  98         // AcceptHandshake accepts the client handshake request and sends
  99         // handshake response back to client.
 100         AcceptHandshake(buf *bufio.Writer) (err os.Error)
 101
 102         // NewServerConn creates a new WebSocket connection.
 103         NewServerConn(buf *bufio.ReadWriter, rwc io.ReadWriteCloser, request *http.Request) (conn *Conn)
 104 }
 105
 106 // frameReader is an interface to read a WebSocket frame.
 107 type frameReader interface {
 108         // Reader is to read payload of the frame.
 109         io.Reader
 110
 111         // PayloadType returns payload type.
 112         PayloadType() byte
 113
 114         // HeaderReader returns a reader to read header of the frame.
 115         HeaderReader() io.Reader
 116
 117         // TrailerReader returns a reader to read trailer of the frame.
 118         // If it returns nil, there is no trailer in the frame.
 119         TrailerReader() io.Reader
 120
 121         // Len returns total length of the frame, including header and trailer.
 122         Len() int
 123 }
 124
 125 // frameReaderFactory is an interface to creates new frame reader.
 126 type frameReaderFactory interface {
 127         NewFrameReader() (r frameReader, err os.Error)
 128 }
 129
 130 // frameWriter is an interface to write a WebSocket frame.
 131 type frameWriter interface {
 132         // Writer is to write playload of the frame.
 133         io.WriteCloser
 134 }
 135
 136 // frameWriterFactory is an interface to create new frame writer.
 137 type frameWriterFactory interface {
 138         NewFrameWriter(payloadType byte) (w frameWriter, err os.Error)
 139 }
 140
 141 type frameHandler interface {
 142         HandleFrame(frame frameReader) (r frameReader, err os.Error)
 143         WriteClose(status int) (err os.Error)
 144 }
 145
 146 // Conn represents a WebSocket connection.
 147 type Conn struct {
 148         config  *Config
 149         request *http.Request
 150
 151         buf *bufio.ReadWriter
 152         rwc io.ReadWriteCloser
 153
 154         rio sync.Mutex
 155         frameReaderFactory
 156         frameReader
 157
 158         wio sync.Mutex
 159         frameWriterFactory
 160
 161         frameHandler
 162         PayloadType        byte
 163         defaultCloseStatus int
 164 }
 165
 166 // Read implements the io.Reader interface:
 167 // it reads data of a frame from the WebSocket connection.
 168 // if msg is not large enough for the frame data, it fills the msg and next Read
 169 // will read the rest of the frame data.
 170 // it reads Text frame or Binary frame.
 171 func (ws *Conn) Read(msg []byte) (n int, err os.Error) {
 172         ws.rio.Lock()
 173         defer ws.rio.Unlock()
 174 again:
 175         if ws.frameReader == nil {
 176                 frame, err := ws.frameReaderFactory.NewFrameReader()
 177                 if err != nil {
 178                         return 0, err
 179                 }
 180                 ws.frameReader, err = ws.frameHandler.HandleFrame(frame)
 181                 if err != nil {
 182                         return 0, err
 183                 }
 184                 if ws.frameReader == nil {
 185                         goto again
 186                 }
 187         }
 188         n, err = ws.frameReader.Read(msg)
 189         if err == os.EOF {
 190                 if trailer := ws.frameReader.TrailerReader(); trailer != nil {
 191                         io.Copy(ioutil.Discard, trailer)
 192                 }
 193                 ws.frameReader = nil
 194                 goto again
 195         }
 196         return n, err
 197 }
 198
 199 // Write implements the io.Writer interface:
 200 // it writes data as a frame to the WebSocket connection.
 201 func (ws *Conn) Write(msg []byte) (n int, err os.Error) {
 202         ws.wio.Lock()
 203         defer ws.wio.Unlock()
 204         w, err := ws.frameWriterFactory.NewFrameWriter(ws.PayloadType)
 205         if err != nil {
 206                 return 0, err
 207         }
 208         n, err = w.Write(msg)
 209         w.Close()
 210         if err != nil {
 211                 return n, err
 212         }
 213         return n, err
 214 }
 215
 216 // Close implements the io.Closer interface.
 217 func (ws *Conn) Close() os.Error {
 218         err := ws.frameHandler.WriteClose(ws.defaultCloseStatus)
 219         if err != nil {
 220                 return err
 221         }
 222         return ws.rwc.Close()
 223 }
 224
 225 func (ws *Conn) IsClientConn() bool { return ws.request == nil }
 226 func (ws *Conn) IsServerConn() bool { return ws.request != nil }
 227
 228 // LocalAddr returns the WebSocket Origin for the connection for client, or
 229 // the WebSocket location for server.
 230 func (ws *Conn) LocalAddr() net.Addr {
 231         if ws.IsClientConn() {
 232                 return &Addr{ws.config.Origin}
 233         }
 234         return &Addr{ws.config.Location}
 235 }
 236
 237 // RemoteAddr returns the WebSocket location for the connection for client, or
 238 // the Websocket Origin for server.
 239 func (ws *Conn) RemoteAddr() net.Addr {
 240         if ws.IsClientConn() {
 241                 return &Addr{ws.config.Location}
 242         }
 243         return &Addr{ws.config.Origin}
 244 }
 245
 246 // SetTimeout sets the connection's network timeout in nanoseconds.
 247 func (ws *Conn) SetTimeout(nsec int64) os.Error {
 248         if conn, ok := ws.rwc.(net.Conn); ok {
 249                 return conn.SetTimeout(nsec)
 250         }
 251         return os.EINVAL
 252 }
 253
 254 // SetReadTimeout sets the connection's network read timeout in nanoseconds.
 255 func (ws *Conn) SetReadTimeout(nsec int64) os.Error {
 256         if conn, ok := ws.rwc.(net.Conn); ok {
 257                 return conn.SetReadTimeout(nsec)
 258         }
 259         return os.EINVAL
 260 }
 261
 262 // SetWriteTimeout sets the connection's network write timeout in nanoseconds.
 263 func (ws *Conn) SetWriteTimeout(nsec int64) os.Error {
 264         if conn, ok := ws.rwc.(net.Conn); ok {
 265                 return conn.SetWriteTimeout(nsec)
 266         }
 267         return os.EINVAL
 268 }
 269
 270 // Config returns the WebSocket config.
 271 func (ws *Conn) Config() *Config { return ws.config }
 272
 273 // Request returns the http request upgraded to the WebSocket.
 274 // It is nil for client side.
 275 func (ws *Conn) Request() *http.Request { return ws.request }
 276
 277 // Codec represents a symmetric pair of functions that implement a codec.
 278 type Codec struct {
 279         Marshal   func(v interface{}) (data []byte, payloadType byte, err os.Error)
 280         Unmarshal func(data []byte, payloadType byte, v interface{}) (err os.Error)
 281 }
 282
 283 // Send sends v marshaled by cd.Marshal as single frame to ws.
 284 func (cd Codec) Send(ws *Conn, v interface{}) (err os.Error) {
 285         if err != nil {
 286                 return err
 287         }
 288         data, payloadType, err := cd.Marshal(v)
 289         if err != nil {
 290                 return err
 291         }
 292         ws.wio.Lock()
 293         defer ws.wio.Unlock()
 294         w, err := ws.frameWriterFactory.NewFrameWriter(payloadType)
 295         _, err = w.Write(data)
 296         w.Close()
 297         return err
 298 }
 299
 300 // Receive receives single frame from ws, unmarshaled by cd.Unmarshal and stores in v.
 301 func (cd Codec) Receive(ws *Conn, v interface{}) (err os.Error) {
 302         ws.rio.Lock()
 303         defer ws.rio.Unlock()
 304         if ws.frameReader != nil {
 305                 _, err = io.Copy(ioutil.Discard, ws.frameReader)
 306                 if err != nil {
 307                         return err
 308                 }
 309                 ws.frameReader = nil
 310         }
 311 again:
 312         frame, err := ws.frameReaderFactory.NewFrameReader()
 313         if err != nil {
 314                 return err
 315         }
 316         frame, err = ws.frameHandler.HandleFrame(frame)
 317         if err != nil {
 318                 return err
 319         }
 320         if frame == nil {
 321                 goto again
 322         }
 323         payloadType := frame.PayloadType()
 324         data, err := ioutil.ReadAll(frame)
 325         if err != nil {
 326                 return err
 327         }
 328         return cd.Unmarshal(data, payloadType, v)
 329 }
 330
 331 func marshal(v interface{}) (msg []byte, payloadType byte, err os.Error) {
 332         switch data := v.(type) {
 333         case string:
 334                 return []byte(data), TextFrame, nil
 335         case []byte:
 336                 return data, BinaryFrame, nil
 337         }
 338         return nil, UnknownFrame, ErrNotSupported
 339 }
 340
 341 func unmarshal(msg []byte, payloadType byte, v interface{}) (err os.Error) {
 342         switch data := v.(type) {
 343         case *string:
 344                 *data = string(msg)
 345                 return nil
 346         case *[]byte:
 347                 *data = msg
 348                 return nil
 349         }
 350         return ErrNotSupported
 351 }
 352
 353 /*
 354 Message is a codec to send/receive text/binary data in a frame on WebSocket connection.
 355 To send/receive text frame, use string type.
 356 To send/receive binary frame, use []byte type.
 357
 358 Trivial usage:
 359
 360         import "websocket"
 361
 362         // receive text frame
 363         var message string
 364         websocket.Message.Receive(ws, &message)
 365
 366         // send text frame
 367         message = "hello"
 368         websocket.Message.Send(ws, message)
 369
 370         // receive binary frame
 371         var data []byte
 372         websocket.Message.Receive(ws, &data)
 373
 374         // send binary frame
 375         data = []byte{0, 1, 2}
 376         websocket.Message.Send(ws, data)
 377
 378 */
 379 var Message = Codec{marshal, unmarshal}
 380
 381 func jsonMarshal(v interface{}) (msg []byte, payloadType byte, err os.Error) {
 382         msg, err = json.Marshal(v)
 383         return msg, TextFrame, err
 384 }
 385
 386 func jsonUnmarshal(msg []byte, payloadType byte, v interface{}) (err os.Error) {
 387         return json.Unmarshal(msg, v)
 388 }
 389
 390 /*
 391 JSON is a codec to send/receive JSON data in a frame from a WebSocket connection.
 392
 393 Trival usage:
 394
 395         import "websocket"
 396
 397         type T struct {
 398                 Msg string
 399                 Count int
 400         }
 401
 402         // receive JSON type T
 403         var data T
 404         websocket.JSON.Receive(ws, &data)
 405
 406         // send JSON type T
 407         websocket.JSON.Send(ws, data)
 408 */
 409 var JSON = Codec{jsonMarshal, jsonUnmarshal}