poc-browser-csr-generation/handlers/websocket.go

package handlers

import (
	"encoding/json"
	"net/http"

	"github.com/gobwas/ws"
	"github.com/gobwas/ws/wsutil"
	log "github.com/sirupsen/logrus"
)

type WebSocketHandler struct {
	requestRegistry *SigningRequestRegistry
}

func NewWebSocketHandler(registry *SigningRequestRegistry) *WebSocketHandler {
	return &WebSocketHandler{requestRegistry: registry}
}

func (w *WebSocketHandler) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
	conn, _, _, err := ws.UpgradeHTTP(request, writer)
	if err != nil {
		http.Error(writer, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
		return
	}

	go func() {
		defer func() { _ = conn.Close() }()

		var (
			reader      = wsutil.NewReader(conn, ws.StateServerSide)
			writer      = wsutil.NewWriter(conn, ws.StateServerSide, ws.OpText)
			jsonDecoder = json.NewDecoder(reader)
			jsonEncoder = json.NewEncoder(writer)
		)

		for {
			header, err := reader.NextFrame()
			if err != nil {
				log.Error(err)
				break
			}
			if header.OpCode == ws.OpClose {
				log.Debug("channel closed")
				break
			}

			type requestType struct {
				RequestId string `json:"request_id"`
			}

			request := &requestType{}
			err = jsonDecoder.Decode(request)
			if err != nil {
				log.Error(err)
				break
			}

			channel, err := w.requestRegistry.GetResponseChannel(request.RequestId)
			if err != nil {
				log.Error(err)
				break
			}

			var response *responseData
			response = <-channel
			if err = jsonEncoder.Encode(response); err != nil {
				log.Error(err)
				break
			}
			close(channel)

			if err = writer.Flush(); err != nil {
				log.Error(err)
				break
			}
		}
	}()
}