github.com/mdaxf/iac@v0.0.0-20240519030858-58a061660378/vendor_skip/nhooyr.io/websocket/netconn.go (about)

     1  package websocket
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"io"
     7  	"math"
     8  	"net"
     9  	"sync"
    10  	"time"
    11  )
    12  
    13  // NetConn converts a *websocket.Conn into a net.Conn.
    14  //
    15  // It's for tunneling arbitrary protocols over WebSockets.
    16  // Few users of the library will need this but it's tricky to implement
    17  // correctly and so provided in the library.
    18  // See https://github.com/nhooyr/websocket/issues/100.
    19  //
    20  // Every Write to the net.Conn will correspond to a message write of
    21  // the given type on *websocket.Conn.
    22  //
    23  // The passed ctx bounds the lifetime of the net.Conn. If cancelled,
    24  // all reads and writes on the net.Conn will be cancelled.
    25  //
    26  // If a message is read that is not of the correct type, the connection
    27  // will be closed with StatusUnsupportedData and an error will be returned.
    28  //
    29  // Close will close the *websocket.Conn with StatusNormalClosure.
    30  //
    31  // When a deadline is hit, the connection will be closed. This is
    32  // different from most net.Conn implementations where only the
    33  // reading/writing goroutines are interrupted but the connection is kept alive.
    34  //
    35  // The Addr methods will return a mock net.Addr that returns "websocket" for Network
    36  // and "websocket/unknown-addr" for String.
    37  //
    38  // A received StatusNormalClosure or StatusGoingAway close frame will be translated to
    39  // io.EOF when reading.
    40  func NetConn(ctx context.Context, c *Conn, msgType MessageType) net.Conn {
    41  	nc := &netConn{
    42  		c:       c,
    43  		msgType: msgType,
    44  	}
    45  
    46  	var cancel context.CancelFunc
    47  	nc.writeContext, cancel = context.WithCancel(ctx)
    48  	nc.writeTimer = time.AfterFunc(math.MaxInt64, cancel)
    49  	if !nc.writeTimer.Stop() {
    50  		<-nc.writeTimer.C
    51  	}
    52  
    53  	nc.readContext, cancel = context.WithCancel(ctx)
    54  	nc.readTimer = time.AfterFunc(math.MaxInt64, cancel)
    55  	if !nc.readTimer.Stop() {
    56  		<-nc.readTimer.C
    57  	}
    58  
    59  	return nc
    60  }
    61  
    62  type netConn struct {
    63  	c       *Conn
    64  	msgType MessageType
    65  
    66  	writeTimer   *time.Timer
    67  	writeContext context.Context
    68  
    69  	readTimer   *time.Timer
    70  	readContext context.Context
    71  
    72  	readMu sync.Mutex
    73  	eofed  bool
    74  	reader io.Reader
    75  }
    76  
    77  var _ net.Conn = &netConn{}
    78  
    79  func (c *netConn) Close() error {
    80  	return c.c.Close(StatusNormalClosure, "")
    81  }
    82  
    83  func (c *netConn) Write(p []byte) (int, error) {
    84  	err := c.c.Write(c.writeContext, c.msgType, p)
    85  	if err != nil {
    86  		return 0, err
    87  	}
    88  	return len(p), nil
    89  }
    90  
    91  func (c *netConn) Read(p []byte) (int, error) {
    92  	c.readMu.Lock()
    93  	defer c.readMu.Unlock()
    94  
    95  	if c.eofed {
    96  		return 0, io.EOF
    97  	}
    98  
    99  	if c.reader == nil {
   100  		typ, r, err := c.c.Reader(c.readContext)
   101  		if err != nil {
   102  			switch CloseStatus(err) {
   103  			case StatusNormalClosure, StatusGoingAway:
   104  				c.eofed = true
   105  				return 0, io.EOF
   106  			}
   107  			return 0, err
   108  		}
   109  		if typ != c.msgType {
   110  			err := fmt.Errorf("unexpected frame type read (expected %v): %v", c.msgType, typ)
   111  			c.c.Close(StatusUnsupportedData, err.Error())
   112  			return 0, err
   113  		}
   114  		c.reader = r
   115  	}
   116  
   117  	n, err := c.reader.Read(p)
   118  	if err == io.EOF {
   119  		c.reader = nil
   120  		err = nil
   121  	}
   122  	return n, err
   123  }
   124  
   125  type websocketAddr struct {
   126  }
   127  
   128  func (a websocketAddr) Network() string {
   129  	return "websocket"
   130  }
   131  
   132  func (a websocketAddr) String() string {
   133  	return "websocket/unknown-addr"
   134  }
   135  
   136  func (c *netConn) RemoteAddr() net.Addr {
   137  	return websocketAddr{}
   138  }
   139  
   140  func (c *netConn) LocalAddr() net.Addr {
   141  	return websocketAddr{}
   142  }
   143  
   144  func (c *netConn) SetDeadline(t time.Time) error {
   145  	c.SetWriteDeadline(t)
   146  	c.SetReadDeadline(t)
   147  	return nil
   148  }
   149  
   150  func (c *netConn) SetWriteDeadline(t time.Time) error {
   151  	if t.IsZero() {
   152  		c.writeTimer.Stop()
   153  	} else {
   154  		c.writeTimer.Reset(t.Sub(time.Now()))
   155  	}
   156  	return nil
   157  }
   158  
   159  func (c *netConn) SetReadDeadline(t time.Time) error {
   160  	if t.IsZero() {
   161  		c.readTimer.Stop()
   162  	} else {
   163  		c.readTimer.Reset(t.Sub(time.Now()))
   164  	}
   165  	return nil
   166  }