helpers.go 2.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106
  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. //
  23. // var (
  24. // t = BindingRequest
  25. // username = NewUsername("username")
  26. // nonce = NewNonce("nonce")
  27. // realm = NewRealm("example.org")
  28. // )
  29. // m := new(Message)
  30. // m.Build(t, username, nonce, realm) // 4 allocations
  31. // m.Build(&t, &username, &nonce, &realm) // 0 allocations
  32. //
  33. // See BenchmarkBuildOverhead.
  34. func (m *Message) Build(setters ...Setter) error {
  35. m.Reset()
  36. m.WriteHeader()
  37. for _, s := range setters {
  38. if err := s.AddTo(m); err != nil {
  39. return err
  40. }
  41. }
  42. return nil
  43. }
  44. // Check applies checkers to message in batch, returning on first error.
  45. func (m *Message) Check(checkers ...Checker) error {
  46. for _, c := range checkers {
  47. if err := c.Check(m); err != nil {
  48. return err
  49. }
  50. }
  51. return nil
  52. }
  53. // Parse applies getters to message in batch, returning on first error.
  54. func (m *Message) Parse(getters ...Getter) error {
  55. for _, c := range getters {
  56. if err := c.GetFrom(m); err != nil {
  57. return err
  58. }
  59. }
  60. return nil
  61. }
  62. // MustBuild wraps Build call and panics on error.
  63. func MustBuild(setters ...Setter) *Message {
  64. m, err := Build(setters...)
  65. if err != nil {
  66. panic(err) // nolint
  67. }
  68. return m
  69. }
  70. // Build wraps Message.Build method.
  71. func Build(setters ...Setter) (*Message, error) {
  72. m := new(Message)
  73. if err := m.Build(setters...); err != nil {
  74. return nil, err
  75. }
  76. return m, nil
  77. }
  78. // ForEach is helper that iterates over message attributes allowing to call
  79. // Getter in f callback to get all attributes of type t and returning on first
  80. // f error.
  81. //
  82. // The m.Get method inside f will be returning next attribute on each f call.
  83. // Does not error if there are no results.
  84. func (m *Message) ForEach(t AttrType, f func(m *Message) error) error {
  85. attrs := m.Attributes
  86. defer func() {
  87. m.Attributes = attrs
  88. }()
  89. for i, a := range attrs {
  90. if a.Type != t {
  91. continue
  92. }
  93. m.Attributes = attrs[i:]
  94. if err := f(m); err != nil {
  95. return err
  96. }
  97. }
  98. return nil
  99. }