encode.go 3.0 KB

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