stream.go raw

   1  package ws
   2  
   3  import (
   4  	"errors"
   5  	"io"
   6  	"net"
   7  	"net/url"
   8  	"sync"
   9  	"time"
  10  )
  11  
  12  // StreamConn adapts a WebSocket to net.Conn: Read returns the payloads of
  13  // binary or text messages in order, Write sends one binary message. A P2P
  14  // stream is a byte stream, so message boundaries are irrelevant above this.
  15  type StreamConn struct {
  16  	c      *Conn
  17  	local  net.Addr
  18  	remote net.Addr
  19  
  20  	mu  sync.Mutex
  21  	buf []byte
  22  	err error
  23  }
  24  
  25  // DialStream opens a WebSocket and returns it as a net.Conn.
  26  func DialStream(rawURL string) (*StreamConn, error) {
  27  	c, err := Dial(rawURL)
  28  	if err != nil {
  29  		return nil, err
  30  	}
  31  	host := "ws"
  32  	if u, err := url.Parse(rawURL); err == nil && u.Host != "" {
  33  		host = u.Host
  34  	}
  35  	return &StreamConn{
  36  		c:      c,
  37  		local:  streamAddr("stream:local"),
  38  		remote: streamAddr(host),
  39  	}, nil
  40  }
  41  
  42  // Read implements net.Conn.
  43  func (s *StreamConn) Read(p []byte) (int, error) {
  44  	if len(p) == 0 {
  45  		return 0, nil
  46  	}
  47  	for {
  48  		s.mu.Lock()
  49  		if len(s.buf) > 0 {
  50  			n := copy(p, s.buf)
  51  			s.buf = s.buf[n:]
  52  			s.mu.Unlock()
  53  			return n, nil
  54  		}
  55  		err := s.err
  56  		s.mu.Unlock()
  57  		if err != nil {
  58  			return 0, err
  59  		}
  60  		msg, err := s.c.ReadMessage()
  61  		if err != nil {
  62  			if errors.Is(err, io.EOF) {
  63  				err = io.EOF
  64  			}
  65  			s.mu.Lock()
  66  			s.err = err
  67  			s.mu.Unlock()
  68  			return 0, err
  69  		}
  70  		if len(msg) == 0 {
  71  			continue
  72  		}
  73  		s.mu.Lock()
  74  		s.buf = msg
  75  		s.mu.Unlock()
  76  	}
  77  }
  78  
  79  // Write implements net.Conn.
  80  func (s *StreamConn) Write(p []byte) (int, error) {
  81  	if err := s.c.WriteBinary(p); err != nil {
  82  		return 0, err
  83  	}
  84  	return len(p), nil
  85  }
  86  
  87  // Close implements net.Conn.
  88  func (s *StreamConn) Close() error { return s.c.Close() }
  89  
  90  // LocalAddr implements net.Conn.
  91  func (s *StreamConn) LocalAddr() net.Addr { return s.local }
  92  
  93  // RemoteAddr implements net.Conn.
  94  func (s *StreamConn) RemoteAddr() net.Addr { return s.remote }
  95  
  96  // SetDeadline implements net.Conn.
  97  func (s *StreamConn) SetDeadline(t time.Time) error { return s.c.SetDeadline(t) }
  98  
  99  // SetReadDeadline implements net.Conn.
 100  func (s *StreamConn) SetReadDeadline(t time.Time) error { return s.c.SetReadDeadline(t) }
 101  
 102  // SetWriteDeadline implements net.Conn.
 103  func (s *StreamConn) SetWriteDeadline(t time.Time) error { return s.c.SetWriteDeadline(t) }
 104  
 105  type streamAddr string
 106  
 107  func (a streamAddr) Network() string { return "ws" }
 108  func (a streamAddr) String() string  { return string(a) }
 109