connection.go

Documentation: github.com/moby/spdystream/ws

     1  /*
     2     Copyright 2014-2021 Docker Inc.
     3  
     4     Licensed under the Apache License, Version 2.0 (the "License");
     5     you may not use this file except in compliance with the License.
     6     You may obtain a copy of the License at
     7  
     8         http://www.apache.org/licenses/LICENSE-2.0
     9  
    10     Unless required by applicable law or agreed to in writing, software
    11     distributed under the License is distributed on an "AS IS" BASIS,
    12     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13     See the License for the specific language governing permissions and
    14     limitations under the License.
    15  */
    16  
    17  package ws
    18  
    19  import (
    20  	"fmt"
    21  	"io"
    22  	"time"
    23  
    24  	"github.com/gorilla/websocket"
    25  )
    26  
    27  // Wrap an HTTP2 connection over WebSockets and
    28  // use the underlying WebSocket framing for proxy
    29  // compatibility.
    30  type Conn struct {
    31  	*websocket.Conn
    32  	reader io.Reader
    33  }
    34  
    35  func NewConnection(w *websocket.Conn) *Conn {
    36  	return &Conn{Conn: w}
    37  }
    38  
    39  func (c *Conn) Write(b []byte) (int, error) {
    40  	err := c.WriteMessage(websocket.BinaryMessage, b)
    41  	if err != nil {
    42  		return 0, err
    43  	}
    44  	return len(b), nil
    45  }
    46  
    47  func (c *Conn) Read(b []byte) (int, error) {
    48  	if c.reader == nil {
    49  		if err := c.nextReader(); err != nil {
    50  			return 0, err
    51  		}
    52  	}
    53  
    54  	for {
    55  		n, err := c.reader.Read(b)
    56  		if err != nil {
    57  			if err != io.EOF {
    58  				return n, err
    59  			}
    60  
    61  			// get next reader if there is no data in the current one
    62  			if err := c.nextReader(); err != nil {
    63  				return 0, err
    64  			}
    65  			continue
    66  		}
    67  		return n, nil
    68  	}
    69  }
    70  
    71  func (c *Conn) nextReader() error {
    72  	t, r, err := c.NextReader()
    73  	if err != nil {
    74  		return err
    75  	}
    76  
    77  	if t != websocket.BinaryMessage {
    78  		return fmt.Errorf("ws: non-binary message in stream")
    79  	}
    80  	c.reader = r
    81  	return nil
    82  }
    83  
    84  func (c *Conn) SetDeadline(t time.Time) error {
    85  	if err := c.Conn.SetReadDeadline(t); err != nil {
    86  		return err
    87  	}
    88  	if err := c.Conn.SetWriteDeadline(t); err != nil {
    89  		return err
    90  	}
    91  	return nil
    92  }
    93  
    94  func (c *Conn) Close() error {
    95  	err := c.Conn.Close()
    96  	return err
    97  }
    98
View as plain text