Session's contexts can now be updated

This is needed for DTLS restart, as mentioned in pion/webrtc#1636.

README.md

@@ -49,6 +49,7 @@ Check out the **[contributing wiki](https://github.com/pion/webrtc/wiki/Contribu
 * [Mission Liao](https://github.com/mission-liao)
 * [Orlando](https://github.com/OrlandoCo)
 * [Tarrence van As](https://github.com/tarrencev)
+* [Antoine Baché](https://github.com/Antonito)
 
 ### License
 MIT License - see [LICENSE](LICENSE) for full text

session.go

@@ -4,6 +4,7 @@ import (
 	"io"
 	"net"
 	"sync"
+	"sync/atomic"
 
 	"github.com/pion/logging"
 	"github.com/pion/transport/packetio"
@@ -17,7 +18,8 @@ type streamSession interface {
 
 type session struct {
 	localContextMutex           sync.Mutex
-	localContext, remoteContext *Context
+	localContext                *Context
+	remoteContext               atomic.Value // *Context
 	localOptions, remoteOptions []ContextOption
 
 	newStream chan readStream
@@ -106,13 +108,15 @@ func (s *session) close() error {
 }
 
 func (s *session) start(localMasterKey, localMasterSalt, remoteMasterKey, remoteMasterSalt []byte, profile ProtectionProfile, child streamSession) error {
-	var err error
-	s.localContext, err = CreateContext(localMasterKey, localMasterSalt, profile, s.localOptions...)
-	if err != nil {
-		return err
-	}
-
-	s.remoteContext, err = CreateContext(remoteMasterKey, remoteMasterSalt, profile, s.remoteOptions...)
+	err := s.UpdateContext(&Config{
+		Keys: SessionKeys{
+			LocalMasterKey:   localMasterKey,
+			LocalMasterSalt:  localMasterSalt,
+			RemoteMasterKey:  remoteMasterKey,
+			RemoteMasterSalt: remoteMasterSalt,
+		},
+		Profile: profile,
+	})
 	if err != nil {
 		return err
 	}
@@ -148,3 +152,23 @@ func (s *session) start(localMasterKey, localMasterSalt, remoteMasterKey, remote
 
 	return nil
 }
+
+// UpdateContext updates the local and remote context of the session.
+func (s *session) UpdateContext(config *Config) error {
+	localContext, err := CreateContext(config.Keys.LocalMasterKey, config.Keys.LocalMasterSalt, config.Profile, s.localOptions...)
+	if err != nil {
+		return err
+	}
+	remoteContext, err := CreateContext(config.Keys.RemoteMasterKey, config.Keys.RemoteMasterSalt, config.Profile, s.remoteOptions...)
+	if err != nil {
+		return err
+	}
+
+	s.localContextMutex.Lock()
+	s.localContext = localContext
+	s.localContextMutex.Unlock()
+
+	s.remoteContext.Store(remoteContext)
+
+	return nil
+}

session_srtcp.go

@@ -147,7 +147,9 @@ func destinationSSRC(pkts []rtcp.Packet) []uint32 {
 }
 
 func (s *SessionSRTCP) decrypt(buf []byte) error {
-	decrypted, err := s.remoteContext.DecryptRTCP(buf, buf, nil)
+	// Safe since remoteContext always contains a *Context.
+	remoteContext := s.remoteContext.Load().(*Context)
+	decrypted, err := remoteContext.DecryptRTCP(buf, buf, nil)
 	if err != nil {
 		return err
 	}

session_srtp.go

@@ -157,7 +157,9 @@ func (s *SessionSRTP) decrypt(buf []byte) error {
 		return errFailedTypeAssertion
 	}
 
-	decrypted, err := s.remoteContext.decryptRTP(buf, buf, h)
+	// Safe since remoteContext always contains a *Context.
+	remoteContext := s.remoteContext.Load().(*Context)
+	decrypted, err := remoteContext.decryptRTP(buf, buf, h)
 	if err != nil {
 		return err
 	}