github.com/profzone/eden-framework@v1.0.10/pkg/courier/transport_http/websocket.go (about)

     1  package transport_http
     2  
     3  import (
     4  	"encoding/json"
     5  	"fmt"
     6  	"reflect"
     7  	"strconv"
     8  	"time"
     9  
    10  	"github.com/fatih/color"
    11  	"github.com/gorilla/websocket"
    12  	"github.com/sirupsen/logrus"
    13  
    14  	"github.com/profzone/eden-framework/pkg/strings"
    15  )
    16  
    17  type IWebSocket interface {
    18  	SubscribeOn(conn *websocket.Conn)
    19  }
    20  
    21  type TypeAndListener struct {
    22  	Listener Listener
    23  	Type     reflect.Type
    24  }
    25  
    26  func typeIndirect(tpe reflect.Type) reflect.Type {
    27  	for tpe.Kind() == reflect.Ptr {
    28  		tpe = tpe.Elem()
    29  	}
    30  	return tpe
    31  }
    32  
    33  type Listener func(v interface{}, ws *WSClient) error
    34  
    35  type Listeners map[string]TypeAndListener
    36  
    37  func (listeners Listeners) On(v interface{}, listener Listener) Listeners {
    38  	tpe := typeIndirect(reflect.TypeOf(v))
    39  	listeners[tpe.Name()] = TypeAndListener{
    40  		Listener: listener,
    41  		Type:     tpe,
    42  	}
    43  	return listeners
    44  }
    45  
    46  func (listeners Listeners) SubscribeOn(conn *websocket.Conn) {
    47  	defer conn.Close()
    48  	ws := ConnWS(conn)
    49  	logger := logrus.WithField("remote_addr", conn.RemoteAddr().String())
    50  
    51  	for !ws.Closed {
    52  		msg := Msg{}
    53  		err := conn.ReadJSON(&msg)
    54  		if err != nil {
    55  			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway) {
    56  				logger.Errorf("%s", err.Error())
    57  			}
    58  			break
    59  		}
    60  		logrus.Debugf(
    61  			`{"type":"%s","payload": %s}`,
    62  			color.RedString(msg.Type),
    63  			color.GreenString("%s", strconv.Quote(string(msg.Payload))),
    64  		)
    65  
    66  		if typeAndListener, exists := listeners[msg.Type]; exists {
    67  			m := reflect.New(typeAndListener.Type).Interface()
    68  			errForUnmarshal := PayloadUnmarshal(msg.Payload, m)
    69  			if errForUnmarshal != nil {
    70  				logger.Errorf("%s", errForUnmarshal)
    71  				break
    72  			}
    73  			errForOnMsg := typeAndListener.Listener(m, ws)
    74  			if errForOnMsg != nil {
    75  				logger.Errorf("%s", errForOnMsg)
    76  				break
    77  			}
    78  		}
    79  	}
    80  }
    81  
    82  type Msg struct {
    83  	Type    string `json:"type"`
    84  	Payload []byte `json:"payload"`
    85  }
    86  
    87  func ConnWS(conn *websocket.Conn) *WSClient {
    88  	return &WSClient{conn: conn}
    89  }
    90  
    91  type WSClient struct {
    92  	conn   *websocket.Conn
    93  	Closed bool
    94  }
    95  
    96  func (c *WSClient) Close() error {
    97  	defer c.conn.Close()
    98  	err := c.conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
    99  	if err != nil {
   100  		return err
   101  	}
   102  	time.Sleep(1 * time.Second)
   103  	return nil
   104  }
   105  
   106  func (c WSClient) Send(payload interface{}) (err error) {
   107  	msg := &Msg{}
   108  	msg.Type = typeIndirect(reflect.TypeOf(payload)).Name()
   109  	data, errForMarshal := PayloadMarshal(payload)
   110  	if errForMarshal != nil {
   111  		return errForMarshal
   112  	}
   113  	msg.Payload = data
   114  	return c.conn.WriteJSON(msg)
   115  }
   116  
   117  func PayloadMarshal(payload interface{}) ([]byte, error) {
   118  	tpe := typeIndirect(reflect.TypeOf(payload))
   119  	if tpe.Kind() == reflect.Struct || tpe.Kind() == reflect.Array || tpe.Kind() == reflect.Slice {
   120  		return json.Marshal(payload)
   121  	}
   122  	if marshaler, ok := payload.(json.Marshaler); ok {
   123  		return marshaler.MarshalJSON()
   124  	}
   125  	return []byte(fmt.Sprintf("%v", payload)), nil
   126  }
   127  
   128  func PayloadUnmarshal(data []byte, payload interface{}) error {
   129  	tpe := typeIndirect(reflect.TypeOf(payload))
   130  	if tpe.Kind() == reflect.Struct || tpe.Kind() == reflect.Array || tpe.Kind() == reflect.Slice {
   131  		return json.Unmarshal(data, payload)
   132  	}
   133  	if unmarshaler, ok := payload.(json.Unmarshaler); ok {
   134  		return unmarshaler.UnmarshalJSON(data)
   135  	}
   136  	rv := reflect.ValueOf(payload)
   137  	return str.ConvertFromStr(string(data), rv)
   138  }