message.go 1.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081
  1. package simple
  2. import (
  3. "bytes"
  4. "encoding/binary"
  5. "io"
  6. "net"
  7. cerr "github.com/mhaya/error"
  8. )
  9. var (
  10. NoneMessage = Message{} // none message
  11. headLength = 8 // ID uint32(4 bytes) + DataLen uint32(4 bytes)
  12. dataLength uint32 = 4096 // data length
  13. )
  14. type Message struct {
  15. MID uint32
  16. Len uint32
  17. Data []byte
  18. }
  19. func ReadMessage(conn net.Conn) (Message, bool, error) {
  20. header, err := io.ReadAll(io.LimitReader(conn, int64(headLength)))
  21. if err != nil {
  22. return NoneMessage, true, err
  23. }
  24. // if the header has no data, we can consider it as a closed connection
  25. if len(header) == 0 {
  26. return NoneMessage, true, cerr.PacketConnectClosed
  27. }
  28. msg, err := parseHeader(header)
  29. if err != nil {
  30. return NoneMessage, true, err
  31. }
  32. msgData, err := io.ReadAll(io.LimitReader(conn, int64(msg.Len)))
  33. if err != nil {
  34. return NoneMessage, true, err
  35. }
  36. msg.Data = msgData
  37. return msg, false, nil
  38. }
  39. func parseHeader(header []byte) (Message, error) {
  40. msg := Message{}
  41. if len(header) != headLength {
  42. return msg, cerr.PacketInvalidHeader
  43. }
  44. bytesReader := bytes.NewReader(header)
  45. err := binary.Read(bytesReader, endian, &msg.MID)
  46. if err != nil {
  47. return msg, err
  48. }
  49. err = binary.Read(bytesReader, endian, &msg.Len)
  50. if err != nil {
  51. return msg, err
  52. }
  53. if msg.Len > dataLength {
  54. return msg, cerr.PacketSizeExceed
  55. }
  56. return msg, nil
  57. }
  58. func pack(mid uint32, data []byte) ([]byte, error) {
  59. pkg := bytes.NewBuffer([]byte{})
  60. binary.Write(pkg, endian, mid)
  61. binary.Write(pkg, endian, uint32(len(data)))
  62. binary.Write(pkg, endian, data)
  63. return pkg.Bytes(), nil
  64. }