github.com/ungtb10d/cli/v2@v2.0.0-20221110210412-98537dd9d6a1/pkg/liveshare/socket.go (about)

     1  package liveshare
     2  
     3  import (
     4  	"context"
     5  	"crypto/tls"
     6  	"io"
     7  	"net"
     8  	"net/http"
     9  	"time"
    10  
    11  	"github.com/gorilla/websocket"
    12  )
    13  
    14  type socket struct {
    15  	addr      string
    16  	tlsConfig *tls.Config
    17  
    18  	conn   *websocket.Conn
    19  	reader io.Reader
    20  }
    21  
    22  func newSocket(uri string, tlsConfig *tls.Config) *socket {
    23  	return &socket{addr: uri, tlsConfig: tlsConfig}
    24  }
    25  
    26  func (s *socket) connect(ctx context.Context) error {
    27  	dialer := websocket.Dialer{
    28  		Proxy:            http.ProxyFromEnvironment,
    29  		HandshakeTimeout: 45 * time.Second,
    30  		TLSClientConfig:  s.tlsConfig,
    31  	}
    32  	ws, _, err := dialer.Dial(s.addr, nil)
    33  	if err != nil {
    34  		return err
    35  	}
    36  	s.conn = ws
    37  	return nil
    38  }
    39  
    40  func (s *socket) Read(b []byte) (int, error) {
    41  	if s.reader == nil {
    42  		_, reader, err := s.conn.NextReader()
    43  		if err != nil {
    44  			return 0, err
    45  		}
    46  
    47  		s.reader = reader
    48  	}
    49  
    50  	bytesRead, err := s.reader.Read(b)
    51  	if err != nil {
    52  		s.reader = nil
    53  
    54  		if err == io.EOF {
    55  			err = nil
    56  		}
    57  	}
    58  
    59  	return bytesRead, err
    60  }
    61  
    62  func (s *socket) Write(b []byte) (int, error) {
    63  	nextWriter, err := s.conn.NextWriter(websocket.BinaryMessage)
    64  	if err != nil {
    65  		return 0, err
    66  	}
    67  
    68  	bytesWritten, err := nextWriter.Write(b)
    69  	nextWriter.Close()
    70  
    71  	return bytesWritten, err
    72  }
    73  
    74  func (s *socket) Close() error {
    75  	return s.conn.Close()
    76  }
    77  
    78  func (s *socket) LocalAddr() net.Addr {
    79  	return s.conn.LocalAddr()
    80  }
    81  
    82  func (s *socket) RemoteAddr() net.Addr {
    83  	return s.conn.RemoteAddr()
    84  }
    85  
    86  func (s *socket) SetDeadline(t time.Time) error {
    87  	if err := s.SetReadDeadline(t); err != nil {
    88  		return err
    89  	}
    90  
    91  	return s.SetWriteDeadline(t)
    92  }
    93  
    94  func (s *socket) SetReadDeadline(t time.Time) error {
    95  	return s.conn.SetReadDeadline(t)
    96  }
    97  
    98  func (s *socket) SetWriteDeadline(t time.Time) error {
    99  	return s.conn.SetWriteDeadline(t)
   100  }