github.com/lastbackend/toolkit@v0.0.0-20241020043710-cafa37b95aad/pkg/server/http/websockets/client.go (about)

     1  /*
     2  Copyright [2014] - [2023] The Last.Backend authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package websockets
    18  
    19  import (
    20  	"context"
    21  	"encoding/json"
    22  	"github.com/lastbackend/toolkit/pkg/runtime/logger"
    23  	"time"
    24  
    25  	"github.com/gorilla/websocket"
    26  )
    27  
    28  // The message types are defined in RFC 6455, section 11.8.
    29  const (
    30  	TextMessage   = websocket.TextMessage
    31  	BinaryMessage = websocket.BinaryMessage
    32  	CloseMessage  = websocket.CloseMessage
    33  	PingMessage   = websocket.PingMessage
    34  	PongMessage   = websocket.PongMessage
    35  )
    36  
    37  type ClientList map[*Client]bool
    38  
    39  // Client is a websocket client, basically a frontend visitor
    40  type Client struct {
    41  	ctx        context.Context
    42  	log        logger.Logger
    43  	connection *websocket.Conn
    44  	manager    *Manager
    45  }
    46  
    47  var (
    48  	// pongWait is how long we will await a pong response from client
    49  	pongWait = 10 * time.Second
    50  	// pingInterval has to be less than pongWait, We cant multiply by 0.9 to get 90% of time
    51  	// Because that can make decimals, so instead *9 / 10 to get 90%
    52  	// The reason why it has to be less than PingRequency is because otherwise it will send
    53  	// a new Ping before getting response
    54  	pingInterval = (pongWait * 9) / 10
    55  )
    56  
    57  // NewClient is used to initialize a new Client with all required values initialized
    58  func NewClient(ctx context.Context, log logger.Logger, conn *websocket.Conn, manager *Manager) *Client {
    59  	return &Client{
    60  		ctx:        ctx,
    61  		log:        log,
    62  		connection: conn,
    63  		manager:    manager,
    64  	}
    65  }
    66  
    67  // WriteMessage is a helper method for getting a writer using NextWriter.
    68  func (c *Client) WriteMessage(messageType int, data []byte) error {
    69  	if err := c.connection.WriteMessage(messageType, data); err != nil {
    70  		c.log.Errorf("failed write message: %v", err)
    71  		return err
    72  	}
    73  	return nil
    74  }
    75  
    76  // WriteJSON writes the JSON encoding of v to the connection.
    77  func (c *Client) WriteJSON(i interface{}) error {
    78  	if err := c.connection.WriteJSON(i); err != nil {
    79  		c.log.Errorf("failed write JSON message: %v", err)
    80  		return err
    81  	}
    82  	return nil
    83  }
    84  
    85  // readMessages will start the client to read messages and handle them
    86  // appropriately.
    87  // This is supposed to be run as a goroutine
    88  func (c *Client) readMessages() {
    89  	defer func() {
    90  		c.manager.removeClient(c)
    91  	}()
    92  
    93  	c.connection.SetReadLimit(512)
    94  
    95  	if err := c.connection.SetReadDeadline(time.Now().Add(pongWait)); err != nil {
    96  		c.log.Errorf("configure Wait time for Pong response failed %v", err)
    97  		return
    98  	}
    99  
   100  	c.connection.SetPongHandler(c.pongHandler)
   101  
   102  	for {
   103  		_, payload, err := c.connection.ReadMessage()
   104  
   105  		if err != nil {
   106  			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
   107  				c.log.Errorf("failed reading message: %v", err)
   108  			}
   109  			break
   110  		}
   111  
   112  		var request Event
   113  		if err := json.Unmarshal(payload, &request); err != nil {
   114  			c.log.Errorf("failed marshalling message: %v", err)
   115  			c.connection.WriteMessage(TextMessage, []byte(`{"error":"invalid JSON format"}`))
   116  			continue
   117  		}
   118  
   119  		if err := c.manager.routeEvent(request, c); err != nil {
   120  			c.log.Errorf("failed handling message: %v", err)
   121  		}
   122  	}
   123  }
   124  
   125  // pongHandler is used to handle PongMessages for the Client
   126  func (c *Client) pongHandler(_ string) error {
   127  	return c.connection.SetReadDeadline(time.Now().Add(pongWait))
   128  }
   129  
   130  // writeMessages is a process that listens for new messages to output to the Client
   131  func (c *Client) writeMessages() {
   132  	ticker := time.NewTicker(pingInterval)
   133  	defer func() {
   134  		ticker.Stop()
   135  		c.manager.removeClient(c)
   136  	}()
   137  
   138  	for {
   139  		select {
   140  		case <-ticker.C:
   141  			if err := c.connection.WriteMessage(websocket.PingMessage, []byte{}); err != nil {
   142  				c.log.Errorf("failed write ping message: %v", err)
   143  				return
   144  			}
   145  		}
   146  
   147  	}
   148  }