Prepare for stun v2

Using github.com/gortc/stun for now.
Should be easy to change it later.

agent.go

@@ -3,7 +3,7 @@
 package ice
 
 import (
-	"bytes"
+	"fmt"
 	"math/rand"
 	"net"
 	"sort"
@@ -11,8 +11,8 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/gortc/stun"
 	"github.com/pion/logging"
-	"github.com/pion/stun"
 	"github.com/pion/transport/packetio"
 )
 
@@ -31,8 +31,6 @@ const (
 
 	// the number of outbound binding requests we cache
 	maxPendingBindingRequests = 50
-
-	stunAttrHeaderLength = 4
 )
 
 type candidatePairs []*candidatePair
@@ -48,7 +46,7 @@ func (bp byPairPriority) Less(i, j int) bool {
 }
 
 type bindingRequest struct {
-	transactionID  []byte
+	transactionID  [stun.TransactionIDSize]byte
 	destination    net.Addr
 	isUseCandidate bool
 }
@@ -562,7 +560,7 @@ func (a *Agent) sendBindingRequest(m *stun.Message, local, remote *Candidate) {
 		a.pendingBindingRequests = a.pendingBindingRequests[overflow:]
 	}
 
-	_, useCandidate := m.GetOneAttribute(stun.AttrUseCandidate)
+	useCandidate := m.Contains(stun.AttrUseCandidate)
 
 	a.pendingBindingRequests = append(a.pendingBindingRequests, bindingRequest{
 		transactionID:  m.TransactionID,
@@ -575,17 +573,13 @@ func (a *Agent) sendBindingRequest(m *stun.Message, local, remote *Candidate) {
 
 func (a *Agent) sendBindingSuccess(m *stun.Message, local, remote *Candidate) {
 	base := remote
-	if out, err := stun.Build(stun.ClassSuccessResponse, stun.MethodBinding, m.TransactionID,
-		&stun.XorMappedAddress{
-			XorAddress: stun.XorAddress{
-				IP:   base.IP,
-				Port: base.Port,
-			},
+	if out, err := stun.Build(m, stun.BindingSuccess,
+		&stun.XORMappedAddress{
+			IP:   base.IP,
+			Port: base.Port,
 		},
-		&stun.MessageIntegrity{
-			Key: []byte(a.localPwd),
-		},
-		&stun.Fingerprint{},
+		stun.NewShortTermIntegrity(a.localPwd),
+		stun.Fingerprint,
 	); err != nil {
 		a.log.Warnf("Failed to handle inbound ICE from: %s to: %s error: %s", local, remote, err)
 	} else {
@@ -595,9 +589,9 @@ func (a *Agent) sendBindingSuccess(m *stun.Message, local, remote *Candidate) {
 
 // Assert that the passed TransactionID is in our pendingBindingRequests and returns the destination
 // If the bindingRequest was valid remove it from our pending cache
-func (a *Agent) handleInboundBindingSuccess(id []byte) (bool, *bindingRequest) {
+func (a *Agent) handleInboundBindingSuccess(id [stun.TransactionIDSize]byte) (bool, *bindingRequest) {
 	for i := range a.pendingBindingRequests {
-		if bytes.Equal(a.pendingBindingRequests[i].transactionID, id) {
+		if a.pendingBindingRequests[i].transactionID == id {
 			validBindingRequest := a.pendingBindingRequests[i]
 			a.pendingBindingRequests = append(a.pendingBindingRequests[:i], a.pendingBindingRequests[i+1:]...)
 			return true, &validBindingRequest
@@ -613,31 +607,31 @@ func (a *Agent) handleInbound(m *stun.Message, local *Candidate, remote net.Addr
 		return
 	}
 
-	if m.Method != stun.MethodBinding ||
-		!(m.Class == stun.ClassSuccessResponse ||
-			m.Class == stun.ClassRequest ||
-			m.Class == stun.ClassIndication) {
-		a.log.Tracef("unhandled STUN from %s to %s class(%s) method(%s)", remote.String(), local.String(), m.Class.String(), m.Method.String())
+	if m.Type.Method != stun.MethodBinding ||
+		!(m.Type.Class == stun.ClassSuccessResponse ||
+			m.Type.Class == stun.ClassRequest ||
+			m.Type.Class == stun.ClassIndication) {
+		a.log.Tracef("unhandled STUN from %s to %s class(%s) method(%s)", remote, local, m.Type.Class, m.Type.Method)
 		return
 	}
 
 	if a.isControlling {
-		if _, isControlling := m.GetOneAttribute(stun.AttrIceControlling); isControlling {
+		if m.Contains(stun.AttrICEControlling) {
 			a.log.Debug("inbound isControlling && a.isControlling == true")
 			return
-		} else if _, useCandidate := m.GetOneAttribute(stun.AttrUseCandidate); useCandidate {
+		} else if m.Contains(stun.AttrUseCandidate) {
 			a.log.Debug("useCandidate && a.isControlling == true")
 			return
 		}
 	} else {
-		if _, isControlled := m.GetOneAttribute(stun.AttrIceControlled); isControlled {
+		if m.Contains(stun.AttrICEControlled) {
 			a.log.Debug("inbound isControlled && a.isControlling == false")
 			return
 		}
 	}
 
 	remoteCandidate := a.findRemoteCandidate(local.NetworkType, remote)
-	if m.Class == stun.ClassSuccessResponse {
+	if m.Type.Class == stun.ClassSuccessResponse {
 		if err = assertInboundMessageIntegrity(m, []byte(a.remotePwd)); err != nil {
 			a.log.Warnf("discard message from (%s), %v", remote, err)
 			return
@@ -649,7 +643,7 @@ func (a *Agent) handleInbound(m *stun.Message, local *Candidate, remote net.Addr
 		}
 
 		a.selector.HandleSucessResponse(m, local, remoteCandidate, remote)
-	} else if m.Class == stun.ClassRequest {
+	} else if m.Type.Class == stun.ClassRequest {
 		if err = assertInboundUsername(m, a.localUfrag+":"+a.remoteUfrag); err != nil {
 			a.log.Warnf("discard message from (%s), %v", remote, err)
 			return
@@ -721,3 +715,41 @@ func (a *Agent) getSelectedPair() (*candidatePair, error) {
 
 	return out, nil
 }
+
+// Role represents ICE agent role, which can be controlling or controlled.
+type Role byte
+
+// UnmarshalText implements TextUnmarshaler.
+func (r *Role) UnmarshalText(text []byte) error {
+	switch string(text) {
+	case "controlling":
+		*r = Controlling
+	case "controlled":
+		*r = Controlled
+	default:
+		return fmt.Errorf("unknown role %q", text)
+	}
+	return nil
+}
+
+// MarshalText implements TextMarshaler.
+func (r Role) MarshalText() (text []byte, err error) {
+	return []byte(r.String()), nil
+}
+
+func (r Role) String() string {
+	switch r {
+	case Controlling:
+		return "controlling"
+	case Controlled:
+		return "controlled"
+	default:
+		return "unknown"
+	}
+}
+
+// Possible ICE agent roles.
+const (
+	Controlling Role = iota
+	Controlled
+)

agent_test.go

@@ -6,7 +6,8 @@ import (
 	"testing"
 	"time"
 
-	"github.com/pion/stun"
+	"github.com/gortc/stun"
+	"github.com/pion/logging"
 	"github.com/pion/transport/test"
 )
 
@@ -205,15 +206,13 @@ func TestHandlePeerReflexive(t *testing.T) {
 
 			remote := &net.UDPAddr{IP: net.ParseIP("172.17.0.3"), Port: 999}
 
-			msg, err := stun.Build(stun.ClassRequest, stun.MethodBinding, stun.GenerateTransactionID(),
-				&stun.Username{Username: a.localUfrag + ":" + a.remoteUfrag},
-				&stun.UseCandidate{},
-				&stun.IceControlling{TieBreaker: a.tieBreaker},
-				&stun.Priority{Priority: local.Priority()},
-				&stun.MessageIntegrity{
-					Key: []byte(a.localPwd),
-				},
-				&stun.Fingerprint{},
+			msg, err := stun.Build(stun.BindingRequest, stun.TransactionID,
+				stun.NewUsername(a.localUfrag+":"+a.remoteUfrag),
+				UseCandidate,
+				AttrControlling(a.tieBreaker),
+				PriorityAttr(local.Priority()),
+				stun.NewShortTermIntegrity(a.localPwd),
+				stun.Fingerprint,
 			)
 			if err != nil {
 				t.Fatal(err)
@@ -282,8 +281,10 @@ func TestHandlePeerReflexive(t *testing.T) {
 		var config AgentConfig
 		runAgentTest(t, &config, func(a *Agent) {
 			a.selector = &controllingSelector{agent: a, log: a.log}
+			tID := [stun.TransactionIDSize]byte{}
+			copy(tID[:], []byte("ABC"))
 			a.pendingBindingRequests = []bindingRequest{
-				{[]byte("ABC"), &net.UDPAddr{}, false},
+				{tID, &net.UDPAddr{}, false},
 			}
 
 			local, err := NewCandidateHost("udp", net.ParseIP("192.168.0.2"), 777, 1)
@@ -293,11 +294,9 @@ func TestHandlePeerReflexive(t *testing.T) {
 			}
 
 			remote := &net.UDPAddr{IP: net.ParseIP("172.17.0.3"), Port: 999}
-			msg, err := stun.Build(stun.ClassSuccessResponse, stun.MethodBinding, []byte("ABC"),
-				&stun.MessageIntegrity{
-					Key: []byte(a.remotePwd),
-				},
-				&stun.Fingerprint{},
+			msg, err := stun.Build(stun.BindingSuccess, stun.NewTransactionIDSetter(tID),
+				stun.NewShortTermIntegrity(a.remotePwd),
+				stun.Fingerprint,
 			)
 			if err != nil {
 				t.Fatal(err)
@@ -321,6 +320,7 @@ func TestConnectivityOnStartup(t *testing.T) {
 		Urls:             []*URL{},
 		NetworkTypes:     supportedNetworkTypes,
 		taskLoopInterval: time.Hour,
+		LoggerFactory:    logging.NewDefaultLoggerFactory(),
 	}
 
 	aNotifier, aConnected := onConnected()
@@ -352,12 +352,10 @@ func TestConnectivityOnStartup(t *testing.T) {
 
 func TestInboundValidity(t *testing.T) {
 	buildMsg := func(class stun.MessageClass, username, key string) *stun.Message {
-		msg, err := stun.Build(class, stun.MethodBinding, stun.GenerateTransactionID(),
-			&stun.Username{Username: username},
-			&stun.MessageIntegrity{
-				Key: []byte(key),
-			},
-			&stun.Fingerprint{},
+		msg, err := stun.Build(stun.NewType(stun.MethodBinding, class), stun.TransactionID,
+			stun.NewUsername(username),
+			stun.NewShortTermIntegrity(key),
+			stun.Fingerprint,
 		)
 		if err != nil {
 			t.Fatal(err)
@@ -437,11 +435,9 @@ func TestInboundValidity(t *testing.T) {
 		var config AgentConfig
 		runAgentTest(t, &config, func(a *Agent) {
 			a.selector = &controllingSelector{agent: a, log: a.log}
-			msg, err := stun.Build(stun.ClassRequest, stun.MethodBinding, stun.GenerateTransactionID(),
-				&stun.Username{Username: a.localUfrag + ":" + a.remoteUfrag},
-				&stun.MessageIntegrity{
-					Key: []byte(a.localPwd),
-				},
+			msg, err := stun.Build(stun.BindingRequest, stun.TransactionID,
+				stun.NewUsername(a.localUfrag+":"+a.remoteUfrag),
+				stun.NewShortTermIntegrity(a.localPwd),
 			)
 			if err != nil {
 				t.Fatal(err)
@@ -467,11 +463,11 @@ func TestInboundValidity(t *testing.T) {
 		}
 
 		remote := &net.UDPAddr{IP: net.ParseIP("172.17.0.3"), Port: 999}
-		msg, err := stun.Build(stun.ClassSuccessResponse, stun.MethodBinding, []byte("ABC"),
-			&stun.MessageIntegrity{
-				Key: []byte(a.remotePwd),
-			},
-			&stun.Fingerprint{},
+		tID := [stun.TransactionIDSize]byte{}
+		copy(tID[:], []byte("ABC"))
+		msg, err := stun.Build(stun.BindingSuccess, stun.NewTransactionIDSetter(tID),
+			stun.NewShortTermIntegrity(a.remotePwd),
+			stun.Fingerprint,
 		)
 		if err != nil {
 			t.Fatal(err)

candidate.go

@@ -6,7 +6,7 @@ import (
 	"sync"
 	"time"
 
-	"github.com/pion/stun"
+	"github.com/gortc/stun"
 )
 
 const (
@@ -140,10 +140,13 @@ func (c *Candidate) recvLoop() {
 			return
 		}
 
-		if stun.IsSTUN(buffer[:n]) {
-			var m *stun.Message
-			m, err = stun.NewMessage(buffer[:n])
-			if err != nil {
+		if stun.IsMessage(buffer[:n]) {
+			m := &stun.Message{
+				Raw: make([]byte, n),
+			}
+			// Explicitly copy raw buffer so Message can own the memory.
+			copy(m.Raw, buffer[:n])
+			if err := m.Decode(); err != nil {
 				log.Warnf("Failed to handle decode ICE from %s to %s: %v", c.addr(), srcAddr, err)
 				continue
 			}

candidatepair.go

@@ -3,7 +3,7 @@ package ice
 import (
 	"fmt"
 
-	"github.com/pion/stun"
+	"github.com/gortc/stun"
 )
 
 func newCandidatePair(local, remote *Candidate, controlling bool) *candidatePair {
@@ -84,12 +84,10 @@ func (p *candidatePair) Write(b []byte) (int, error) {
 
 // keepaliveCandidate sends a STUN Binding Indication to the remote candidate
 func (a *Agent) keepaliveCandidate(local, remote *Candidate) {
-	msg, err := stun.Build(stun.ClassIndication, stun.MethodBinding, stun.GenerateTransactionID(),
-		&stun.Username{Username: a.remoteUfrag + ":" + a.localUfrag},
-		&stun.MessageIntegrity{
-			Key: []byte(a.remotePwd),
-		},
-		&stun.Fingerprint{},
+	msg, err := stun.Build(stun.NewType(stun.MethodBinding, stun.ClassIndication), stun.TransactionID,
+		stun.NewUsername(a.remoteUfrag+":"+a.localUfrag),
+		stun.NewShortTermIntegrity(a.remotePwd),
+		stun.Fingerprint,
 	)
 
 	if err != nil {
@@ -101,7 +99,7 @@ func (a *Agent) keepaliveCandidate(local, remote *Candidate) {
 }
 
 func (a *Agent) sendSTUN(msg *stun.Message, local, remote *Candidate) {
-	_, err := local.writeTo(msg.Pack(), remote)
+	_, err := local.writeTo(msg.Raw, remote)
 	if err != nil {
 		a.log.Tracef("failed to send STUN message: %s", err)
 	}

gather.go

@@ -5,6 +5,7 @@ import (
 	"net"
 	"time"
 
+	// TODO(ar): Merge
 	"github.com/pion/stun"
 )
 

go.mod

@@ -3,6 +3,7 @@ module github.com/pion/ice
 go 1.12
 
 require (
+	github.com/gortc/stun v1.19.0
 	github.com/pion/logging v0.2.1
 	github.com/pion/stun v0.2.2
 	github.com/pion/transport v0.7.0

go.sum

@@ -1,5 +1,7 @@
 github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/gortc/stun v1.19.0 h1:6qy7zGGk0tdMOdEzK7hLeAVZEHllJC8+OOBpPAyjY1c=
+github.com/gortc/stun v1.19.0/go.mod h1:dZ0O/fYCkg9Z0Pvl6WDpNhRFTAU0X1CPOsJiZqn6EHo=
 github.com/pion/logging v0.2.1 h1:LwASkBKZ+2ysGJ+jLv1E/9H1ge0k1nTfi1X+5zirkDk=
 github.com/pion/logging v0.2.1/go.mod h1:k0/tDVsRCX2Mb2ZEmTqNa7CWsQPc+YYCB7Q+5pahoms=
 github.com/pion/stun v0.2.2 h1:0IJCwJFOdEmHzz4oxl9SBGLlJbnNbF+0h6XSOmuE034=

icecontrol.go

@@ -0,0 +1,83 @@
+package ice
+
+import "github.com/gortc/stun"
+
+// tiebreaker is common helper for ICE-{CONTROLLED,CONTROLLING}
+// and represents the so-called tiebreaker number.
+type tiebreaker uint64
+
+const tiebreakerSize = 8 // 64 bit
+
+// AddToAs adds tiebreaker value to m as t attribute.
+func (a tiebreaker) AddToAs(m *stun.Message, t stun.AttrType) error {
+	v := make([]byte, tiebreakerSize)
+	bin.PutUint64(v, uint64(a))
+	m.Add(t, v)
+	return nil
+}
+
+// GetFromAs decodes tiebreaker value in message getting it as for t type.
+func (a *tiebreaker) GetFromAs(m *stun.Message, t stun.AttrType) error {
+	v, err := m.Get(t)
+	if err != nil {
+		return err
+	}
+	if err = stun.CheckSize(t, len(v), tiebreakerSize); err != nil {
+		return err
+	}
+	*a = tiebreaker(bin.Uint64(v))
+	return nil
+}
+
+// AttrControlled represents ICE-CONTROLLED attribute.
+type AttrControlled uint64
+
+// AddTo adds ICE-CONTROLLED to message.
+func (c AttrControlled) AddTo(m *stun.Message) error {
+	return tiebreaker(c).AddToAs(m, stun.AttrICEControlled)
+}
+
+// GetFrom decodes ICE-CONTROLLED from message.
+func (c *AttrControlled) GetFrom(m *stun.Message) error {
+	return (*tiebreaker)(c).GetFromAs(m, stun.AttrICEControlled)
+}
+
+// AttrControlling represents ICE-CONTROLLING attribute.
+type AttrControlling uint64
+
+// AddTo adds ICE-CONTROLLING to message.
+func (c AttrControlling) AddTo(m *stun.Message) error {
+	return tiebreaker(c).AddToAs(m, stun.AttrICEControlling)
+}
+
+// GetFrom decodes ICE-CONTROLLING from message.
+func (c *AttrControlling) GetFrom(m *stun.Message) error {
+	return (*tiebreaker)(c).GetFromAs(m, stun.AttrICEControlling)
+}
+
+// AttrControl is helper that wraps ICE-{CONTROLLED,CONTROLLING}.
+type AttrControl struct {
+	Role       Role
+	Tiebreaker uint64
+}
+
+// AddTo adds ICE-CONTROLLED or ICE-CONTROLLING attribute depending on Role.
+func (c AttrControl) AddTo(m *stun.Message) error {
+	if c.Role == Controlling {
+		return tiebreaker(c.Tiebreaker).AddToAs(m, stun.AttrICEControlling)
+	}
+	return tiebreaker(c.Tiebreaker).AddToAs(m, stun.AttrICEControlled)
+}
+
+// GetFrom decodes Role and Tiebreaker value from message.
+func (c *AttrControl) GetFrom(m *stun.Message) error {
+	if m.Contains(stun.AttrICEControlling) {
+		c.Role = Controlling
+		return (*tiebreaker)(&c.Tiebreaker).GetFromAs(m, stun.AttrICEControlling)
+	}
+	if m.Contains(stun.AttrICEControlled) {
+		c.Role = Controlled
+		return (*tiebreaker)(&c.Tiebreaker).GetFromAs(m, stun.AttrICEControlled)
+	}
+	return stun.ErrAttributeNotFound
+}

icecontrol_test.go

@@ -0,0 +1,139 @@
+package ice
+
+import (
+	"testing"
+
+	"github.com/gortc/stun"
+)
+
+func TestControlled_GetFrom(t *testing.T) {
+	m := new(stun.Message)
+	var c AttrControlled
+	if err := c.GetFrom(m); err != stun.ErrAttributeNotFound {
+		t.Error("unexpected error")
+	}
+	if err := m.Build(stun.BindingRequest, &c); err != nil {
+		t.Error(err)
+	}
+	m1 := new(stun.Message)
+	if _, err := m1.Write(m.Raw); err != nil {
+		t.Error(err)
+	}
+	var c1 AttrControlled
+	if err := c1.GetFrom(m1); err != nil {
+		t.Error(err)
+	}
+	if c1 != c {
+		t.Error("not equal")
+	}
+	t.Run("IncorrectSize", func(t *testing.T) {
+		m3 := new(stun.Message)
+		m3.Add(stun.AttrICEControlled, make([]byte, 100))
+		var c2 AttrControlled
+		if err := c2.GetFrom(m3); !stun.IsAttrSizeInvalid(err) {
+			t.Error("should error")
+		}
+	})
+}
+
+func TestControlling_GetFrom(t *testing.T) {
+	m := new(stun.Message)
+	var c AttrControlling
+	if err := c.GetFrom(m); err != stun.ErrAttributeNotFound {
+		t.Error("unexpected error")
+	}
+	if err := m.Build(stun.BindingRequest, &c); err != nil {
+		t.Error(err)
+	}
+	m1 := new(stun.Message)
+	if _, err := m1.Write(m.Raw); err != nil {
+		t.Error(err)
+	}
+	var c1 AttrControlling
+	if err := c1.GetFrom(m1); err != nil {
+		t.Error(err)
+	}
+	if c1 != c {
+		t.Error("not equal")
+	}
+	t.Run("IncorrectSize", func(t *testing.T) {
+		m3 := new(stun.Message)
+		m3.Add(stun.AttrICEControlling, make([]byte, 100))
+		var c2 AttrControlling
+		if err := c2.GetFrom(m3); !stun.IsAttrSizeInvalid(err) {
+			t.Error("should error")
+		}
+	})
+}
+
+func TestControl_GetFrom(t *testing.T) {
+	t.Run("Blank", func(t *testing.T) {
+		m := new(stun.Message)
+		var c AttrControl
+		if err := c.GetFrom(m); err != stun.ErrAttributeNotFound {
+			t.Error("unexpected error")
+		}
+	})
+	t.Run("Controlling", func(t *testing.T) {
+		m := new(stun.Message)
+		var c AttrControl
+		if err := c.GetFrom(m); err != stun.ErrAttributeNotFound {
+			t.Error("unexpected error")
+		}
+		c.Role = Controlling
+		c.Tiebreaker = 4321
+		if err := m.Build(stun.BindingRequest, &c); err != nil {
+			t.Error(err)
+		}
+		m1 := new(stun.Message)
+		if _, err := m1.Write(m.Raw); err != nil {
+			t.Error(err)
+		}
+		var c1 AttrControl
+		if err := c1.GetFrom(m1); err != nil {
+			t.Error(err)
+		}
+		if c1 != c {
+			t.Error("not equal")
+		}
+		t.Run("IncorrectSize", func(t *testing.T) {
+			m3 := new(stun.Message)
+			m3.Add(stun.AttrICEControlling, make([]byte, 100))
+			var c2 AttrControl
+			if err := c2.GetFrom(m3); !stun.IsAttrSizeInvalid(err) {
+				t.Error("should error")
+			}
+		})
+	})
+	t.Run("Controlled", func(t *testing.T) {
+		m := new(stun.Message)
+		var c AttrControl
+		if err := c.GetFrom(m); err != stun.ErrAttributeNotFound {
+			t.Error("unexpected error")
+		}
+		c.Role = Controlled
+		c.Tiebreaker = 1234
+		if err := m.Build(stun.BindingRequest, &c); err != nil {
+			t.Error(err)
+		}
+		m1 := new(stun.Message)
+		if _, err := m1.Write(m.Raw); err != nil {
+			t.Error(err)
+		}
+		var c1 AttrControl
+		if err := c1.GetFrom(m1); err != nil {
+			t.Error(err)
+		}
+		if c1 != c {
+			t.Error("not equal")
+		}
+		t.Run("IncorrectSize", func(t *testing.T) {
+			m3 := new(stun.Message)
+			m3.Add(stun.AttrICEControlling, make([]byte, 100))
+			var c2 AttrControl
+			if err := c2.GetFrom(m3); !stun.IsAttrSizeInvalid(err) {
+				t.Error("should error")
+			}
+		})
+	})
+}

priority.go

@@ -0,0 +1,29 @@
+package ice
+
+import "github.com/gortc/stun"
+
+// PriorityAttr represents PRIORITY attribute.
+type PriorityAttr uint32
+
+const prioritySize = 4 // 32 bit
+
+// AddTo adds PRIORITY attribute to message.
+func (p PriorityAttr) AddTo(m *stun.Message) error {
+	v := make([]byte, prioritySize)
+	bin.PutUint32(v, uint32(p))
+	m.Add(stun.AttrPriority, v)
+	return nil
+}
+
+// GetFrom decodes PRIORITY attribute from message.
+func (p *PriorityAttr) GetFrom(m *stun.Message) error {
+	v, err := m.Get(stun.AttrPriority)
+	if err != nil {
+		return err
+	}
+	if err = stun.CheckSize(stun.AttrPriority, len(v), prioritySize); err != nil {
+		return err
+	}
+	*p = PriorityAttr(bin.Uint32(v))
+	return nil
+}

priority_test.go

@@ -0,0 +1,37 @@
+package ice
+
+import (
+	"testing"
+
+	"github.com/gortc/stun"
+)
+
+func TestPriority_GetFrom(t *testing.T) {
+	m := new(stun.Message)
+	var p PriorityAttr
+	if err := p.GetFrom(m); err != stun.ErrAttributeNotFound {
+		t.Error("unexpected error")
+	}
+	if err := m.Build(stun.BindingRequest, &p); err != nil {
+		t.Error(err)
+	}
+	m1 := new(stun.Message)
+	if _, err := m1.Write(m.Raw); err != nil {
+		t.Error(err)
+	}
+	var p1 PriorityAttr
+	if err := p1.GetFrom(m1); err != nil {
+		t.Error(err)
+	}
+	if p1 != p {
+		t.Error("not equal")
+	}
+	t.Run("IncorrectSize", func(t *testing.T) {
+		m3 := new(stun.Message)
+		m3.Add(stun.AttrPriority, make([]byte, 100))
+		var p2 PriorityAttr
+		if err := p2.GetFrom(m3); !stun.IsAttrSizeInvalid(err) {
+			t.Error("should error")
+		}
+	})
+}

selection.go

@@ -3,8 +3,8 @@ package ice
 import (
 	"net"
 
+	"github.com/gortc/stun"
 	"github.com/pion/logging"
-	"github.com/pion/stun"
 )
 
 type pairCandidateSelector interface {
@@ -40,21 +40,17 @@ func (s *controllingSelector) ContactCandidates() {
 }
 
 func (s *controllingSelector) nominatePair(pair *candidatePair) {
-	transactionID := stun.GenerateTransactionID()
-
 	// The controlling agent MUST include the USE-CANDIDATE attribute in
 	// order to nominate a candidate pair (Section 8.1.1).  The controlled
 	// agent MUST NOT include the USE-CANDIDATE attribute in a Binding
 	// request.
-	msg, err := stun.Build(stun.ClassRequest, stun.MethodBinding, transactionID,
-		&stun.Username{Username: s.agent.remoteUfrag + ":" + s.agent.localUfrag},
-		&stun.UseCandidate{},
-		&stun.IceControlling{TieBreaker: s.agent.tieBreaker},
-		&stun.Priority{Priority: pair.local.Priority()},
-		&stun.MessageIntegrity{
-			Key: []byte(s.agent.remotePwd),
-		},
-		&stun.Fingerprint{},
+	msg, err := stun.Build(stun.BindingRequest, stun.TransactionID,
+		stun.NewUsername(s.agent.remoteUfrag+":"+s.agent.localUfrag),
+		UseCandidate,
+		AttrControlling(s.agent.tieBreaker),
+		PriorityAttr(pair.local.Priority()),
+		stun.NewShortTermIntegrity(s.agent.remotePwd),
+		stun.Fingerprint,
 	)
 
 	if err != nil {
@@ -101,16 +97,12 @@ func (s *controllingSelector) HandleSucessResponse(m *stun.Message, local, remot
 }
 
 func (s *controllingSelector) PingCandidate(local, remote *Candidate) {
-	transactionID := stun.GenerateTransactionID()
-
-	msg, err := stun.Build(stun.ClassRequest, stun.MethodBinding, transactionID,
-		&stun.Username{Username: s.agent.remoteUfrag + ":" + s.agent.localUfrag},
-		&stun.IceControlling{TieBreaker: s.agent.tieBreaker},
-		&stun.Priority{Priority: local.Priority()},
-		&stun.MessageIntegrity{
-			Key: []byte(s.agent.remotePwd),
-		},
-		&stun.Fingerprint{},
+	msg, err := stun.Build(stun.BindingRequest, stun.TransactionID,
+		stun.NewUsername(s.agent.remoteUfrag+":"+s.agent.localUfrag),
+		AttrControlling(s.agent.tieBreaker),
+		PriorityAttr(local.Priority()),
+		stun.NewShortTermIntegrity(s.agent.remotePwd),
+		stun.Fingerprint,
 	)
 
 	if err != nil {
@@ -141,16 +133,12 @@ func (s *controlledSelector) ContactCandidates() {
 }
 
 func (s *controlledSelector) PingCandidate(local, remote *Candidate) {
-	transactionID := stun.GenerateTransactionID()
-
-	msg, err := stun.Build(stun.ClassRequest, stun.MethodBinding, transactionID,
-		&stun.Username{Username: s.agent.remoteUfrag + ":" + s.agent.localUfrag},
-		&stun.IceControlled{TieBreaker: s.agent.tieBreaker},
-		&stun.Priority{Priority: local.Priority()},
-		&stun.MessageIntegrity{
-			Key: []byte(s.agent.remotePwd),
-		},
-		&stun.Fingerprint{},
+	msg, err := stun.Build(stun.BindingRequest, stun.TransactionID,
+		stun.NewUsername(s.agent.remoteUfrag+":"+s.agent.localUfrag),
+		AttrControlled(s.agent.tieBreaker),
+		PriorityAttr(local.Priority()),
+		stun.NewShortTermIntegrity(s.agent.remotePwd),
+		stun.Fingerprint,
 	)
 
 	if err != nil {
@@ -189,9 +177,7 @@ func (s *controlledSelector) HandleSucessResponse(m *stun.Message, local, remote
 }
 
 func (s *controlledSelector) HandleBindingRequest(m *stun.Message, local, remote *Candidate) {
-	_, useCandidate := m.GetOneAttribute(stun.AttrUseCandidate)
-
-	if useCandidate {
+	if m.Contains(stun.AttrUseCandidate) {
 		// https://tools.ietf.org/html/rfc8445#section-7.3.1.5
 		p := s.agent.findValidPair(local, remote)
 

stun.go

@@ -1,65 +1,28 @@
 package ice
 
 import (
-	"bytes"
 	"encoding/binary"
 	"fmt"
 
-	"github.com/pion/stun"
+	"github.com/gortc/stun"
 )
 
-func assertInboundUsername(m *stun.Message, expectedUsername string) error {
-	usernameAttr := &stun.Username{}
-	usernameRawAttr, usernameFound := m.GetOneAttribute(stun.AttrUsername)
+// bin is shorthand for BigEndian.
+var bin = binary.BigEndian
 
-	if !usernameFound {
-		return fmt.Errorf("inbound packet missing Username")
-	} else if err := usernameAttr.Unpack(m, usernameRawAttr); err != nil {
+func assertInboundUsername(m *stun.Message, expectedUsername string) error {
+	var username stun.Username
+	if err := username.GetFrom(m); err != nil {
 		return err
 	}
-
-	if usernameAttr.Username != expectedUsername {
-		return fmt.Errorf("username mismatch expected(%x) actual(%x)", expectedUsername, usernameAttr.Username)
+	if string(username) != expectedUsername {
+		return fmt.Errorf("username mismatch expected(%x) actual(%x)", expectedUsername, string(username))
 	}
 
 	return nil
 }
 
 func assertInboundMessageIntegrity(m *stun.Message, key []byte) error {
-	messageIntegrityAttr := &stun.MessageIntegrity{}
-	messageIntegrityRawAttr, messageIntegrityAttrFound := m.GetOneAttribute(stun.AttrMessageIntegrity)
-
-	if !messageIntegrityAttrFound {
-		return fmt.Errorf("inbound packet missing MessageIntegrity")
-	} else if err := messageIntegrityAttr.Unpack(m, messageIntegrityRawAttr); err != nil {
-		return err
-	}
-
-	tailLength := messageIntegrityRawAttr.Length + stunAttrHeaderLength
-	rawCopy := make([]byte, len(m.Raw))
-	copy(rawCopy, m.Raw)
-
-	// If we have a fingerprint we need to exclude it from the MessageIntegrity computation
-	if rawFingerprint, hasFingerprint := m.GetOneAttribute(stun.AttrFingerprint); hasFingerprint {
-		fingerprintLength := rawFingerprint.Length + stunAttrHeaderLength
-		tailLength += fingerprintLength
-
-		// Rewrite the packet header to be new length (excluding values we don't care about)
-		currLength := binary.BigEndian.Uint16(rawCopy[2:4])
-		binary.BigEndian.PutUint16(rawCopy[2:], currLength-fingerprintLength)
-	}
-
-	lengthToHash := len(rawCopy) - int(tailLength)
-	if lengthToHash < 1 {
-		return fmt.Errorf("unable to assert MessageIntegrity, length calculation failed (%d)", lengthToHash)
-	}
-
-	computedMessageIntegrity, err := stun.MessageIntegrityCalculateHMAC(key, rawCopy[:lengthToHash])
-	if err != nil {
-		return err
-	} else if !bytes.Equal(computedMessageIntegrity, messageIntegrityRawAttr.Value) {
-		return fmt.Errorf("messageIntegrity mismatch expected(%x) actual(%x)", computedMessageIntegrity, messageIntegrityRawAttr.Value)
-	}
-
-	return nil
+	messageIntegrityAttr := stun.MessageIntegrity(key)
+	return messageIntegrityAttr.Check(m)
 }

transport.go

@@ -6,7 +6,7 @@ import (
 	"net"
 	"time"
 
-	"github.com/pion/stun"
+	"github.com/gortc/stun"
 )
 
 // Dial connects to the remote agent, acting as the controlling ice agent.
@@ -73,7 +73,7 @@ func (c *Conn) Write(p []byte) (int, error) {
 		return 0, err
 	}
 
-	if stun.IsSTUN(p) {
+	if stun.IsMessage(p) {
 		return 0, errors.New("the ICE conn can't write STUN messages")
 	}
 

usecandidate.go

@@ -0,0 +1,21 @@
+package ice
+
+import "github.com/gortc/stun"
+
+// UseCandidateAttr represents USE-CANDIDATE attribute.
+type UseCandidateAttr struct{}
+
+// AddTo adds USE-CANDIDATE attribute to message.
+func (UseCandidateAttr) AddTo(m *stun.Message) error {
+	m.Add(stun.AttrUseCandidate, nil)
+	return nil
+}
+
+// IsSet returns true if USE-CANDIDATE attribute is set.
+func (UseCandidateAttr) IsSet(m *stun.Message) bool {
+	_, err := m.Get(stun.AttrUseCandidate)
+	return err == nil
+}
+
+// UseCandidate is shorthand for UseCandidateAttr.
+var UseCandidate UseCandidateAttr

usecandidate_test.go

@@ -0,0 +1,24 @@
+package ice
+
+import (
+	"testing"
+
+	"github.com/gortc/stun"
+)
+
+func TestUseCandidateAttr_AddTo(t *testing.T) {
+	m := new(stun.Message)
+	if UseCandidate.IsSet(m) {
+		t.Error("should not be set")
+	}
+	if err := m.Build(stun.BindingRequest, UseCandidate); err != nil {
+		t.Error(err)
+	}
+	m1 := new(stun.Message)
+	if _, err := m1.Write(m.Raw); err != nil {
+		t.Error(err)
+	}
+	if !UseCandidate.IsSet(m1) {
+		t.Error("should be set")
+	}
+}