encode.go 2.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102
  1. package wire
  2. import (
  3. "bytes"
  4. "encoding/binary"
  5. "errors"
  6. "fmt"
  7. "io"
  8. "reflect"
  9. )
  10. var ErrMarshalFailure = errors.New("failed to marshal")
  11. func Marshal(v any, w io.Writer) error {
  12. return marshal(reflect.TypeOf(v), reflect.ValueOf(v), "", w)
  13. }
  14. func marshal(t reflect.Type, v reflect.Value, tag reflect.StructTag, w io.Writer) error {
  15. switch t.Kind() {
  16. case reflect.Struct:
  17. for i := 0; i < t.NumField(); i++ {
  18. if err := marshal(t.Field(i).Type, v.Field(i), t.Field(i).Tag, w); err != nil {
  19. return err
  20. }
  21. }
  22. return nil
  23. case reflect.String:
  24. if lenTag, ok := tag.Lookup("len_prefix"); ok {
  25. switch lenTag {
  26. case "uint8":
  27. if err := binary.Write(w, binary.BigEndian, uint8(len(v.String()))); err != nil {
  28. return err
  29. }
  30. case "uint16":
  31. if err := binary.Write(w, binary.BigEndian, uint16(len(v.String()))); err != nil {
  32. return err
  33. }
  34. default:
  35. return fmt.Errorf("%w: unsupported len_prefix type %s. allowed types: uint8, uint16", ErrMarshalFailure, lenTag)
  36. }
  37. }
  38. return binary.Write(w, binary.BigEndian, []byte(v.String()))
  39. case reflect.Slice:
  40. // todo: only write to temporary buffer if len_prefix is set
  41. buf := &bytes.Buffer{}
  42. if t.Elem().Kind() == reflect.Struct {
  43. for j := 0; j < v.Len(); j++ {
  44. element := v.Index(j)
  45. if err := Marshal(element.Interface(), buf); err != nil {
  46. return err
  47. }
  48. }
  49. } else {
  50. if err := binary.Write(buf, binary.BigEndian, v.Interface()); err != nil {
  51. return fmt.Errorf("%w: error marshalling %s", ErrMarshalFailure, t.Elem().Kind())
  52. }
  53. }
  54. var hasLenPrefix bool
  55. if l, ok := tag.Lookup("len_prefix"); ok {
  56. hasLenPrefix = true
  57. switch l {
  58. case "uint8":
  59. if err := binary.Write(w, binary.BigEndian, uint8(buf.Len())); err != nil {
  60. return err
  61. }
  62. case "uint16":
  63. if err := binary.Write(w, binary.BigEndian, uint16(buf.Len())); err != nil {
  64. return err
  65. }
  66. default:
  67. return fmt.Errorf("%w: unsupported len_prefix type %s. allowed types: uint8, uint16", ErrMarshalFailure, l)
  68. }
  69. }
  70. if l, ok := tag.Lookup("count_prefix"); ok {
  71. if hasLenPrefix {
  72. return fmt.Errorf("%w: struct elem has both len_prefix and count_prefix: ", ErrMarshalFailure)
  73. }
  74. switch l {
  75. case "uint8":
  76. if err := binary.Write(w, binary.BigEndian, uint8(v.Len())); err != nil {
  77. return err
  78. }
  79. case "uint16":
  80. if err := binary.Write(w, binary.BigEndian, uint16(v.Len())); err != nil {
  81. return err
  82. }
  83. default:
  84. return fmt.Errorf("%w: unsupported count_prefix type %s. allowed types: uint8, uint16", ErrMarshalFailure, l)
  85. }
  86. }
  87. if buf.Len() > 0 {
  88. _, err := w.Write(buf.Bytes())
  89. return err
  90. }
  91. return nil
  92. case reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
  93. return binary.Write(w, binary.BigEndian, v.Interface())
  94. default:
  95. return fmt.Errorf("%w: unsupported type %v", ErrMarshalFailure, t.Kind())
  96. }
  97. }