helpers.go 2.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105
  1. package stun
  2. // Interfaces that are implemented by message attributes, shorthands for them,
  3. // or helpers for message fields as type or transaction id.
  4. type (
  5. // Setter sets *Message attribute.
  6. Setter interface {
  7. AddTo(m *Message) error
  8. }
  9. // Getter parses attribute from *Message.
  10. Getter interface {
  11. GetFrom(m *Message) error
  12. }
  13. // Checker checks *Message attribute.
  14. Checker interface {
  15. Check(m *Message) error
  16. }
  17. )
  18. // Build resets message and applies setters to it in batch, returning on
  19. // first error. To prevent allocations, pass pointers to values.
  20. //
  21. // Example:
  22. // var (
  23. // t = BindingRequest
  24. // username = NewUsername("username")
  25. // nonce = NewNonce("nonce")
  26. // realm = NewRealm("example.org")
  27. // )
  28. // m := new(Message)
  29. // m.Build(t, username, nonce, realm) // 4 allocations
  30. // m.Build(&t, &username, &nonce, &realm) // 0 allocations
  31. //
  32. // See BenchmarkBuildOverhead.
  33. func (m *Message) Build(setters ...Setter) error {
  34. m.Reset()
  35. m.WriteHeader()
  36. for _, s := range setters {
  37. if err := s.AddTo(m); err != nil {
  38. return err
  39. }
  40. }
  41. return nil
  42. }
  43. // Check applies checkers to message in batch, returning on first error.
  44. func (m *Message) Check(checkers ...Checker) error {
  45. for _, c := range checkers {
  46. if err := c.Check(m); err != nil {
  47. return err
  48. }
  49. }
  50. return nil
  51. }
  52. // Parse applies getters to message in batch, returning on first error.
  53. func (m *Message) Parse(getters ...Getter) error {
  54. for _, c := range getters {
  55. if err := c.GetFrom(m); err != nil {
  56. return err
  57. }
  58. }
  59. return nil
  60. }
  61. // MustBuild wraps Build call and panics on error.
  62. func MustBuild(setters ...Setter) *Message {
  63. m, err := Build(setters...)
  64. if err != nil {
  65. panic(err) // nolint
  66. }
  67. return m
  68. }
  69. // Build wraps Message.Build method.
  70. func Build(setters ...Setter) (*Message, error) {
  71. m := new(Message)
  72. if err := m.Build(setters...); err != nil {
  73. return nil, err
  74. }
  75. return m, nil
  76. }
  77. // ForEach is helper that iterates over message attributes allowing to call
  78. // Getter in f callback to get all attributes of type t and returning on first
  79. // f error.
  80. //
  81. // The m.Get method inside f will be returning next attribute on each f call.
  82. // Does not error if there are no results.
  83. func (m *Message) ForEach(t AttrType, f func(m *Message) error) error {
  84. attrs := m.Attributes
  85. defer func() {
  86. m.Attributes = attrs
  87. }()
  88. for i, a := range attrs {
  89. if a.Type != t {
  90. continue
  91. }
  92. m.Attributes = attrs[i:]
  93. if err := f(m); err != nil {
  94. return err
  95. }
  96. }
  97. return nil
  98. }