feat(websocket): improve websocket code sharing BE-11340 (#61)

api/http/handler/websocket/hijack.go

@@ -1,163 +0,0 @@
-package websocket
-
-import (
-	"bufio"
-	"errors"
-	"fmt"
-	"io"
-	"net"
-	"net/http"
-	"sync"
-	"time"
-
-	"github.com/gorilla/websocket"
-	"github.com/rs/zerolog/log"
-)
-
-const (
-	// Time allowed to write a message to the peer
-	writeWait = 10 * time.Second
-
-	// Send pings to peer with this period
-	pingPeriod = 50 * time.Second
-)
-
-func hijackRequest(
-	websocketConn *websocket.Conn,
-	conn net.Conn,
-	request *http.Request,
-	token string,
-) error {
-	resp, err := sendHTTPRequest(conn, request)
-	if err != nil {
-		return err
-	}
-	defer resp.Body.Close()
-
-	// Check if the response status code indicates an upgrade (101 Switching Protocols)
-	if resp.StatusCode != http.StatusSwitchingProtocols {
-		return fmt.Errorf("unexpected response status code: %d", resp.StatusCode)
-	}
-
-	errorChan := make(chan error, 1)
-	go readWebSocketToTCP(websocketConn, conn, errorChan)
-	go writeTCPToWebSocket(websocketConn, conn, errorChan)
-
-	err = <-errorChan
-	if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNoStatusReceived) {
-		log.Debug().Msgf("Unexpected close error: %v\n", err)
-		return err
-	}
-
-	log.Debug().Msgf("session ended")
-	return nil
-}
-
-// sendHTTPRequest sends an HTTP request over the provided net.Conn and parses the response.
-func sendHTTPRequest(conn net.Conn, req *http.Request) (*http.Response, error) {
-	// Send the HTTP request to the server
-	if err := req.Write(conn); err != nil {
-		return nil, fmt.Errorf("error writing request: %w", err)
-	}
-
-	// Read the response from the server
-	resp, err := http.ReadResponse(bufio.NewReader(conn), req)
-	if err != nil {
-		return nil, fmt.Errorf("error reading response: %w", err)
-	}
-
-	return resp, nil
-}
-
-func readWebSocketToTCP(websocketConn *websocket.Conn, tcpConn net.Conn, errorChan chan error) {
-	for {
-		messageType, p, err := websocketConn.ReadMessage()
-		if err != nil {
-			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway) {
-				log.Debug().Msgf("Unexpected close error: %v\n", err)
-			}
-			errorChan <- err
-			return
-		}
-
-		if messageType == websocket.TextMessage || messageType == websocket.BinaryMessage {
-			_, err := tcpConn.Write(p)
-			if err != nil {
-				log.Debug().Msgf("Error writing to TCP connection: %v\n", err)
-				errorChan <- err
-				return
-			}
-		}
-	}
-}
-
-func writeTCPToWebSocket(websocketConn *websocket.Conn, tcpConn net.Conn, errorChan chan error) {
-	var mu sync.Mutex
-	out := make([]byte, readerBufferSize)
-	input := make(chan string)
-	pingTicker := time.NewTicker(pingPeriod)
-	defer pingTicker.Stop()
-	defer websocketConn.Close()
-
-	websocketConn.SetReadLimit(2048)
-	websocketConn.SetPongHandler(func(string) error {
-		return nil
-	})
-
-	websocketConn.SetPingHandler(func(data string) error {
-		websocketConn.SetWriteDeadline(time.Now().Add(writeWait))
-		return websocketConn.WriteMessage(websocket.PongMessage, []byte(data))
-	})
-
-	reader := bufio.NewReader(tcpConn)
-
-	go func() {
-		for {
-			n, err := reader.Read(out)
-			if err != nil {
-				errorChan <- err
-				if !errors.Is(err, io.EOF) {
-					log.Debug().Msgf("error reading from server: %v", err)
-				}
-				return
-			}
-
-			processedOutput := validString(string(out[:n]))
-			input <- processedOutput
-		}
-	}()
-
-	for {
-		select {
-		case msg := <-input:
-			err := wswrite(websocketConn, &mu, msg)
-			if err != nil {
-				log.Debug().Msgf("error writing to websocket: %v", err)
-				errorChan <- err
-				return
-			}
-		case <-pingTicker.C:
-			if err := wsping(websocketConn, &mu); err != nil {
-				log.Debug().Msgf("error writing to websocket during pong response: %v", err)
-				errorChan <- err
-				return
-			}
-		}
-	}
-}
-
-func wswrite(websocketConn *websocket.Conn, mu *sync.Mutex, msg string) error {
-	mu.Lock()
-	defer mu.Unlock()
-
-	websocketConn.SetWriteDeadline(time.Now().Add(writeWait))
-	return websocketConn.WriteMessage(websocket.TextMessage, []byte(msg))
-}
-
-func wsping(websocketConn *websocket.Conn, mu *sync.Mutex) error {
-	mu.Lock()
-	defer mu.Unlock()
-
-	websocketConn.SetWriteDeadline(time.Now().Add(writeWait))
-	return websocketConn.WriteMessage(websocket.PingMessage, nil)
-}