vendor/github.com/gorilla/websocket/client.go

   1 // Copyright 2013 The Gorilla WebSocket 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
   6
   7 import (
   8         "bytes"
   9         "context"
  10         "crypto/tls"
  11         "errors"
  12         "io"
  13         "io/ioutil"
  14         "net"
  15         "net/http"
  16         "net/http/httptrace"
  17         "net/url"
  18         "strings"
  19         "time"
  20 )
  21
  22 // ErrBadHandshake is returned when the server response to opening handshake is
  23 // invalid.
  24 var ErrBadHandshake = errors.New("websocket: bad handshake")
  25
  26 var errInvalidCompression = errors.New("websocket: invalid compression negotiation")
  27
  28 // NewClient creates a new client connection using the given net connection.
  29 // The URL u specifies the host and request URI. Use requestHeader to specify
  30 // the origin (Origin), subprotocols (Sec-WebSocket-Protocol) and cookies
  31 // (Cookie). Use the response.Header to get the selected subprotocol
  32 // (Sec-WebSocket-Protocol) and cookies (Set-Cookie).
  33 //
  34 // If the WebSocket handshake fails, ErrBadHandshake is returned along with a
  35 // non-nil *http.Response so that callers can handle redirects, authentication,
  36 // etc.
  37 //
  38 // Deprecated: Use Dialer instead.
  39 func NewClient(netConn net.Conn, u *url.URL, requestHeader http.Header, readBufSize, writeBufSize int) (c *Conn, response *http.Response, err error) {
  40         d := Dialer{
  41                 ReadBufferSize:  readBufSize,
  42                 WriteBufferSize: writeBufSize,
  43                 NetDial: func(net, addr string) (net.Conn, error) {
  44                         return netConn, nil
  45                 },
  46         }
  47         return d.Dial(u.String(), requestHeader)
  48 }
  49
  50 // A Dialer contains options for connecting to WebSocket server.
  51 type Dialer struct {
  52         // NetDial specifies the dial function for creating TCP connections. If
  53         // NetDial is nil, net.Dial is used.
  54         NetDial func(network, addr string) (net.Conn, error)
  55
  56         // NetDialContext specifies the dial function for creating TCP connections. If
  57         // NetDialContext is nil, net.DialContext is used.
  58         NetDialContext func(ctx context.Context, network, addr string) (net.Conn, error)
  59
  60         // Proxy specifies a function to return a proxy for a given
  61         // Request. If the function returns a non-nil error, the
  62         // request is aborted with the provided error.
  63         // If Proxy is nil or returns a nil *URL, no proxy is used.
  64         Proxy func(*http.Request) (*url.URL, error)
  65
  66         // TLSClientConfig specifies the TLS configuration to use with tls.Client.
  67         // If nil, the default configuration is used.
  68         TLSClientConfig *tls.Config
  69
  70         // HandshakeTimeout specifies the duration for the handshake to complete.
  71         HandshakeTimeout time.Duration
  72
  73         // ReadBufferSize and WriteBufferSize specify I/O buffer sizes in bytes. If a buffer
  74         // size is zero, then a useful default size is used. The I/O buffer sizes
  75         // do not limit the size of the messages that can be sent or received.
  76         ReadBufferSize, WriteBufferSize int
  77
  78         // WriteBufferPool is a pool of buffers for write operations. If the value
  79         // is not set, then write buffers are allocated to the connection for the
  80         // lifetime of the connection.
  81         //
  82         // A pool is most useful when the application has a modest volume of writes
  83         // across a large number of connections.
  84         //
  85         // Applications should use a single pool for each unique value of
  86         // WriteBufferSize.
  87         WriteBufferPool BufferPool
  88
  89         // Subprotocols specifies the client's requested subprotocols.
  90         Subprotocols []string
  91
  92         // EnableCompression specifies if the client should attempt to negotiate
  93         // per message compression (RFC 7692). Setting this value to true does not
  94         // guarantee that compression will be supported. Currently only "no context
  95         // takeover" modes are supported.
  96         EnableCompression bool
  97
  98         // Jar specifies the cookie jar.
  99         // If Jar is nil, cookies are not sent in requests and ignored
 100         // in responses.
 101         Jar http.CookieJar
 102 }
 103
 104 // Dial creates a new client connection by calling DialContext with a background context.
 105 func (d *Dialer) Dial(urlStr string, requestHeader http.Header) (*Conn, *http.Response, error) {
 106         return d.DialContext(context.Background(), urlStr, requestHeader)
 107 }
 108
 109 var errMalformedURL = errors.New("malformed ws or wss URL")
 110
 111 func hostPortNoPort(u *url.URL) (hostPort, hostNoPort string) {
 112         hostPort = u.Host
 113         hostNoPort = u.Host
 114         if i := strings.LastIndex(u.Host, ":"); i > strings.LastIndex(u.Host, "]") {
 115                 hostNoPort = hostNoPort[:i]
 116         } else {
 117                 switch u.Scheme {
 118                 case "wss":
 119                         hostPort += ":443"
 120                 case "https":
 121                         hostPort += ":443"
 122                 default:
 123                         hostPort += ":80"
 124                 }
 125         }
 126         return hostPort, hostNoPort
 127 }
 128
 129 // DefaultDialer is a dialer with all fields set to the default values.
 130 var DefaultDialer = &Dialer{
 131         Proxy:            http.ProxyFromEnvironment,
 132         HandshakeTimeout: 45 * time.Second,
 133 }
 134
 135 // nilDialer is dialer to use when receiver is nil.
 136 var nilDialer = *DefaultDialer
 137
 138 // DialContext creates a new client connection. Use requestHeader to specify the
 139 // origin (Origin), subprotocols (Sec-WebSocket-Protocol) and cookies (Cookie).
 140 // Use the response.Header to get the selected subprotocol
 141 // (Sec-WebSocket-Protocol) and cookies (Set-Cookie).
 142 //
 143 // The context will be used in the request and in the Dialer.
 144 //
 145 // If the WebSocket handshake fails, ErrBadHandshake is returned along with a
 146 // non-nil *http.Response so that callers can handle redirects, authentication,
 147 // etcetera. The response body may not contain the entire response and does not
 148 // need to be closed by the application.
 149 func (d *Dialer) DialContext(ctx context.Context, urlStr string, requestHeader http.Header) (*Conn, *http.Response, error) {
 150         if d == nil {
 151                 d = &nilDialer
 152         }
 153
 154         challengeKey, err := generateChallengeKey()
 155         if err != nil {
 156                 return nil, nil, err
 157         }
 158
 159         u, err := url.Parse(urlStr)
 160         if err != nil {
 161                 return nil, nil, err
 162         }
 163
 164         switch u.Scheme {
 165         case "ws":
 166                 u.Scheme = "http"
 167         case "wss":
 168                 u.Scheme = "https"
 169         default:
 170                 return nil, nil, errMalformedURL
 171         }
 172
 173         if u.User != nil {
 174                 // User name and password are not allowed in websocket URIs.
 175                 return nil, nil, errMalformedURL
 176         }
 177
 178         req := &http.Request{
 179                 Method:     "GET",
 180                 URL:        u,
 181                 Proto:      "HTTP/1.1",
 182                 ProtoMajor: 1,
 183                 ProtoMinor: 1,
 184                 Header:     make(http.Header),
 185                 Host:       u.Host,
 186         }
 187         req = req.WithContext(ctx)
 188
 189         // Set the cookies present in the cookie jar of the dialer
 190         if d.Jar != nil {
 191                 for _, cookie := range d.Jar.Cookies(u) {
 192                         req.AddCookie(cookie)
 193                 }
 194         }
 195
 196         // Set the request headers using the capitalization for names and values in
 197         // RFC examples. Although the capitalization shouldn't matter, there are
 198         // servers that depend on it. The Header.Set method is not used because the
 199         // method canonicalizes the header names.
 200         req.Header["Upgrade"] = []string{"websocket"}
 201         req.Header["Connection"] = []string{"Upgrade"}
 202         req.Header["Sec-WebSocket-Key"] = []string{challengeKey}
 203         req.Header["Sec-WebSocket-Version"] = []string{"13"}
 204         if len(d.Subprotocols) > 0 {
 205                 req.Header["Sec-WebSocket-Protocol"] = []string{strings.Join(d.Subprotocols, ", ")}
 206         }
 207         for k, vs := range requestHeader {
 208                 switch {
 209                 case k == "Host":
 210                         if len(vs) > 0 {
 211                                 req.Host = vs[0]
 212                         }
 213                 case k == "Upgrade" ||
 214                         k == "Connection" ||
 215                         k == "Sec-Websocket-Key" ||
 216                         k == "Sec-Websocket-Version" ||
 217                         k == "Sec-Websocket-Extensions" ||
 218                         (k == "Sec-Websocket-Protocol" && len(d.Subprotocols) > 0):
 219                         return nil, nil, errors.New("websocket: duplicate header not allowed: " + k)
 220                 case k == "Sec-Websocket-Protocol":
 221                         req.Header["Sec-WebSocket-Protocol"] = vs
 222                 default:
 223                         req.Header[k] = vs
 224                 }
 225         }
 226
 227         if d.EnableCompression {
 228                 req.Header["Sec-WebSocket-Extensions"] = []string{"permessage-deflate; server_no_context_takeover; client_no_context_takeover"}
 229         }
 230
 231         if d.HandshakeTimeout != 0 {
 232                 var cancel func()
 233                 ctx, cancel = context.WithTimeout(ctx, d.HandshakeTimeout)
 234                 defer cancel()
 235         }
 236
 237         // Get network dial function.
 238         var netDial func(network, add string) (net.Conn, error)
 239
 240         if d.NetDialContext != nil {
 241                 netDial = func(network, addr string) (net.Conn, error) {
 242                         return d.NetDialContext(ctx, network, addr)
 243                 }
 244         } else if d.NetDial != nil {
 245                 netDial = d.NetDial
 246         } else {
 247                 netDialer := &net.Dialer{}
 248                 netDial = func(network, addr string) (net.Conn, error) {
 249                         return netDialer.DialContext(ctx, network, addr)
 250                 }
 251         }
 252
 253         // If needed, wrap the dial function to set the connection deadline.
 254         if deadline, ok := ctx.Deadline(); ok {
 255                 forwardDial := netDial
 256                 netDial = func(network, addr string) (net.Conn, error) {
 257                         c, err := forwardDial(network, addr)
 258                         if err != nil {
 259                                 return nil, err
 260                         }
 261                         err = c.SetDeadline(deadline)
 262                         if err != nil {
 263                                 c.Close()
 264                                 return nil, err
 265                         }
 266                         return c, nil
 267                 }
 268         }
 269
 270         // If needed, wrap the dial function to connect through a proxy.
 271         if d.Proxy != nil {
 272                 proxyURL, err := d.Proxy(req)
 273                 if err != nil {
 274                         return nil, nil, err
 275                 }
 276                 if proxyURL != nil {
 277                         dialer, err := proxy_FromURL(proxyURL, netDialerFunc(netDial))
 278                         if err != nil {
 279                                 return nil, nil, err
 280                         }
 281                         netDial = dialer.Dial
 282                 }
 283         }
 284
 285         hostPort, hostNoPort := hostPortNoPort(u)
 286         trace := httptrace.ContextClientTrace(ctx)
 287         if trace != nil && trace.GetConn != nil {
 288                 trace.GetConn(hostPort)
 289         }
 290
 291         netConn, err := netDial("tcp", hostPort)
 292         if trace != nil && trace.GotConn != nil {
 293                 trace.GotConn(httptrace.GotConnInfo{
 294                         Conn: netConn,
 295                 })
 296         }
 297         if err != nil {
 298                 return nil, nil, err
 299         }
 300
 301         defer func() {
 302                 if netConn != nil {
 303                         netConn.Close()
 304                 }
 305         }()
 306
 307         if u.Scheme == "https" {
 308                 cfg := cloneTLSConfig(d.TLSClientConfig)
 309                 if cfg.ServerName == "" {
 310                         cfg.ServerName = hostNoPort
 311                 }
 312                 tlsConn := tls.Client(netConn, cfg)
 313                 netConn = tlsConn
 314
 315                 var err error
 316                 if trace != nil {
 317                         err = doHandshakeWithTrace(trace, tlsConn, cfg)
 318                 } else {
 319                         err = doHandshake(tlsConn, cfg)
 320                 }
 321
 322                 if err != nil {
 323                         return nil, nil, err
 324                 }
 325         }
 326
 327         conn := newConn(netConn, false, d.ReadBufferSize, d.WriteBufferSize, d.WriteBufferPool, nil, nil)
 328
 329         if err := req.Write(netConn); err != nil {
 330                 return nil, nil, err
 331         }
 332
 333         if trace != nil && trace.GotFirstResponseByte != nil {
 334                 if peek, err := conn.br.Peek(1); err == nil && len(peek) == 1 {
 335                         trace.GotFirstResponseByte()
 336                 }
 337         }
 338
 339         resp, err := http.ReadResponse(conn.br, req)
 340         if err != nil {
 341                 return nil, nil, err
 342         }
 343
 344         if d.Jar != nil {
 345                 if rc := resp.Cookies(); len(rc) > 0 {
 346                         d.Jar.SetCookies(u, rc)
 347                 }
 348         }
 349
 350         if resp.StatusCode != 101 ||
 351                 !strings.EqualFold(resp.Header.Get("Upgrade"), "websocket") ||
 352                 !strings.EqualFold(resp.Header.Get("Connection"), "upgrade") ||
 353                 resp.Header.Get("Sec-Websocket-Accept") != computeAcceptKey(challengeKey) {
 354                 // Before closing the network connection on return from this
 355                 // function, slurp up some of the response to aid application
 356                 // debugging.
 357                 buf := make([]byte, 1024)
 358                 n, _ := io.ReadFull(resp.Body, buf)
 359                 resp.Body = ioutil.NopCloser(bytes.NewReader(buf[:n]))
 360                 return nil, resp, ErrBadHandshake
 361         }
 362
 363         for _, ext := range parseExtensions(resp.Header) {
 364                 if ext[""] != "permessage-deflate" {
 365                         continue
 366                 }
 367                 _, snct := ext["server_no_context_takeover"]
 368                 _, cnct := ext["client_no_context_takeover"]
 369                 if !snct || !cnct {
 370                         return nil, resp, errInvalidCompression
 371                 }
 372                 conn.newCompressionWriter = compressNoContextTakeover
 373                 conn.newDecompressionReader = decompressNoContextTakeover
 374                 break
 375         }
 376
 377         resp.Body = ioutil.NopCloser(bytes.NewReader([]byte{}))
 378         conn.subprotocol = resp.Header.Get("Sec-Websocket-Protocol")
 379
 380         netConn.SetDeadline(time.Time{})
 381         netConn = nil // to avoid close in defer.
 382         return conn, resp, nil
 383 }
 384
 385 func doHandshake(tlsConn *tls.Conn, cfg *tls.Config) error {
 386         if err := tlsConn.Handshake(); err != nil {
 387                 return err
 388         }
 389         if !cfg.InsecureSkipVerify {
 390                 if err := tlsConn.VerifyHostname(cfg.ServerName); err != nil {
 391                         return err
 392                 }
 393         }
 394         return nil
 395 }