github.com/anycable/anycable-go@v1.5.1/ws/connection.go (about)

     1  package ws
     2  
     3  import (
     4  	"time"
     5  
     6  	"github.com/gorilla/websocket"
     7  )
     8  
     9  // Connection is a WebSocket implementation of Connection
    10  type Connection struct {
    11  	conn *websocket.Conn
    12  }
    13  
    14  func NewConnection(conn *websocket.Conn) *Connection {
    15  	return &Connection{conn}
    16  }
    17  
    18  // Write writes a text message to a WebSocket
    19  func (ws Connection) Write(msg []byte, deadline time.Time) error {
    20  	if err := ws.conn.SetWriteDeadline(deadline); err != nil {
    21  		return err
    22  	}
    23  
    24  	w, err := ws.conn.NextWriter(websocket.TextMessage)
    25  
    26  	if err != nil {
    27  		return err
    28  	}
    29  
    30  	if _, err = w.Write(msg); err != nil {
    31  		return err
    32  	}
    33  
    34  	return w.Close()
    35  }
    36  
    37  // WriteBinary writes a binary message to a WebSocket
    38  func (ws Connection) WriteBinary(msg []byte, deadline time.Time) error {
    39  	if err := ws.conn.SetWriteDeadline(deadline); err != nil {
    40  		return err
    41  	}
    42  
    43  	w, err := ws.conn.NextWriter(websocket.BinaryMessage)
    44  
    45  	if err != nil {
    46  		return err
    47  	}
    48  
    49  	if _, err = w.Write(msg); err != nil {
    50  		return err
    51  	}
    52  
    53  	return w.Close()
    54  }
    55  
    56  func (ws Connection) Read() ([]byte, error) {
    57  	_, message, err := ws.conn.ReadMessage()
    58  	return message, err
    59  }
    60  
    61  // Close sends close frame with a given code and a reason
    62  func (ws Connection) Close(code int, reason string) {
    63  	CloseWithReason(ws.conn, code, reason)
    64  }
    65  
    66  // CloseWithReason closes WebSocket connection with the specified close code and reason
    67  func CloseWithReason(ws *websocket.Conn, code int, reason string) {
    68  	deadline := time.Now().Add(time.Second)
    69  	msg := websocket.FormatCloseMessage(code, reason)
    70  	ws.WriteControl(websocket.CloseMessage, msg, deadline) //nolint:errcheck
    71  	ws.Close()
    72  }