ws_proxy.go

package main

import (
	"github.com/gorilla/websocket"
	"github.com/sirupsen/logrus"
	"net/http"
)

/**
gin.engine.Any("/ws",gin.WrapF(WsProxy))
**/

var upgrader = websocket.Upgrader{
	ReadBufferSize:  4096,
	WriteBufferSize: 4096,
	CheckOrigin:     func(r *http.Request) bool { return true },
}

func WsProxy(w http.ResponseWriter, r *http.Request) {
	l := logrus.WithField("proxy", "websocket port proxy")
	logrus.Infof("websocket proxy incoming.....")
	c, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		l.Error(err)
		w.WriteHeader(http.StatusBadRequest)
		return
	}
	defer c.Close()

	backendConn, _, err := websocket.DefaultDialer.Dial("ws://localhost:4000", nil)
	if err != nil {
		l.Error(err)
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	go func() {
		defer func() {
			backendConn.Close()
			c.Close()
		}()
		for {
			msgType, message, err := backendConn.ReadMessage()
			if err != nil {
				l.Error(err)
				return
			}
			l.Tracef("backendConn recv: %s", message)
			err = c.WriteMessage(msgType, message)
			if err != nil {
				l.Error(err)
				return
			}
		}
	}()
	for {
		mt, message, err := c.ReadMessage()
		if err != nil {
			l.Error("read:", err)
			break
		}
		l.Tracef("recv: %s", message)
		err = backendConn.WriteMessage(mt, message)
		if err != nil {
			l.Error(err)
			break
		}
	}
}