github.com/imgk/caddy-trojan@v0.0.0-20221206043256-2631719e16c8/websocket/websocket.go (about)

     1  package websocket
     2  
     3  import (
     4  	"errors"
     5  	"io"
     6  	"net/http"
     7  	"time"
     8  
     9  	"github.com/gorilla/websocket"
    10  )
    11  
    12  // Upgrader is ...
    13  type Upgrader struct {
    14  	websocket.Upgrader
    15  }
    16  
    17  // IsWebSocketUpgrade is ...
    18  func IsWebSocketUpgrade(r *http.Request) bool {
    19  	return websocket.IsWebSocketUpgrade(r)
    20  }
    21  
    22  // eofReader is ...
    23  type eofReader struct{}
    24  
    25  // Read is ...
    26  func (*eofReader) Read(b []byte) (int, error) {
    27  	return 0, io.EOF
    28  }
    29  
    30  // Conn is ...
    31  type Conn struct {
    32  	*websocket.Conn
    33  	Reader io.Reader
    34  }
    35  
    36  // NewConn is ...
    37  func NewConn(c *websocket.Conn) *Conn {
    38  	return &Conn{
    39  		Conn:   c,
    40  		Reader: (*eofReader)(nil),
    41  	}
    42  }
    43  
    44  // Read is ...
    45  func (c *Conn) Read(b []byte) (n int, err error) {
    46  	n, _ = c.Reader.Read(b)
    47  	if n > 0 {
    48  		return n, nil
    49  	}
    50  
    51  	_, c.Reader, err = c.Conn.NextReader()
    52  	if err != nil {
    53  		if ce := (*websocket.CloseError)(nil); errors.As(err, &ce) {
    54  			return 0, io.EOF
    55  		}
    56  		return 0, err
    57  	}
    58  
    59  	n, _ = c.Reader.Read(b)
    60  	return n, nil
    61  }
    62  
    63  // Write is ...
    64  func (c *Conn) Write(b []byte) (int, error) {
    65  	err := c.Conn.WriteMessage(websocket.BinaryMessage, b)
    66  	if err != nil {
    67  		if ce := (*websocket.CloseError)(nil); errors.As(err, &ce) {
    68  			return 0, io.EOF
    69  		}
    70  		return 0, err
    71  	}
    72  	return len(b), nil
    73  }
    74  
    75  // Close is ...
    76  func (c *Conn) Close() error {
    77  	msg := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")
    78  	c.Conn.WriteControl(websocket.CloseMessage, msg, time.Now().Add(time.Second*5))
    79  	return c.Conn.Close()
    80  }