helpers.go 2.5 KB

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