Repository: pion/dtls
Pull request #66 (2019-05-22 07:58:41)

Allow user to specify PSK
Add configuration option for user to pass PSK. When a user
passes a PSK then we only allow CipherSuites that do PSK.

If user passes a PSK and a certificate return an error.

Relates to #45


codecov-io: # [Codecov](https://codecov.io/gh/pion/dtls/pull/66?src=pr&el=h1) Report
> Merging [#66](https://codecov.io/gh/pion/dtls/pull/66?src=pr&el=desc) into [master](https://codecov.io/gh/pion/dtls/commit/1f6822ac4d23ebce40356edb0baf2cf415217e4a?src=pr&el=desc) will **increase** coverage by `0.61%`.
> The diff coverage is `82.43%`.

[![Impacted file tree graph](https://codecov.io/gh/pion/dtls/pull/66/graphs/tree.svg?width=650&token=bcCB2jB88I&height=150&src=pr)](https://codecov.io/gh/pion/dtls/pull/66?src=pr&el=tree)

```diff
@@            Coverage Diff            @@
##           master     #66      +/-   ##
=========================================
+ Coverage   70.49%   71.1%   +0.61%     
=========================================
  Files          48      48              
  Lines        2596    2627      +31     
=========================================
+ Hits         1830    1868      +38     
+ Misses        554     551       -3     
+ Partials      212     208       -4
```


| [Impacted Files](https://codecov.io/gh/pion/dtls/pull/66?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [...pher\_suite\_tls\_ecdhe\_ecdsa\_with\_aes\_256\_cbc\_sha.go](https://codecov.io/gh/pion/dtls/pull/66/diff?src=pr&el=tree#diff-Y2lwaGVyX3N1aXRlX3Rsc19lY2RoZV9lY2RzYV93aXRoX2Flc18yNTZfY2JjX3NoYS5nbw==) | `84.21% <100%> (+6.43%)` | :arrow_up: |
| [cipher\_suite.go](https://codecov.io/gh/pion/dtls/pull/66/diff?src=pr&el=tree#diff-Y2lwaGVyX3N1aXRlLmdv) | `100% <100%> (ø)` | :arrow_up: |
| [...r\_suite\_tls\_ecdhe\_ecdsa\_with\_aes\_128\_gcm\_sha256.go](https://codecov.io/gh/pion/dtls/pull/66/diff?src=pr&el=tree#diff-Y2lwaGVyX3N1aXRlX3Rsc19lY2RoZV9lY2RzYV93aXRoX2Flc18xMjhfZ2NtX3NoYTI1Ni5nbw==) | `81.25% <100%> (+7.91%)` | :arrow_up: |
| [conn.go](https://codecov.io/gh/pion/dtls/pull/66/diff?src=pr&el=tree#diff-Y29ubi5nbw==) | `76.04% <68.29%> (+0.52%)` | :arrow_up: |
| [flight.go](https://codecov.io/gh/pion/dtls/pull/66/diff?src=pr&el=tree#diff-ZmxpZ2h0Lmdv) | `65.71% <0%> (+2.85%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/pion/dtls/pull/66?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/pion/dtls/pull/66?src=pr&el=footer). Last update [1f6822a...2000f89](https://codecov.io/gh/pion/dtls/pull/66?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Patch:
diff --git a/cipher_suite.go b/cipher_suite.go
index 53ae3a6..c7abb83 100644
--- a/cipher_suite.go
+++ b/cipher_suite.go
@@ -2,6 +2,7 @@ package dtls
 
 import (
 	"encoding/binary"
+	"fmt"
 	"hash"
 )
 
@@ -24,6 +25,7 @@ type cipherSuite interface {
 	ID() CipherSuiteID
 	certificateType() clientCertificateType
 	hashFunc() func() hash.Hash
+	isPSK() bool
 
 	// Generate the internal encryption state
 	init(masterSecret, clientRandom, serverRandom []byte, isClient bool) error
@@ -85,3 +87,46 @@ func encodeCipherSuites(c []cipherSuite) []byte {
 
 	return out
 }
+
+func parseCipherSuites(userSelectedSuites []CipherSuiteID, excludePSK, excludeNonPSK bool) ([]cipherSuite, error) {
+	cipherSuitesForIDs := func(ids []CipherSuiteID) ([]cipherSuite, error) {
+		cipherSuites := []cipherSuite{}
+		for _, id := range ids {
+			c := cipherSuiteForID(id)
+			if c == nil {
+				return nil, fmt.Errorf("CipherSuite with id(%d) is not valid", id)
+			}
+			cipherSuites = append(cipherSuites, c)
+		}
+		return cipherSuites, nil
+	}
+
+	var (
+		cipherSuites []cipherSuite
+		err          error
+		i            int
+	)
+	if len(userSelectedSuites) != 0 {
+		cipherSuites, err = cipherSuitesForIDs(userSelectedSuites)
+		if err != nil {
+			return nil, err
+		}
+	} else {
+		cipherSuites = defaultCipherSuites()
+	}
+
+	for _, c := range cipherSuites {
+		if excludePSK && c.isPSK() || excludeNonPSK && !c.isPSK() {
+			continue
+		}
+		cipherSuites[i] = c
+		i++
+	}
+
+	cipherSuites = cipherSuites[:i]
+	if len(cipherSuites) == 0 {
+		return nil, errNoAvailableCipherSuites
+	}
+
+	return cipherSuites, nil
+}
diff --git a/cipher_suite_tls_ecdhe_ecdsa_with_aes_128_gcm_sha256.go b/cipher_suite_tls_ecdhe_ecdsa_with_aes_128_gcm_sha256.go
index 566453d..7dd7486 100644
--- a/cipher_suite_tls_ecdhe_ecdsa_with_aes_128_gcm_sha256.go
+++ b/cipher_suite_tls_ecdhe_ecdsa_with_aes_128_gcm_sha256.go
@@ -26,6 +26,10 @@ func (c cipherSuiteTLSEcdheEcdsaWithAes128GcmSha256) hashFunc() func() hash.Hash
 	return sha256.New
 }
 
+func (c cipherSuiteTLSEcdheEcdsaWithAes128GcmSha256) isPSK() bool {
+	return false
+}
+
 func (c *cipherSuiteTLSEcdheEcdsaWithAes128GcmSha256) init(masterSecret, clientRandom, serverRandom []byte, isClient bool) error {
 	const (
 		prfMacLen = 0
diff --git a/cipher_suite_tls_ecdhe_ecdsa_with_aes_256_cbc_sha.go b/cipher_suite_tls_ecdhe_ecdsa_with_aes_256_cbc_sha.go
index 028013f..5d5771f 100644
--- a/cipher_suite_tls_ecdhe_ecdsa_with_aes_256_cbc_sha.go
+++ b/cipher_suite_tls_ecdhe_ecdsa_with_aes_256_cbc_sha.go
@@ -26,6 +26,10 @@ func (c cipherSuiteTLSEcdheEcdsaWithAes256CbcSha) hashFunc() func() hash.Hash {
 	return sha256.New
 }
 
+func (c cipherSuiteTLSEcdheEcdsaWithAes256CbcSha) isPSK() bool {
+	return false
+}
+
 func (c *cipherSuiteTLSEcdheEcdsaWithAes256CbcSha) init(masterSecret, clientRandom, serverRandom []byte, isClient bool) error {
 	const (
 		prfMacLen = 20
diff --git a/config.go b/config.go
index aabb7ec..7716221 100644
--- a/config.go
+++ b/config.go
@@ -11,10 +11,9 @@ import (
 // Config is used to configure a DTLS client or server.
 // After a Config is passed to a DTLS function it must not be modified.
 type Config struct {
-	// Certificates contains certificate chain to present to
-	// the other side of the connection. Server MUST set this,
-	// client SHOULD sets this so CertificateRequests
-	// can be handled
+	// Certificates contains certificate chain to present to the other side of the connection.
+	// Server MUST set this if PSK is non-nil
+	// client SHOULD sets this so CertificateRequests can be handled if PSK is non-nil
 	Certificate *x509.Certificate
 
 	// PrivateKey contains matching private key for the certificate
@@ -38,9 +37,18 @@ type Config struct {
 	// defaults to time.Second
 	FlightInterval time.Duration
 
+	// PSK sets the pre-shared key used by this DTLS connection
+	// If PSK is non-nil only PSK CipherSuites will be used
+	PSK             PSKCallback
+	PSKIdentityHint []byte
+
 	LoggerFactory logging.LoggerFactory
 }
 
+// PSKCallback is called once we have the remote's PSKIdentityHint.
+// If the remote provided none it will be nil
+type PSKCallback func([]byte) ([]byte, error)
+
 // ClientAuthType declares the policy the server will follow for
 // TLS Client Authentication.
 type ClientAuthType int
diff --git a/conn.go b/conn.go
index 3f85a96..e566827 100644
--- a/conn.go
+++ b/conn.go
@@ -5,7 +5,6 @@ import (
 	"crypto/ecdsa"
 	"crypto/rand"
 	"crypto/x509"
-	"errors"
 	"fmt"
 	"net"
 	"sync"
@@ -15,9 +14,12 @@ import (
 	"github.com/pion/logging"
 )
 
-const initialTickerInterval = time.Second
-const cookieLength = 20
-const defaultNamedCurve = namedCurveX25519
+const (
+	initialTickerInterval = time.Second
+	cookieLength          = 20
+	defaultNamedCurve     = namedCurveX25519
+	inboundBufferSize     = 8192
+)
 
 var invalidKeyingLabels = map[string]bool{
 	"client finished": true,
@@ -54,6 +56,9 @@ type Conn struct {
 	localKeypair, remoteKeypair *namedCurveKeypair
 	cookie                      []byte
 
+	localPSKCallback     PSKCallback
+	localPSKIdentityHint []byte
+
 	localCertificateVerify    []byte // cache CertificateVerify
 	localVerifyData           []byte // cached VerifyData
 	localKeySignature         []byte // cached keySignature
@@ -68,34 +73,24 @@ type Conn struct {
 }
 
 func createConn(nextConn net.Conn, flightHandler flightHandler, handshakeMessageHandler handshakeMessageHandler, config *Config, isClient bool) (*Conn, error) {
-	if config == nil {
-		return nil, errors.New("no config provided")
-	}
-
-	loggerFactory := config.LoggerFactory
-	if loggerFactory == nil {
-		loggerFactory = logging.NewDefaultLoggerFactory()
+	switch {
+	case config == nil:
+		return nil, errNoConfigProvided
+	case nextConn == nil:
+		return nil, errNilNextConn
+	case config.Certificate != nil && (config.PSK != nil || config.PSKIdentityHint != nil):
+		return nil, errPSKAndCertificate
 	}
 
 	if config.PrivateKey != nil {
 		if _, ok := config.PrivateKey.(*ecdsa.PrivateKey); !ok {
 			return nil, errInvalidPrivateKey
 		}
-	} else if nextConn == nil {
-		return nil, errNilNextConn
 	}
 
-	cipherSuites := []cipherSuite{}
-	if len(config.CipherSuites) != 0 {
-		for _, id := range config.CipherSuites {
-			c := cipherSuiteForID(id)
-			if c == nil {
-				return nil, fmt.Errorf("CipherSuite with id(%d) is not valid", id)
-			}
-			cipherSuites = append(cipherSuites, c)
-		}
-	} else {
-		cipherSuites = defaultCipherSuites()
+	cipherSuites, err := parseCipherSuites(config.CipherSuites, config.PSK == nil, config.PSK != nil)
+	if err != nil {
+		return nil, err
 	}
 
 	workerInterval := initialTickerInterval
@@ -103,6 +98,11 @@ func createConn(nextConn net.Conn, flightHandler flightHandler, handshakeMessage
 		workerInterval = config.FlightInterval
 	}
 
+	loggerFactory := config.LoggerFactory
+	if loggerFactory == nil {
+		loggerFactory = logging.NewDefaultLoggerFactory()
+	}
+
 	c := &Conn{
 		nextConn:                    nextConn,
 		currFlight:                  newFlight(isClient),
@@ -117,19 +117,21 @@ func createConn(nextConn net.Conn, flightHandler flightHandler, handshakeMessage
 		localCipherSuites:           cipherSuites,
 		namedCurve:                  defaultNamedCurve,
 
+		localPSKCallback:     config.PSK,
+		localPSKIdentityHint: config.PSKIdentityHint,
+
 		decrypted:          make(chan []byte),
 		workerTicker:       time.NewTicker(workerInterval),
 		handshakeCompleted: make(chan bool),
 		log:                loggerFactory.NewLogger("dtls"),
 	}
-	c.state.isClient = isClient
 
 	var zeroEpoch uint16
 	c.state.localEpoch.Store(zeroEpoch)
 	c.state.remoteEpoch.Store(zeroEpoch)
+	c.state.isClient = isClient
 
-	err := c.state.localRandom.populate()
-	if err != nil {
+	if err = c.state.localRandom.populate(); err != nil {
 		return nil, err
 	}
 	if !isClient {
@@ -143,27 +145,7 @@ func createConn(nextConn net.Conn, flightHandler flightHandler, handshakeMessage
 	c.startHandshakeOutbound()
 
 	// Handle inbound
-	go func() {
-		defer func() {
-			close(c.decrypted)
-		}()
-
-		b := make([]byte, 8192)
-		for {
-			i, err := c.nextConn.Read(b)
-			if err != nil {
-				c.stopWithError(err)
-				return
-			} else if c.getConnErr() != nil {
-				return
-			}
-
-			if err := c.handleIncoming(b[:i]); err != nil {
-				c.stopWithError(err)
-				return
-			}
-		}
-	}()
+	go c.inboundLoop()
 
 	<-c.handshakeCompleted
 	c.log.Trace("Handshake Completed")
@@ -186,9 +168,12 @@ func Client(conn net.Conn, config *Config) (*Conn, error) {
 
 // Server listens for incoming DTLS connections
 func Server(conn net.Conn, config *Config) (*Conn, error) {
-	if config == nil || config.Certificate == nil {
+	if config == nil {
+		return nil, errNoConfigProvided
+	} else if config.PSK == nil && config.Certificate == nil {
 		return nil, errServerMustHaveCertificate
 	}
+
 	return createConn(conn, serverFlightHandler, serverHandshakeHandler, config, false)
 }
 
@@ -320,19 +305,35 @@ func (c *Conn) internalSend(pkt *recordLayer, shouldEncrypt bool) {
 	}
 }
 
-func (c *Conn) handleIncoming(buf []byte) error {
-	pkts, err := unpackDatagram(buf)
-	if err != nil {
-		return err
-	}
+func (c *Conn) inboundLoop() {
+	defer func() {
+		close(c.decrypted)
+	}()
 
-	for _, p := range pkts {
-		err := c.handleIncomingPacket(p)
+	b := make([]byte, inboundBufferSize)
+	for {
+		i, err := c.nextConn.Read(b)
 		if err != nil {
-			return err
+			c.stopWithError(err)
+			return
+		} else if c.getConnErr() != nil {
+			return
+		}
+
+		pkts, err := unpackDatagram(b[:i])
+		if err != nil {
+			c.stopWithError(err)
+			return
+		}
+
+		for _, p := range pkts {
+			err := c.handleIncomingPacket(p)
+			if err != nil {
+				c.stopWithError(err)
+				return
+			}
 		}
 	}
-	return nil
 }
 
 func (c *Conn) handleIncomingPacket(buf []byte) error {
diff --git a/errors.go b/errors.go
index ace7a8a..5c0e473 100644
--- a/errors.go
+++ b/errors.go
@@ -44,4 +44,7 @@ var (
 	errServerMustHaveCertificate         = errors.New("dtls: Certificate is mandatory for server")
 	errUnableToMarshalFragmented         = errors.New("dtls: unable to marshal fragmented handshakes")
 	errVerifyDataMismatch                = errors.New("dtls: Expected and actual verify data does not match")
+	errNoConfigProvided                  = errors.New("dtls: No config provided")
+	errPSKAndCertificate                 = errors.New("dtls: Certificate and PSK or PSK Identity Hint provided")
+	errNoAvailableCipherSuites           = errors.New("dtls: Connection can not be created, no CipherSuites satisfy this Config")
 )


Test patch:
diff --git a/conn_test.go b/conn_test.go
index 2539aad..064af44 100644
--- a/conn_test.go
+++ b/conn_test.go
@@ -83,12 +83,12 @@ func pipeMemory() (*Conn, *Conn, error) {
 
 	// Setup client
 	go func() {
-		client, err := testClient(ca, &Config{SRTPProtectionProfiles: []SRTPProtectionProfile{SRTP_AES128_CM_HMAC_SHA1_80}})
+		client, err := testClient(ca, &Config{SRTPProtectionProfiles: []SRTPProtectionProfile{SRTP_AES128_CM_HMAC_SHA1_80}}, true)
 		c <- result{client, err}
 	}()
 
 	// Setup server
-	server, err := testServer(cb, &Config{SRTPProtectionProfiles: []SRTPProtectionProfile{SRTP_AES128_CM_HMAC_SHA1_80}})
+	server, err := testServer(cb, &Config{SRTPProtectionProfiles: []SRTPProtectionProfile{SRTP_AES128_CM_HMAC_SHA1_80}}, true)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -102,23 +102,27 @@ func pipeMemory() (*Conn, *Conn, error) {
 	return res.c, server, nil
 }
 
-func testClient(c net.Conn, cfg *Config) (*Conn, error) {
-	clientCert, clientKey, err := GenerateSelfSigned()
-	if err != nil {
-		return nil, err
+func testClient(c net.Conn, cfg *Config, generateCertificate bool) (*Conn, error) {
+	if generateCertificate {
+		clientCert, clientKey, err := GenerateSelfSigned()
+		if err != nil {
+			return nil, err
+		}
+		cfg.PrivateKey = clientKey
+		cfg.Certificate = clientCert
 	}
-	cfg.PrivateKey = clientKey
-	cfg.Certificate = clientCert
 	return Client(c, cfg)
 }
 
-func testServer(c net.Conn, cfg *Config) (*Conn, error) {
-	serverCert, serverKey, err := GenerateSelfSigned()
-	if err != nil {
-		return nil, err
+func testServer(c net.Conn, cfg *Config, generateCertificate bool) (*Conn, error) {
+	if generateCertificate {
+		serverCert, serverKey, err := GenerateSelfSigned()
+		if err != nil {
+			return nil, err
+		}
+		cfg.PrivateKey = serverKey
+		cfg.Certificate = serverCert
 	}
-	cfg.PrivateKey = serverKey
-	cfg.Certificate = serverCert
 	return Server(c, cfg)
 }
 
@@ -222,11 +226,11 @@ func TestSRTPConfiguration(t *testing.T) {
 		c := make(chan result)
 
 		go func() {
-			client, err := testClient(ca, &Config{SRTPProtectionProfiles: test.ClientSRTP})
+			client, err := testClient(ca, &Config{SRTPProtectionProfiles: test.ClientSRTP}, true)
 			c <- result{client, err}
 		}()
 
-		server, err := testServer(cb, &Config{SRTPProtectionProfiles: test.ServerSRTP})
+		server, err := testServer(cb, &Config{SRTPProtectionProfiles: test.ServerSRTP}, true)
 		if err != nil || test.WantServerError != nil {
 			if !(err != nil && test.WantServerError != nil && err.Error() == test.WantServerError.Error()) {
 				t.Errorf("TestSRTPConfiguration: Server Error Mismatch '%s': expected(%v) actual(%v)", test.Name, test.WantServerError, err)
@@ -263,12 +267,12 @@ func TestClientCertificate(t *testing.T) {
 
 	go func() {
 		conf := &Config{ClientAuth: RequireAnyClientCert}
-		client, err := testClient(ca, conf)
+		client, err := testClient(ca, conf, true)
 		c <- result{client, conf, err}
 	}()
 
 	serverCfg := &Config{ClientAuth: RequireAnyClientCert}
-	server, err := testServer(cb, serverCfg)
+	server, err := testServer(cb, serverCfg, true)
 	if err != nil {
 		t.Errorf("TestClientCertificate: Server failed(%v)", err)
 	}
@@ -342,11 +346,68 @@ func TestCipherSuiteConfiguration(t *testing.T) {
 		c := make(chan result)
 
 		go func() {
-			client, err := testClient(ca, &Config{CipherSuites: test.ClientCipherSuites})
+			client, err := testClient(ca, &Config{CipherSuites: test.ClientCipherSuites}, true)
+			c <- result{client, err}
+		}()
+
+		_, err := testServer(cb, &Config{CipherSuites: test.ServerCipherSuites}, true)
+		if err != nil || test.WantServerError != nil {
+			if !(err != nil && test.WantServerError != nil && err.Error() == test.WantServerError.Error()) {
+				t.Errorf("TestCipherSuiteConfiguration: Server Error Mismatch '%s': expected(%v) actual(%v)", test.Name, test.WantServerError, err)
+			}
+		}
+
+		res := <-c
+		if res.err != nil || test.WantClientError != nil {
+			if !(res.err != nil && test.WantClientError != nil && err.Error() == test.WantClientError.Error()) {
+				t.Errorf("TestSRTPConfiguration: Client Error Mismatch '%s': expected(%v) actual(%v)", test.Name, test.WantClientError, err)
+			}
+		}
+	}
+}
+
+func TestPSKConfiguration(t *testing.T) {
+	for _, test := range []struct {
+		Name                 string
+		ClientHasCertificate bool
+		ServerHasCertificate bool
+		ClientPSK            PSKCallback
+		ServerPSK            PSKCallback
+		WantClientError      error
+		WantServerError      error
+	}{
+		{
+			Name:                 "PSK specified",
+			ClientHasCertificate: false,
+			ServerHasCertificate: false,
+			ClientPSK:            func([]byte) ([]byte, error) { return []byte{0x00, 0x01, 0x02}, nil },
+			ServerPSK:            func([]byte) ([]byte, error) { return []byte{0x00, 0x01, 0x02}, nil },
+			WantClientError:      errNoAvailableCipherSuites, // TODO (should be nil when PSK CipherSuite is added)
+			WantServerError:      errNoAvailableCipherSuites, // TODO (should be nil when PSK CipherSuite is added)
+		},
+		{
+			Name:                 "PSK and certificate specified",
+			ClientHasCertificate: true,
+			ServerHasCertificate: true,
+			ClientPSK:            func([]byte) ([]byte, error) { return []byte{0x00, 0x01, 0x02}, nil },
+			ServerPSK:            func([]byte) ([]byte, error) { return []byte{0x00, 0x01, 0x02}, nil },
+			WantClientError:      errPSKAndCertificate,
+			WantServerError:      errPSKAndCertificate,
+		},
+	} {
+		ca, cb := net.Pipe()
+		type result struct {
+			c   *Conn
+			err error
+		}
+		c := make(chan result)
+
+		go func() {
+			client, err := testClient(ca, &Config{PSK: test.ClientPSK}, test.ClientHasCertificate)
 			c <- result{client, err}
 		}()
 
-		_, err := testServer(cb, &Config{CipherSuites: test.ServerCipherSuites})
+		_, err := testServer(cb, &Config{PSK: test.ServerPSK}, test.ServerHasCertificate)
 		if err != nil || test.WantServerError != nil {
 			if !(err != nil && test.WantServerError != nil && err.Error() == test.WantServerError.Error()) {
 				t.Errorf("TestCipherSuiteConfiguration: Server Error Mismatch '%s': expected(%v) actual(%v)", test.Name, test.WantServerError, err)