goyave.dev/goyave/v5@v5.0.0-rc9.0.20240517145003-d3f977d0b9f3/websocket/conn.go (about)

     1  package websocket
     2  
     3  import (
     4  	"context"
     5  	stderrors "errors"
     6  	"net/http"
     7  	"strings"
     8  	"sync"
     9  	"time"
    10  
    11  	ws "github.com/gorilla/websocket"
    12  
    13  	"goyave.dev/goyave/v5/util/errors"
    14  )
    15  
    16  // Conn represents a WebSocket connection.
    17  type Conn struct {
    18  	*ws.Conn
    19  	waitClose    chan struct{}
    20  	closeTimeout time.Duration
    21  	closeOnce    sync.Once
    22  }
    23  
    24  func newConn(c *ws.Conn, closeTimeout time.Duration) *Conn {
    25  	conn := &Conn{
    26  		Conn:         c,
    27  		waitClose:    make(chan struct{}, 1),
    28  		closeTimeout: closeTimeout,
    29  	}
    30  	c.SetCloseHandler(conn.closeHandler)
    31  	return conn
    32  }
    33  
    34  // SetCloseHandshakeTimeout set the timeout used when writing and reading
    35  // close frames during the close handshake.
    36  func (c *Conn) SetCloseHandshakeTimeout(timeout time.Duration) {
    37  	c.closeTimeout = timeout
    38  }
    39  
    40  // GetCloseHandshakeTimeout return the timeout used when writing and reading
    41  // close frames during the close handshake.
    42  func (c *Conn) GetCloseHandshakeTimeout() time.Duration {
    43  	return c.closeTimeout
    44  }
    45  
    46  func (c *Conn) closeHandler(_ int, _ string) error {
    47  	c.waitClose <- struct{}{}
    48  	return nil
    49  }
    50  
    51  // CloseNormal performs the closing handshake as specified by
    52  // RFC 6455 Section 1.4. Sends status code 1000 (normal closure) and
    53  // message "Server closed connection".
    54  //
    55  // This function expects another goroutine to be reading the connection,
    56  // expecting the close frame in response. This waiting can time out. If so,
    57  // Close will just close the connection.
    58  //
    59  // Calling this function multiple times is safe and only the first call will
    60  // write the close frame to the connection.
    61  func (c *Conn) CloseNormal() error {
    62  	return c.Close(ws.CloseNormalClosure, NormalClosureMessage)
    63  }
    64  
    65  // CloseWithError performs the closing handshake as specified by
    66  // RFC 6455 Section 1.4 because a server error occurred.
    67  // Sends status code 1011 (internal server error) and
    68  // message "Internal server error".
    69  //
    70  // This function starts another goroutine to read the connection,
    71  // expecting the close frame in response. This waiting can time out. If so,
    72  // Close will just close the connection. Therefore, it is not safe to call
    73  // this function if there is already an active reader.
    74  func (c *Conn) CloseWithError(_ error) error {
    75  	return c.internalClose(ws.CloseInternalServerErr, http.StatusText(http.StatusInternalServerError))
    76  }
    77  
    78  // Close performs the closing handshake as specified by RFC 6455 Section 1.4.
    79  //
    80  // This function expects another goroutine to be reading the connection,
    81  // expecting the close frame in response. This waiting can time out. If so,
    82  // Close will just close the connection.
    83  //
    84  // Calling this function multiple times is safe and only the first call will
    85  // write the close frame to the connection.
    86  func (c *Conn) Close(code int, message string) error {
    87  	var err error
    88  	c.closeOnce.Do(func() {
    89  		deadline := time.Now().Add(c.closeTimeout)
    90  		m := ws.FormatCloseMessage(code, message)
    91  		writeErr := c.WriteControl(ws.CloseMessage, m, deadline)
    92  		if writeErr != nil && !stderrors.Is(writeErr, ws.ErrCloseSent) {
    93  			if strings.Contains(writeErr.Error(), "use of closed network connection") {
    94  				err = errors.New(writeErr)
    95  			}
    96  			return
    97  		}
    98  
    99  		ctx, cancel := context.WithDeadline(context.Background(), deadline)
   100  		defer cancel()
   101  
   102  		select {
   103  		case <-ctx.Done():
   104  		case <-c.waitClose:
   105  			close(c.waitClose)
   106  		}
   107  
   108  		err = errors.New(c.Conn.Close())
   109  	})
   110  
   111  	return err
   112  }
   113  
   114  // internalClose performs the close handshake. Starts a goroutine reading in the connection,
   115  // expecting a close frame response for the close handshake. This function should only be
   116  // used if the server wants to initiate the close handshake.
   117  func (c *Conn) internalClose(code int, message string) error {
   118  	go func() {
   119  		for {
   120  			if _, _, err := c.ReadMessage(); err != nil {
   121  				return
   122  			}
   123  		}
   124  	}()
   125  	return c.Close(code, message)
   126  }