github.com/Uhtred009/v2ray-core-1@v4.31.2+incompatible/transport/internet/websocket/connection.go (about)

     1  // +build !confonly
     2  
     3  package websocket
     4  
     5  import (
     6  	"io"
     7  	"net"
     8  	"time"
     9  
    10  	"github.com/gorilla/websocket"
    11  	"v2ray.com/core/common/buf"
    12  	"v2ray.com/core/common/errors"
    13  	"v2ray.com/core/common/serial"
    14  )
    15  
    16  var (
    17  	_ buf.Writer = (*connection)(nil)
    18  )
    19  
    20  // connection is a wrapper for net.Conn over WebSocket connection.
    21  type connection struct {
    22  	conn       *websocket.Conn
    23  	reader     io.Reader
    24  	remoteAddr net.Addr
    25  }
    26  
    27  func newConnection(conn *websocket.Conn, remoteAddr net.Addr) *connection {
    28  	return &connection{
    29  		conn:       conn,
    30  		remoteAddr: remoteAddr,
    31  	}
    32  }
    33  
    34  // Read implements net.Conn.Read()
    35  func (c *connection) Read(b []byte) (int, error) {
    36  	for {
    37  		reader, err := c.getReader()
    38  		if err != nil {
    39  			return 0, err
    40  		}
    41  
    42  		nBytes, err := reader.Read(b)
    43  		if errors.Cause(err) == io.EOF {
    44  			c.reader = nil
    45  			continue
    46  		}
    47  		return nBytes, err
    48  	}
    49  }
    50  
    51  func (c *connection) getReader() (io.Reader, error) {
    52  	if c.reader != nil {
    53  		return c.reader, nil
    54  	}
    55  
    56  	_, reader, err := c.conn.NextReader()
    57  	if err != nil {
    58  		return nil, err
    59  	}
    60  	c.reader = reader
    61  	return reader, nil
    62  }
    63  
    64  // Write implements io.Writer.
    65  func (c *connection) Write(b []byte) (int, error) {
    66  	if err := c.conn.WriteMessage(websocket.BinaryMessage, b); err != nil {
    67  		return 0, err
    68  	}
    69  	return len(b), nil
    70  }
    71  
    72  func (c *connection) WriteMultiBuffer(mb buf.MultiBuffer) error {
    73  	mb = buf.Compact(mb)
    74  	mb, err := buf.WriteMultiBuffer(c, mb)
    75  	buf.ReleaseMulti(mb)
    76  	return err
    77  }
    78  
    79  func (c *connection) Close() error {
    80  	var errors []interface{}
    81  	if err := c.conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), time.Now().Add(time.Second*5)); err != nil {
    82  		errors = append(errors, err)
    83  	}
    84  	if err := c.conn.Close(); err != nil {
    85  		errors = append(errors, err)
    86  	}
    87  	if len(errors) > 0 {
    88  		return newError("failed to close connection").Base(newError(serial.Concat(errors...)))
    89  	}
    90  	return nil
    91  }
    92  
    93  func (c *connection) LocalAddr() net.Addr {
    94  	return c.conn.LocalAddr()
    95  }
    96  
    97  func (c *connection) RemoteAddr() net.Addr {
    98  	return c.remoteAddr
    99  }
   100  
   101  func (c *connection) SetDeadline(t time.Time) error {
   102  	if err := c.SetReadDeadline(t); err != nil {
   103  		return err
   104  	}
   105  	return c.SetWriteDeadline(t)
   106  }
   107  
   108  func (c *connection) SetReadDeadline(t time.Time) error {
   109  	return c.conn.SetReadDeadline(t)
   110  }
   111  
   112  func (c *connection) SetWriteDeadline(t time.Time) error {
   113  	return c.conn.SetWriteDeadline(t)
   114  }