github.com/isyscore/isc-gobase@v1.5.3-0.20231218061332-cbc7451899e9/websocket/message.go (about)

     1  package websocket
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/binary"
     6  	"encoding/json"
     7  	"errors"
     8  	"fmt"
     9  	"strconv"
    10  
    11  	"github.com/valyala/bytebufferpool"
    12  )
    13  
    14  type messageType uint8
    15  
    16  func (m messageType) String() string {
    17  	return strconv.Itoa(int(m))
    18  }
    19  
    20  func (m messageType) Name() string {
    21  	switch m {
    22  	case messageTypeString:
    23  		return "string"
    24  	case messageTypeInt:
    25  		return "int"
    26  	case messageTypeBool:
    27  		return "bool"
    28  	case messageTypeBytes:
    29  		return "[]byte"
    30  	case messageTypeJSON:
    31  		return "json"
    32  	default:
    33  		return "Invalid(" + m.String() + ")"
    34  	}
    35  }
    36  
    37  const (
    38  	messageTypeString messageType = iota
    39  	messageTypeInt
    40  	messageTypeBool
    41  	messageTypeBytes
    42  	messageTypeJSON
    43  
    44  	messageSeparator = ";"
    45  )
    46  
    47  var messageSeparatorByte = messageSeparator[0]
    48  
    49  type messageSerializer struct {
    50  	prefix []byte
    51  
    52  	prefixLen       int
    53  	separatorLen    int
    54  	prefixAndSepIdx int
    55  	prefixIdx       int
    56  	separatorIdx    int
    57  
    58  	buf *bytebufferpool.Pool
    59  }
    60  
    61  func newMessageSerializer(messagePrefix []byte) *messageSerializer {
    62  	return &messageSerializer{
    63  		prefix:          messagePrefix,
    64  		prefixLen:       len(messagePrefix),
    65  		separatorLen:    len(messageSeparator),
    66  		prefixAndSepIdx: len(messagePrefix) + len(messageSeparator) - 1,
    67  		prefixIdx:       len(messagePrefix) - 1,
    68  		separatorIdx:    len(messageSeparator) - 1,
    69  
    70  		buf: new(bytebufferpool.Pool),
    71  	}
    72  }
    73  
    74  var (
    75  	boolTrueB  = []byte("true")
    76  	boolFalseB = []byte("false")
    77  )
    78  
    79  func (ms *messageSerializer) serialize(event string, data any) ([]byte, error) {
    80  	b := ms.buf.Get()
    81  	_, _ = b.Write(ms.prefix)
    82  	_, _ = b.WriteString(event)
    83  	_ = b.WriteByte(messageSeparatorByte)
    84  
    85  	switch v := data.(type) {
    86  	case string:
    87  		_, _ = b.WriteString(messageTypeString.String())
    88  		_ = b.WriteByte(messageSeparatorByte)
    89  		_, _ = b.WriteString(v)
    90  	case int:
    91  		_, _ = b.WriteString(messageTypeInt.String())
    92  		_ = b.WriteByte(messageSeparatorByte)
    93  		_ = binary.Write(b, binary.LittleEndian, v)
    94  	case bool:
    95  		_, _ = b.WriteString(messageTypeBool.String())
    96  		_ = b.WriteByte(messageSeparatorByte)
    97  		if v {
    98  			_, _ = b.Write(boolTrueB)
    99  		} else {
   100  			_, _ = b.Write(boolFalseB)
   101  		}
   102  	case []byte:
   103  		_, _ = b.WriteString(messageTypeBytes.String())
   104  		_ = b.WriteByte(messageSeparatorByte)
   105  		_, _ = b.Write(v)
   106  	default:
   107  		//we suppose is json
   108  		res, err := json.Marshal(data)
   109  		if err != nil {
   110  			return nil, err
   111  		}
   112  		_, _ = b.WriteString(messageTypeJSON.String())
   113  		_ = b.WriteByte(messageSeparatorByte)
   114  		_, _ = b.Write(res)
   115  	}
   116  
   117  	message := b.Bytes()
   118  	ms.buf.Put(b)
   119  
   120  	return message, nil
   121  }
   122  
   123  var errInvalidTypeMessage = "Type %s is invalid for message: %s"
   124  
   125  func (ms *messageSerializer) deserialize(event []byte, websocketMessage []byte) (any, error) {
   126  	dataStartIdx := ms.prefixAndSepIdx + len(event) + 3
   127  	if len(websocketMessage) <= dataStartIdx {
   128  		return nil, errors.New("websocket invalid message: " + string(websocketMessage))
   129  	}
   130  
   131  	typ, err := strconv.Atoi(string(websocketMessage[ms.prefixAndSepIdx+len(event)+1 : ms.prefixAndSepIdx+len(event)+2])) // in order to iris-websocket-message;user;-> 4
   132  	if err != nil {
   133  		return nil, err
   134  	}
   135  
   136  	data := websocketMessage[dataStartIdx:]
   137  
   138  	switch messageType(typ) {
   139  	case messageTypeString:
   140  		return string(data), nil
   141  	case messageTypeInt:
   142  		msg, err := strconv.Atoi(string(data))
   143  		if err != nil {
   144  			return nil, err
   145  		}
   146  		return msg, nil
   147  	case messageTypeBool:
   148  		if bytes.Equal(data, boolTrueB) {
   149  			return true, nil
   150  		}
   151  		return false, nil
   152  	case messageTypeBytes:
   153  		return data, nil
   154  	case messageTypeJSON:
   155  		var msg any
   156  		err := json.Unmarshal(data, &msg)
   157  		return msg, err
   158  	default:
   159  		return nil, errors.New(fmt.Sprintf(errInvalidTypeMessage, messageType(typ).Name(), websocketMessage))
   160  	}
   161  }
   162  
   163  func (ms *messageSerializer) getWebsocketCustomEvent(websocketMessage []byte) []byte {
   164  	if len(websocketMessage) < ms.prefixAndSepIdx {
   165  		return nil
   166  	}
   167  	s := websocketMessage[ms.prefixAndSepIdx:]
   168  	evt := s[:bytes.IndexByte(s, messageSeparatorByte)]
   169  	return evt
   170  }