github.com/moqsien/xraycore@v1.8.5/transport/internet/websocket/connection.go (about)

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