Repository: pion/dtls
Pull request #76 (2019-06-30 10:04:44)

Start TLS_PSK_WITH_AES_128_CCM8
@daenney This works 100% !`OpenSSL <-> Pion` and `Pion <-> Pion`

I am going on vacation on Tuesday, mind helping me catch corner cases etc...? Would love to get this merged and tagged :) 

Sean-Der: To use 
```
diff --git a/examples/listen/main.go b/examples/listen/main.go
index 37b8261..6310ce2 100644
--- a/examples/listen/main.go
+++ b/examples/listen/main.go
@@ -16,16 +16,13 @@ func main() {
        // Prepare the IP to connect to
        addr := &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 4444}

-       // Generate a certificate and private key to secure the connection
-       certificate, privateKey, genErr := dtls.GenerateSelfSigned()
-       util.Check(genErr)
-
-       //
-       // Everything below is the pion-DTLS API! Thanks for using it ❤️ .
-       //
-
        // Prepare the configuration of the DTLS connection
-       config := &dtls.Config{Certificate: certificate, PrivateKey: privateKey}
+       config := &dtls.Config{
+               PSK: func(hint []byte) ([]byte, error) {
+                       return []byte{0xAB, 0xC1, 0x23}, nil
+               },
+               CipherSuites: []dtls.CipherSuiteID{dtls.TLS_PSK_WITH_AES_128_CCM8},
+       }

        // Connect to a DTLS server
        listener, err := dtls.Listen("udp", addr, config)
```

Against
```
openssl s_client -dtls1_2 -connect 127.0.0.1:4444 -psk abc123 -cipher PSK-AES128-CCM8
```
codecov-io: # [Codecov](https://codecov.io/gh/pion/dtls/pull/76?src=pr&el=h1) Report
> Merging [#76](https://codecov.io/gh/pion/dtls/pull/76?src=pr&el=desc) into [master](https://codecov.io/gh/pion/dtls/commit/b1c2af308c308e012d5101185ca57d1858d9b814?src=pr&el=desc) will **decrease** coverage by `2.25%`.
> The diff coverage is `44.69%`.

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

```diff
@@            Coverage Diff             @@
##           master      #76      +/-   ##
==========================================
- Coverage   69.67%   67.42%   -2.26%     
==========================================
  Files          48       50       +2     
  Lines        2681     2775      +94     
==========================================
+ Hits         1868     1871       +3     
- Misses        579      669      +90     
- Partials      234      235       +1
```


| [Impacted Files](https://codecov.io/gh/pion/dtls/pull/76?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [conn.go](https://codecov.io/gh/pion/dtls/pull/76/diff?src=pr&el=tree#diff-Y29ubi5nbw==) | `76.04% <ø> (ø)` | :arrow_up: |
| [prf.go](https://codecov.io/gh/pion/dtls/pull/76/diff?src=pr&el=tree#diff-cHJmLmdv) | `57.52% <0%> (-3.23%)` | :arrow_down: |
| [cipher\_suite.go](https://codecov.io/gh/pion/dtls/pull/76/diff?src=pr&el=tree#diff-Y2lwaGVyX3N1aXRlLmdv) | `94.02% <0%> (-2.9%)` | :arrow_down: |
| [crypto\_ccm.go](https://codecov.io/gh/pion/dtls/pull/76/diff?src=pr&el=tree#diff-Y3J5cHRvX2NjbS5nbw==) | `0% <0%> (ø)` | |
| [crypto\_gcm.go](https://codecov.io/gh/pion/dtls/pull/76/diff?src=pr&el=tree#diff-Y3J5cHRvX2djbS5nbw==) | `62.5% <100%> (-7.5%)` | :arrow_down: |
| [client\_handlers.go](https://codecov.io/gh/pion/dtls/pull/76/diff?src=pr&el=tree#diff-Y2xpZW50X2hhbmRsZXJzLmdv) | `79.78% <100%> (-0.08%)` | :arrow_down: |
| [crypto.go](https://codecov.io/gh/pion/dtls/pull/76/diff?src=pr&el=tree#diff-Y3J5cHRvLmdv) | `61.11% <100%> (+5.55%)` | :arrow_up: |
| [cipher\_suite\_tls\_psk\_with\_aes\_128\_ccm8.go](https://codecov.io/gh/pion/dtls/pull/76/diff?src=pr&el=tree#diff-Y2lwaGVyX3N1aXRlX3Rsc19wc2tfd2l0aF9hZXNfMTI4X2NjbTguZ28=) | `6.25% <6.25%> (ø)` | |
| [server\_handlers.go](https://codecov.io/gh/pion/dtls/pull/76/diff?src=pr&el=tree#diff-c2VydmVyX2hhbmRsZXJzLmdv) | `82.17% <83.11%> (-1.01%)` | :arrow_down: |
| ... and [1 more](https://codecov.io/gh/pion/dtls/pull/76/diff?src=pr&el=tree-more) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/pion/dtls/pull/76?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/76?src=pr&el=footer). Last update [b1c2af3...eaeafda](https://codecov.io/gh/pion/dtls/pull/76?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

Sean-Der: @daenney @igolaizola 

I am happy with this w/e I can get a review I would really appreciate it! 
daenney: This is awesome, thank you so much! I'm currently also on vacay but I have access to the IoT equipment at home, so I'll try to test this today.
daenney: I'm happy to say I've just managed to ping my IKEA Tradfri gateway with this! So this has been validated in protoduction :smile:.

If anyone's looking on and would like to try it, here's the sample code:

```go
package main

import (
	"fmt"
	"net"

	"github.com/dustin/go-coap"
	"github.com/pion/dtls"
)

func main() {
	addr := &net.UDPAddr{IP: net.ParseIP("gateway-ip"), Port: 5684}
	config := &dtls.Config{
		PSK: func(_ []byte) ([]byte, error) {
			return []byte("psk-on-bottom-of-gateway"), nil
		},
		PSKIdentityHint: []byte("Client_identity"), # For Tradfri Gateway the IdentityHint MUST be Client_identity
		CipherSuites:    []dtls.CipherSuiteID{dtls.TLS_PSK_WITH_AES_128_CCM8},
	}

	dtlsConn, err := dtls.Dial("udp", addr, config)
	if err != nil {
		panic(err)
	}
	defer dtlsConn.Close()

	req := coap.Message{
		Type:      coap.Confirmable,
		Code:      coap.GET,
		MessageID: 1,
		Token:     []byte{},
		Payload:   []byte{},
	} # This is a CoAP Ping, empty confirmable message

	data, err := req.MarshalBinary()
	if err != nil {
		panic(err)
	}

	_, err = dtlsConn.Write(data)
	if err != nil {
		panic(err)
	}

	resp := make([]byte, 2048)
	dtlsConn.Read(resp)
	msg, err := coap.ParseMessage(resp)
	if err != nil {
		panic(err)
	}
	fmt.Println(msg.MessageID)
	fmt.Println(msg.Type)
	fmt.Println(msg.Code)
	fmt.Println(msg.Token)
	fmt.Println(string(msg.Payload))

}
```
It'll output:

```
1
Acknowledgement
NotFound
[]
Not Found
```
Sean-Der: @daenney that is awesome! So glad it works :)

Thanks for doing half the work, really glad we did this. I am hoping we can get `pion/dtls` good enough so it can get wide adoption/get the community around one Go implementation.

Patch:
diff --git a/README.md b/README.md
index 57e80ce..cc1b558 100644
--- a/README.md
+++ b/README.md
@@ -28,10 +28,11 @@ We would love contributes that fall under the 'Planned Features' and fixing any
 
 #### Current features
 * DTLS 1.2 Client/Server
-* Forward secrecy using ECDHE; with curve25519 and nistp256 (non-PFS will not be supported)
-* AES_128_GCM, AES_256_CBC
+* Key Exchange via ECDHE(curve25519 and nistp256) and PSK
+* AES_128_GCM, AES_256_CBC, AES_128_CCM8
 * Packet loss and re-ordering is handled during handshaking
 * Key export (RFC5705)
+* Serialization and Resumption of sessions
 
 #### Planned Features
 * Extended master secret support (RFC7627)
@@ -42,7 +43,9 @@ We would love contributes that fall under the 'Planned Features' and fixing any
 * Renegotiation
 * Compression
 
-### Pion DTLS
+### Using
+
+#### Pion DTLS
 For a DTLS 1.2 Server that listens on 127.0.0.1:4444
 ```sh
 go run examples/listen/main.go
@@ -53,7 +56,7 @@ For a DTLS 1.2 Client that connects to 127.0.0.1:4444
 go run examples/dial/main.go
 ```
 
-### OpenSSL
+#### OpenSSL
 Pion DTLS can connect to itself and OpenSSL.
 ```
   // Generate a certificate
@@ -68,6 +71,28 @@ Pion DTLS can connect to itself and OpenSSL.
   openssl s_client -dtls1_2 -connect 127.0.0.1:4444 -debug -cert cert.pem -key key.pem
 ```
 
+### Using with PSK
+Pion DTLS also comes with examples that do key exchange via PSK
+
+
+#### Pion DTLS
+```sh
+go run examples/listen-psk/main.go
+```
+
+```sh
+go run examples/dial-psk/main.go
+```
+
+#### OpenSSL
+```
+  // Use with examples/dial-psk/main.go
+  openssl s_server -dtls1_2 -accept 4444 -nocert -psk abc123 -cipher PSK-AES128-CCM8
+
+  // Use with examples/listen-psk/main.go
+  openssl s_client -dtls1_2 -connect 127.0.0.1:4444 -psk abc123 -cipher PSK-AES128-CCM8
+```
+
 ### Contributing
 Check out the **[contributing wiki](https://github.com/pion/webrtc/wiki/Contributing)** to join the group of amazing people making this project possible:
 
diff --git a/cipher_suite.go b/cipher_suite.go
index 162d51c..ef05125 100644
--- a/cipher_suite.go
+++ b/cipher_suite.go
@@ -18,6 +18,8 @@ const (
 	// AES-256-CBC-SHA
 	TLS_ECDHE_ECDSA_WITH_AES_256_CBC_SHA CipherSuiteID = 0xc00a
 	TLS_ECDHE_RSA_WITH_AES_256_CBC_SHA   CipherSuiteID = 0x0035
+
+	TLS_PSK_WITH_AES_128_CCM8 CipherSuiteID = 0xc0a8
 )
 
 type cipherSuite interface {
@@ -47,8 +49,9 @@ func cipherSuiteForID(id CipherSuiteID) cipherSuite {
 		return &cipherSuiteTLSEcdheEcdsaWithAes256CbcSha{}
 	case cipherSuiteTLSEcdheRsaWithAes256CbcSha{}.ID():
 		return &cipherSuiteTLSEcdheRsaWithAes256CbcSha{}
+	case cipherSuiteTLSPskWithAes128Ccm8{}.ID():
+		return &cipherSuiteTLSPskWithAes128Ccm8{}
 	}
-
 	return nil
 }
 
diff --git a/cipher_suite_tls_psk_with_aes_128_ccm8.go b/cipher_suite_tls_psk_with_aes_128_ccm8.go
new file mode 100644
index 0000000..aa5d883
--- /dev/null
+++ b/cipher_suite_tls_psk_with_aes_128_ccm8.go
@@ -0,0 +1,67 @@
+package dtls
+
+import (
+	"crypto/sha256"
+	"errors"
+	"hash"
+)
+
+type cipherSuiteTLSPskWithAes128Ccm8 struct {
+	ccm *cryptoCCM
+}
+
+func (c cipherSuiteTLSPskWithAes128Ccm8) certificateType() clientCertificateType {
+	return clientCertificateType(0)
+}
+
+func (c cipherSuiteTLSPskWithAes128Ccm8) ID() CipherSuiteID {
+	return TLS_PSK_WITH_AES_128_CCM8
+}
+
+func (c cipherSuiteTLSPskWithAes128Ccm8) String() string {
+	return "TLS_PSK_WITH_AES_128_CCM8"
+}
+
+func (c cipherSuiteTLSPskWithAes128Ccm8) hashFunc() func() hash.Hash {
+	return sha256.New
+}
+
+func (c cipherSuiteTLSPskWithAes128Ccm8) isPSK() bool {
+	return true
+}
+
+func (c *cipherSuiteTLSPskWithAes128Ccm8) init(masterSecret, clientRandom, serverRandom []byte, isClient bool) error {
+	const (
+		prfMacLen = 0
+		prfKeyLen = 16
+		prfIvLen  = 4
+	)
+
+	keys, err := prfEncryptionKeys(masterSecret, clientRandom, serverRandom, prfMacLen, prfKeyLen, prfIvLen, c.hashFunc())
+	if err != nil {
+		return err
+	}
+
+	if isClient {
+		c.ccm, err = newCryptoCCM(keys.clientWriteKey, keys.clientWriteIV, keys.serverWriteKey, keys.serverWriteIV)
+	} else {
+		c.ccm, err = newCryptoCCM(keys.serverWriteKey, keys.serverWriteIV, keys.clientWriteKey, keys.clientWriteIV)
+	}
+	return err
+}
+
+func (c *cipherSuiteTLSPskWithAes128Ccm8) encrypt(pkt *recordLayer, raw []byte) ([]byte, error) {
+	if c.ccm == nil {
+		return nil, errors.New("CipherSuite has not been initalized, unable to encrypt")
+	}
+
+	return c.ccm.encrypt(pkt, raw)
+}
+
+func (c *cipherSuiteTLSPskWithAes128Ccm8) decrypt(raw []byte) ([]byte, error) {
+	if c.ccm == nil {
+		return nil, errors.New("CipherSuite has not been initalized, unable to decrypt ")
+	}
+
+	return c.ccm.decrypt(raw)
+}
diff --git a/client_handlers.go b/client_handlers.go
index e23e9f4..ca0ec7e 100644
--- a/client_handlers.go
+++ b/client_handlers.go
@@ -6,6 +6,53 @@ import (
 )
 
 func clientHandshakeHandler(c *Conn) error {
+	initalizeCipherSuite := func(h *handshakeMessageServerKeyExchange) error {
+		clientRandom, err := c.state.localRandom.Marshal()
+		if err != nil {
+			return err
+		}
+		serverRandom, err := c.state.remoteRandom.Marshal()
+		if err != nil {
+			return err
+		}
+
+		var preMasterSecret []byte
+		if c.localPSKCallback != nil {
+			var psk []byte
+			if psk, err = c.localPSKCallback(h.identityHint); err != nil {
+				return err
+			}
+
+			preMasterSecret = prfPSKPreMasterSecret(psk)
+
+		} else {
+			if c.localKeypair, err = generateKeypair(h.namedCurve); err != nil {
+				return err
+			}
+
+			if preMasterSecret, err = prfPreMasterSecret(h.publicKey, c.localKeypair.privateKey, c.localKeypair.curve); err != nil {
+				return err
+			}
+		}
+
+		c.state.masterSecret, err = prfMasterSecret(preMasterSecret, clientRandom, serverRandom, c.state.cipherSuite.hashFunc())
+		if err != nil {
+			return err
+		}
+
+		if err := c.state.cipherSuite.init(c.state.masterSecret, clientRandom, serverRandom /* isClient */, true); err != nil {
+			return err
+		}
+
+		if c.localPSKCallback == nil {
+			expectedHash := valueKeySignature(clientRandom, serverRandom, h.publicKey, h.namedCurve, h.hashAlgorithm)
+			if err := verifyKeySignature(expectedHash, h.signature, h.hashAlgorithm, c.state.remoteCertificate); err != nil {
+				return err
+			}
+		}
+		return nil
+	}
+
 	handleSingleHandshake := func(buf []byte) error {
 		rawHandshake := &handshake{}
 		if err := rawHandshake.Unmarshal(buf); err != nil {
@@ -42,41 +89,9 @@ func clientHandshakeHandler(c *Conn) error {
 			c.state.remoteCertificate = h.certificate
 
 		case *handshakeMessageServerKeyExchange:
-			c.remoteKeypair = &namedCurveKeypair{h.namedCurve, h.publicKey, nil}
-
-			clientRandom, err := c.state.localRandom.Marshal()
-			if err != nil {
+			if err := initalizeCipherSuite(h); err != nil {
 				return err
 			}
-			serverRandom, err := c.state.remoteRandom.Marshal()
-			if err != nil {
-				return err
-			}
-
-			c.localKeypair, err = generateKeypair(h.namedCurve)
-			if err != nil {
-				return err
-			}
-
-			preMasterSecret, err := prfPreMasterSecret(c.remoteKeypair.publicKey, c.localKeypair.privateKey, c.localKeypair.curve)
-			if err != nil {
-				return err
-			}
-
-			c.state.masterSecret, err = prfMasterSecret(preMasterSecret, clientRandom, serverRandom, c.state.cipherSuite.hashFunc())
-			if err != nil {
-				return err
-			}
-
-			if err := c.state.cipherSuite.init(c.state.masterSecret, clientRandom, serverRandom /* isClient */, true); err != nil {
-				return err
-			}
-
-			expectedHash := valueKeySignature(clientRandom, serverRandom, h.publicKey, h.namedCurve, h.hashAlgorithm)
-			if err := verifyKeySignature(expectedHash, h.signature, h.hashAlgorithm, c.state.remoteCertificate); err != nil {
-				return err
-			}
-
 		case *handshakeMessageCertificateRequest:
 			c.remoteRequestedCertificate = true
 		case *handshakeMessageServerHelloDone:
@@ -150,7 +165,11 @@ func clientHandshakeHandler(c *Conn) error {
 		expectedSeqnum := expectedMessages[0].messageSequence
 		for i, msg := range expectedMessages {
 			switch {
-			// handshakeMessageCertificateRequest can be nil, just make sure we have no gaps
+			// handshakeTypeCertificate and handshakeTypeServerKeyExchange can be nil
+			// when doing PSK
+			case c.localPSKCallback != nil && (i == 1 || i == 2) && msg == nil:
+				continue
+			// handshakeMessageCertificateRequest can be nil
 			case i == 3 && msg == nil:
 				continue
 			case msg == nil:
@@ -168,6 +187,14 @@ func clientHandshakeHandler(c *Conn) error {
 				}
 			}
 		}
+
+		// handshakeMessageServerKeyExchange is optional for PSK
+		if c.state.masterSecret == nil {
+			if err := initalizeCipherSuite(&handshakeMessageServerKeyExchange{}); err != nil {
+				return err
+			}
+		}
+
 		c.state.localSequenceNumber++
 		c.log.Tracef("[handshake] Flight 3 changed to %s", flight5.String())
 		if err := c.currFlight.set(flight5); err != nil {
@@ -203,12 +230,6 @@ func clientFlightHandler(c *Conn) (bool, error) {
 		fallthrough
 	case flight3:
 		extensions := []extension{
-			&extensionSupportedEllipticCurves{
-				ellipticCurves: []namedCurve{namedCurveX25519, namedCurveP256},
-			},
-			&extensionSupportedPointFormats{
-				pointFormats: []ellipticCurvePointFormat{ellipticCurvePointFormatUncompressed},
-			},
 			&extensionSupportedSignatureAlgorithms{
 				signatureHashAlgorithms: []signatureHashAlgorithm{
 					{HashAlgorithmSHA256, signatureAlgorithmECDSA},
@@ -220,6 +241,17 @@ func clientFlightHandler(c *Conn) (bool, error) {
 				},
 			},
 		}
+		if c.localPSKCallback == nil {
+			extensions = append(extensions, []extension{
+				&extensionSupportedEllipticCurves{
+					ellipticCurves: []namedCurve{namedCurveX25519, namedCurveP256},
+				},
+				&extensionSupportedPointFormats{
+					pointFormats: []ellipticCurvePointFormat{ellipticCurvePointFormatUncompressed},
+				},
+			}...)
+		}
+
 		if len(c.localSRTPProtectionProfiles) > 0 {
 			extensions = append(extensions, &extensionUseSRTP{
 				protectionProfiles: c.localSRTPProtectionProfiles,
@@ -271,6 +303,13 @@ func clientFlightHandler(c *Conn) (bool, error) {
 			sequenceNumber++
 		}
 
+		clientKeyExchange := &handshakeMessageClientKeyExchange{}
+		if c.localPSKCallback == nil {
+			clientKeyExchange.publicKey = c.localKeypair.publicKey
+		} else {
+			clientKeyExchange.identityHint = c.localPSKIdentityHint
+		}
+
 		c.internalSend(&recordLayer{
 			recordLayerHeader: recordLayerHeader{
 				sequenceNumber:  sequenceNumber,
@@ -281,10 +320,10 @@ func clientFlightHandler(c *Conn) (bool, error) {
 				handshakeHeader: handshakeHeader{
 					messageSequence: uint16(sequenceNumber),
 				},
-				handshakeMessage: &handshakeMessageClientKeyExchange{
-					publicKey: c.localKeypair.publicKey,
-				}},
+				handshakeMessage: clientKeyExchange,
+			},
 		}, false)
+
 		sequenceNumber++
 
 		if c.remoteRequestedCertificate {
diff --git a/conn.go b/conn.go
index e566827..dfd196c 100644
--- a/conn.go
+++ b/conn.go
@@ -49,12 +49,12 @@ type Conn struct {
 
 	clientAuth ClientAuthType // If we are a client should we request a client certificate
 
-	currFlight                  *flight
-	namedCurve                  namedCurve
-	localCertificate            *x509.Certificate
-	localPrivateKey             crypto.PrivateKey
-	localKeypair, remoteKeypair *namedCurveKeypair
-	cookie                      []byte
+	currFlight       *flight
+	namedCurve       namedCurve
+	localCertificate *x509.Certificate
+	localPrivateKey  crypto.PrivateKey
+	localKeypair     *namedCurveKeypair
+	cookie           []byte
 
 	localPSKCallback     PSKCallback
 	localPSKIdentityHint []byte
@@ -78,8 +78,10 @@ func createConn(nextConn net.Conn, flightHandler flightHandler, handshakeMessage
 		return nil, errNoConfigProvided
 	case nextConn == nil:
 		return nil, errNilNextConn
-	case config.Certificate != nil && (config.PSK != nil || config.PSKIdentityHint != nil):
+	case config.Certificate != nil && config.PSK != nil:
 		return nil, errPSKAndCertificate
+	case config.PSKIdentityHint != nil && config.PSK == nil:
+		return nil, errIdentityNoPSK
 	}
 
 	if config.PrivateKey != nil {
@@ -163,14 +165,22 @@ func Dial(network string, raddr *net.UDPAddr, config *Config) (*Conn, error) {
 
 // Client establishes a DTLS connection over an existing conn
 func Client(conn net.Conn, config *Config) (*Conn, error) {
+	switch {
+	case config == nil:
+		return nil, errNoConfigProvided
+	case config.PSK != nil && config.PSKIdentityHint == nil:
+		return nil, errPSKAndIdentityMustBeSetForClient
+	}
+
 	return createConn(conn, clientFlightHandler, clientHandshakeHandler, config, true)
 }
 
 // Server listens for incoming DTLS connections
 func Server(conn net.Conn, config *Config) (*Conn, error) {
-	if config == nil {
+	switch {
+	case config == nil:
 		return nil, errNoConfigProvided
-	} else if config.PSK == nil && config.Certificate == nil {
+	case config.PSK == nil && config.Certificate == nil:
 		return nil, errServerMustHaveCertificate
 	}
 
diff --git a/crypto.go b/crypto.go
index 451d314..ccad8b7 100644
--- a/crypto.go
+++ b/crypto.go
@@ -120,3 +120,17 @@ func verifyCertificateVerify(handshakeBodies []byte, hashAlgorithm HashAlgorithm
 
 	return errKeySignatureVerifyUnimplemented
 }
+
+func generateAEADAdditionalData(h *recordLayerHeader, payloadLen int) []byte {
+	var additionalData [13]byte
+	// SequenceNumber MUST be set first
+	// we only want uint48, clobbering an extra 2 (using uint64, Golang doesn't have uint48)
+	binary.BigEndian.PutUint64(additionalData[:], h.sequenceNumber)
+	binary.BigEndian.PutUint16(additionalData[:], h.epoch)
+	additionalData[8] = byte(h.contentType)
+	additionalData[9] = h.protocolVersion.major
+	additionalData[10] = h.protocolVersion.minor
+	binary.BigEndian.PutUint16(additionalData[len(additionalData)-2:], uint16(payloadLen))
+
+	return additionalData[:]
+}
diff --git a/crypto_ccm.go b/crypto_ccm.go
new file mode 100644
index 0000000..3facb24
--- /dev/null
+++ b/crypto_ccm.go
@@ -0,0 +1,92 @@
+package dtls
+
+import (
+	"crypto/aes"
+	"crypto/rand"
+	"encoding/binary"
+	"fmt"
+
+	"github.com/pion/dtls/internal/crypto/ccm"
+)
+
+const (
+	cryptoCCMTagLength   = 8
+	cryptoCCMNonceLength = 12
+)
+
+// State needed to handle encrypted input/output
+type cryptoCCM struct {
+	localCCM, remoteCCM         ccm.CCM
+	localWriteIV, remoteWriteIV []byte
+}
+
+func newCryptoCCM(localKey, localWriteIV, remoteKey, remoteWriteIV []byte) (*cryptoCCM, error) {
+	localBlock, err := aes.NewCipher(localKey)
+	if err != nil {
+		return nil, err
+	}
+	localCCM, err := ccm.NewCCM(localBlock, cryptoCCMTagLength, cryptoCCMNonceLength)
+	if err != nil {
+		return nil, err
+	}
+
+	remoteBlock, err := aes.NewCipher(remoteKey)
+	if err != nil {
+		return nil, err
+	}
+	remoteCCM, err := ccm.NewCCM(remoteBlock, cryptoCCMTagLength, cryptoCCMNonceLength)
+	if err != nil {
+		return nil, err
+	}
+
+	return &cryptoCCM{
+		localCCM:      localCCM,
+		localWriteIV:  localWriteIV,
+		remoteCCM:     remoteCCM,
+		remoteWriteIV: remoteWriteIV,
+	}, nil
+}
+
+func (c *cryptoCCM) encrypt(pkt *recordLayer, raw []byte) ([]byte, error) {
+	payload := raw[recordLayerHeaderSize:]
+	raw = raw[:recordLayerHeaderSize]
+
+	nonce := append(append([]byte{}, c.localWriteIV[:4]...), make([]byte, 8)...)
+	if _, err := rand.Read(nonce[4:]); err != nil {
+		return nil, err
+	}
+
+	additionalData := generateAEADAdditionalData(&pkt.recordLayerHeader, len(payload))
+	encryptedPayload := c.localCCM.Seal(nil, nonce, payload, additionalData)
+
+	encryptedPayload = append(nonce[4:], encryptedPayload...)
+	raw = append(raw, encryptedPayload...)
+
+	// Update recordLayer size to include explicit nonce
+	binary.BigEndian.PutUint16(raw[recordLayerHeaderSize-2:], uint16(len(raw)-recordLayerHeaderSize))
+	return raw, nil
+}
+
+func (c *cryptoCCM) decrypt(in []byte) ([]byte, error) {
+	var h recordLayerHeader
+	err := h.Unmarshal(in)
+	switch {
+	case err != nil:
+		return nil, err
+	case h.contentType == contentTypeChangeCipherSpec:
+		// Nothing to encrypt with ChangeCipherSpec
+		return in, nil
+	case len(in) <= (8 + recordLayerHeaderSize):
+		return nil, errNotEnoughRoomForNonce
+	}
+
+	nonce := append(append([]byte{}, c.remoteWriteIV[:4]...), in[recordLayerHeaderSize:recordLayerHeaderSize+8]...)
+	out := in[recordLayerHeaderSize+8:]
+
+	additionalData := generateAEADAdditionalData(&h, len(out)-cryptoCCMTagLength)
+	out, err = c.remoteCCM.Open(out[:0], nonce, out, additionalData)
+	if err != nil {
+		return nil, fmt.Errorf("decryptPacket: %v", err)
+	}
+	return append(in[:recordLayerHeaderSize], out...), nil
+}
diff --git a/crypto_gcm.go b/crypto_gcm.go
index 58034c0..02c03f3 100644
--- a/crypto_gcm.go
+++ b/crypto_gcm.go
@@ -52,17 +52,8 @@ func (c *cryptoGCM) encrypt(pkt *recordLayer, raw []byte) ([]byte, error) {
 		return nil, err
 	}
 
-	var additionalData [13]byte
-	// SequenceNumber MUST be set first
-	// we only want uint48, clobbering an extra 2 (using uint64, Golang doesn't have uint48)
-	binary.BigEndian.PutUint64(additionalData[:], pkt.recordLayerHeader.sequenceNumber)
-	binary.BigEndian.PutUint16(additionalData[:], pkt.recordLayerHeader.epoch)
-	additionalData[8] = byte(pkt.content.contentType())
-	additionalData[9] = pkt.recordLayerHeader.protocolVersion.major
-	additionalData[10] = pkt.recordLayerHeader.protocolVersion.minor
-	binary.BigEndian.PutUint16(additionalData[len(additionalData)-2:], uint16(len(payload)))
-	encryptedPayload := c.localGCM.Seal(nil, nonce, payload, additionalData[:])
-
+	additionalData := generateAEADAdditionalData(&pkt.recordLayerHeader, len(payload))
+	encryptedPayload := c.localGCM.Seal(nil, nonce, payload, additionalData)
 	encryptedPayload = append(nonce[4:], encryptedPayload...)
 	raw = append(raw, encryptedPayload...)
 
@@ -88,16 +79,8 @@ func (c *cryptoGCM) decrypt(in []byte) ([]byte, error) {
 	nonce := append(append([]byte{}, c.remoteWriteIV[:4]...), in[recordLayerHeaderSize:recordLayerHeaderSize+8]...)
 	out := in[recordLayerHeaderSize+8:]
 
-	var additionalData [13]byte
-	// SequenceNumber MUST be set first
-	// we only want uint48, clobbering an extra 2 (using uint64, Golang doesn't have uint48)
-	binary.BigEndian.PutUint64(additionalData[:], h.sequenceNumber)
-	binary.BigEndian.PutUint16(additionalData[:], h.epoch)
-	additionalData[8] = byte(h.contentType)
-	additionalData[9] = h.protocolVersion.major
-	additionalData[10] = h.protocolVersion.minor
-	binary.BigEndian.PutUint16(additionalData[len(additionalData)-2:], uint16(len(out)-cryptoGCMTagLength))
-	out, err = c.remoteGCM.Open(out[:0], nonce, out, additionalData[:])
+	additionalData := generateAEADAdditionalData(&h, len(out)-cryptoGCMTagLength)
+	out, err = c.remoteGCM.Open(out[:0], nonce, out, additionalData)
 	if err != nil {
 		return nil, fmt.Errorf("decryptPacket: %v", err)
 	}
diff --git a/errors.go b/errors.go
index 5c0e473..d1c434d 100644
--- a/errors.go
+++ b/errors.go
@@ -13,7 +13,7 @@ var (
 	errCertificateVerifyNoCertificate    = errors.New("dtls: client sent certificate verify but we have no certificate to verify")
 	errCipherSuiteNoIntersection         = errors.New("dtls: Client+Server do not support any shared cipher suites")
 	errCipherSuiteUnset                  = errors.New("dtls: server hello can not be created without a cipher suite")
-	errCompressionmethodUnset            = errors.New("dtls: server hello can not be created without a compression method")
+	errCompressionMethodUnset            = errors.New("dtls: server hello can not be created without a compression method")
 	errContextUnsupported                = errors.New("dtls: context is not supported for ExportKeyingMaterial")
 	errCookieMismatch                    = errors.New("dtls: Client+Server cookie does not match")
 	errCookieTooLong                     = errors.New("dtls: cookie must not be longer then 255 bytes")
@@ -45,6 +45,9 @@ var (
 	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")
+	errPSKAndCertificate                 = errors.New("dtls: Certificate and PSK provided")
+	errPSKAndIdentityMustBeSetForClient  = errors.New("dtls: PSK and PSK Identity Hint must both be set for client")
+	errIdentityNoPSK                     = errors.New("dtls: Identity Hint provided but PSK is nil")
 	errNoAvailableCipherSuites           = errors.New("dtls: Connection can not be created, no CipherSuites satisfy this Config")
+	errInvalidClientKeyExchange          = errors.New("dtls: Unable to determine if ClientKeyExchange is a public key or PSK Identity")
 )
diff --git a/examples/dial-psk/main.go b/examples/dial-psk/main.go
new file mode 100644
index 0000000..6644f0b
--- /dev/null
+++ b/examples/dial-psk/main.go
@@ -0,0 +1,40 @@
+package main
+
+import (
+	"fmt"
+	"net"
+
+	"github.com/pion/dtls"
+	"github.com/pion/dtls/examples/util"
+)
+
+func main() {
+	// Prepare the IP to connect to
+	addr := &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 4444}
+
+	//
+	// Everything below is the pion-DTLS API! Thanks for using it ❤️.
+	//
+
+	// Prepare the configuration of the DTLS connection
+	config := &dtls.Config{
+		PSK: func(hint []byte) ([]byte, error) {
+			fmt.Printf("Server's hint: %s \n", hint)
+			return []byte{0xAB, 0xC1, 0x23}, nil
+		},
+		PSKIdentityHint: []byte("Pion DTLS Server"),
+		CipherSuites:    []dtls.CipherSuiteID{dtls.TLS_PSK_WITH_AES_128_CCM8},
+	}
+
+	// Connect to a DTLS server
+	dtlsConn, err := dtls.Dial("udp", addr, config)
+	util.Check(err)
+	defer func() {
+		util.Check(dtlsConn.Close())
+	}()
+
+	fmt.Println("Connected; type 'exit' to shutdown gracefully")
+
+	// Simulate a chat session
+	util.Chat(dtlsConn)
+}
diff --git a/examples/listen-psk/main.go b/examples/listen-psk/main.go
new file mode 100644
index 0000000..059951f
--- /dev/null
+++ b/examples/listen-psk/main.go
@@ -0,0 +1,55 @@
+package main
+
+import (
+	"fmt"
+	"net"
+
+	"github.com/pion/dtls"
+	"github.com/pion/dtls/examples/util"
+)
+
+func main() {
+	// Prepare the IP to connect to
+	addr := &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 4444}
+
+	//
+	// Everything below is the pion-DTLS API! Thanks for using it ❤️.
+	//
+
+	// Prepare the configuration of the DTLS connection
+	config := &dtls.Config{
+		PSK: func(hint []byte) ([]byte, error) {
+			fmt.Printf("Client's hint: %s \n", hint)
+			return []byte{0xAB, 0xC1, 0x23}, nil
+		},
+		PSKIdentityHint: []byte("Pion DTLS Client"),
+		CipherSuites:    []dtls.CipherSuiteID{dtls.TLS_PSK_WITH_AES_128_CCM8},
+	}
+
+	// Connect to a DTLS server
+	listener, err := dtls.Listen("udp", addr, config)
+	util.Check(err)
+	defer func() {
+		util.Check(listener.Close())
+	}()
+
+	fmt.Println("Listening")
+
+	// Simulate a chat session
+	hub := util.NewHub()
+
+	go func() {
+		for {
+			// Wait for a connection.
+			conn, err := listener.Accept()
+			util.Check(err)
+			// defer conn.Close() // TODO: graceful shutdown
+
+			// Register the connection with the chat hub
+			hub.Register(conn)
+		}
+	}()
+
+	// Start chatting
+	hub.Chat()
+}
diff --git a/examples/listen/main.go b/examples/listen/main.go
index 37b8261..1869659 100644
--- a/examples/listen/main.go
+++ b/examples/listen/main.go
@@ -1,12 +1,8 @@
 package main
 
 import (
-	"bufio"
 	"fmt"
 	"net"
-	"os"
-	"strings"
-	"sync"
 
 	"github.com/pion/dtls"
 	"github.com/pion/dtls/examples/util"
@@ -37,7 +33,7 @@ func main() {
 	fmt.Println("Listening")
 
 	// Simulate a chat session
-	hub := newHub()
+	hub := util.NewHub()
 
 	go func() {
 		for {
@@ -47,79 +43,10 @@ func main() {
 			// defer conn.Close() // TODO: graceful shutdown
 
 			// Register the connection with the chat hub
-			hub.register(conn)
+			hub.Register(conn)
 		}
 	}()
 
 	// Start chatting
-	hub.chat()
-}
-
-const bufSize = 8192
-
-// hub is a helper to handle one to many chat
-type hub struct {
-	conns map[string]net.Conn
-	lock  sync.RWMutex
-}
-
-func newHub() *hub {
-	return &hub{conns: make(map[string]net.Conn)}
-}
-
-func (h *hub) register(conn net.Conn) {
-	fmt.Printf("Connected to %s\n", conn.RemoteAddr())
-	h.lock.Lock()
-	defer h.lock.Unlock()
-
-	h.conns[conn.RemoteAddr().String()] = conn
-
-	go h.readLoop(conn)
-}
-
-func (h *hub) readLoop(conn net.Conn) {
-	b := make([]byte, bufSize)
-	for {
-		n, err := conn.Read(b)
-		if err != nil {
-			h.unregister(conn)
-			return
-		}
-		fmt.Printf("Got message: %s\n", string(b[:n]))
-	}
-}
-
-func (h *hub) unregister(conn net.Conn) {
-	h.lock.Lock()
-	defer h.lock.Unlock()
-	delete(h.conns, conn.RemoteAddr().String())
-	err := conn.Close()
-	if err != nil {
-		fmt.Println("Failed to disconnect", conn.RemoteAddr(), err)
-	} else {
-		fmt.Println("Disconnected ", conn.RemoteAddr())
-	}
-}
-
-func (h *hub) broadcast(msg []byte) {
-	h.lock.RLock()
-	defer h.lock.RUnlock()
-	for _, conn := range h.conns {
-		_, err := conn.Write(msg)
-		if err != nil {
-			fmt.Printf("Failed to write message to %s: %v\n", conn.RemoteAddr(), err)
-		}
-	}
-}
-
-func (h *hub) chat() {
-	reader := bufio.NewReader(os.Stdin)
-	for {
-		msg, err := reader.ReadString('\n')
-		util.Check(err)
-		if strings.TrimSpace(msg) == "exit" {
-			return
-		}
-		h.broadcast([]byte(msg))
-	}
+	hub.Chat()
 }
diff --git a/examples/util/hub.go b/examples/util/hub.go
new file mode 100644
index 0000000..ad8e597
--- /dev/null
+++ b/examples/util/hub.go
@@ -0,0 +1,80 @@
+package util
+
+import (
+	"bufio"
+	"fmt"
+	"net"
+	"os"
+	"strings"
+	"sync"
+)
+
+// Hub is a helper to handle one to many chat
+type Hub struct {
+	conns map[string]net.Conn
+	lock  sync.RWMutex
+}
+
+// NewHub builds a new hub
+func NewHub() *Hub {
+	return &Hub{conns: make(map[string]net.Conn)}
+}
+
+// Register adds a new conn to the Hub
+func (h *Hub) Register(conn net.Conn) {
+	fmt.Printf("Connected to %s\n", conn.RemoteAddr())
+	h.lock.Lock()
+	defer h.lock.Unlock()
+
+	h.conns[conn.RemoteAddr().String()] = conn
+
+	go h.readLoop(conn)
+}
+
+func (h *Hub) readLoop(conn net.Conn) {
+	b := make([]byte, bufSize)
+	for {
+		n, err := conn.Read(b)
+		if err != nil {
+			h.unregister(conn)
+			return
+		}
+		fmt.Printf("Got message: %s\n", string(b[:n]))
+	}
+}
+
+func (h *Hub) unregister(conn net.Conn) {
+	h.lock.Lock()
+	defer h.lock.Unlock()
+	delete(h.conns, conn.RemoteAddr().String())
+	err := conn.Close()
+	if err != nil {
+		fmt.Println("Failed to disconnect", conn.RemoteAddr(), err)
+	} else {
+		fmt.Println("Disconnected ", conn.RemoteAddr())
+	}
+}
+
+func (h *Hub) broadcast(msg []byte) {
+	h.lock.RLock()
+	defer h.lock.RUnlock()
+	for _, conn := range h.conns {
+		_, err := conn.Write(msg)
+		if err != nil {
+			fmt.Printf("Failed to write message to %s: %v\n", conn.RemoteAddr(), err)
+		}
+	}
+}
+
+// Chat starts the stdin readloop to dispatch messages to the hub
+func (h *Hub) Chat() {
+	reader := bufio.NewReader(os.Stdin)
+	for {
+		msg, err := reader.ReadString('\n')
+		Check(err)
+		if strings.TrimSpace(msg) == "exit" {
+			return
+		}
+		h.broadcast([]byte(msg))
+	}
+}
diff --git a/handshake_message_client_key_exchange.go b/handshake_message_client_key_exchange.go
index e2bea32..b2ab209 100644
--- a/handshake_message_client_key_exchange.go
+++ b/handshake_message_client_key_exchange.go
@@ -1,7 +1,12 @@
 package dtls
 
+import (
+	"encoding/binary"
+)
+
 type handshakeMessageClientKeyExchange struct {
-	publicKey []byte
+	identityHint []byte
+	publicKey    []byte
 }
 
 func (h handshakeMessageClientKeyExchange) handshakeType() handshakeType {
@@ -9,17 +14,33 @@ func (h handshakeMessageClientKeyExchange) handshakeType() handshakeType {
 }
 
 func (h *handshakeMessageClientKeyExchange) Marshal() ([]byte, error) {
-	return append([]byte{byte(len(h.publicKey))}, h.publicKey...), nil
+	switch {
+	case (h.identityHint != nil && h.publicKey != nil) || (h.identityHint == nil && h.publicKey == nil):
+		return nil, errInvalidClientKeyExchange
+	case h.publicKey != nil:
+		return append([]byte{byte(len(h.publicKey))}, h.publicKey...), nil
+	default:
+		out := append([]byte{0x00, 0x00}, h.identityHint...)
+		binary.BigEndian.PutUint16(out, uint16(len(out)-2))
+		return out, nil
+	}
 }
 
 func (h *handshakeMessageClientKeyExchange) Unmarshal(data []byte) error {
-	if len(data) < 1 {
+	if len(data) < 2 {
 		return errBufferTooSmall
 	}
-	publicKeyLength := int(data[0])
-	if len(data) <= publicKeyLength {
+
+	// If parsed as PSK return early and only populate PSK Identity Hint
+	if pskLength := binary.BigEndian.Uint16(data); len(data) == int(pskLength+2) {
+		h.identityHint = append([]byte{}, data[2:]...)
+		return nil
+	}
+
+	if publicKeyLength := int(data[0]); len(data) != publicKeyLength+1 {
 		return errBufferTooSmall
 	}
+
 	h.publicKey = append([]byte{}, data[1:]...)
 	return nil
 }
diff --git a/handshake_message_server_hello.go b/handshake_message_server_hello.go
index 4837abc..b5b4e00 100644
--- a/handshake_message_server_hello.go
+++ b/handshake_message_server_hello.go
@@ -30,7 +30,7 @@ func (h *handshakeMessageServerHello) Marshal() ([]byte, error) {
 	if h.cipherSuite == nil {
 		return nil, errCipherSuiteUnset
 	} else if h.compressionMethod == nil {
-		return nil, errCompressionmethodUnset
+		return nil, errCompressionMethodUnset
 	}
 
 	out := make([]byte, handshakeMessageServerHelloVariableWidthStart)
diff --git a/handshake_message_server_key_exchange.go b/handshake_message_server_key_exchange.go
index 93c2966..fd5427c 100644
--- a/handshake_message_server_key_exchange.go
+++ b/handshake_message_server_key_exchange.go
@@ -4,8 +4,10 @@ import (
 	"encoding/binary"
 )
 
-// Structure only supports ECDH
+// Structure supports ECDH and PSK
 type handshakeMessageServerKeyExchange struct {
+	identityHint []byte
+
 	ellipticCurveType  ellipticCurveType
 	namedCurve         namedCurve
 	publicKey          []byte
@@ -19,6 +21,12 @@ func (h handshakeMessageServerKeyExchange) handshakeType() handshakeType {
 }
 
 func (h *handshakeMessageServerKeyExchange) Marshal() ([]byte, error) {
+	if h.identityHint != nil {
+		out := append([]byte{0x00, 0x00}, h.identityHint...)
+		binary.BigEndian.PutUint16(out, uint16(len(out)-2))
+		return out, nil
+	}
+
 	out := []byte{byte(h.ellipticCurveType), 0x00, 0x00}
 	binary.BigEndian.PutUint16(out[1:], uint16(h.namedCurve))
 
@@ -34,9 +42,16 @@ func (h *handshakeMessageServerKeyExchange) Marshal() ([]byte, error) {
 }
 
 func (h *handshakeMessageServerKeyExchange) Unmarshal(data []byte) error {
-	if len(data) < 1 {
+	if len(data) < 2 {
 		return errBufferTooSmall
 	}
+
+	// If parsed as PSK return early and only populate PSK Identity Hint
+	if pskLength := binary.BigEndian.Uint16(data); len(data) == int(pskLength+2) {
+		h.identityHint = append([]byte{}, data[2:]...)
+		return nil
+	}
+
 	if _, ok := ellipticCurveTypes[ellipticCurveType(data[0])]; ok {
 		h.ellipticCurveType = ellipticCurveType(data[0])
 	} else {
diff --git a/prf.go b/prf.go
index 7244238..e085c50 100644
--- a/prf.go
+++ b/prf.go
@@ -50,6 +50,21 @@ func (e *encryptionKeys) String() string {
 		e.serverWriteIV)
 }
 
+// The premaster secret is formed as follows: if the PSK is N octets
+// long, concatenate a uint16 with the value N, N zero octets, a second
+// uint16 with the value N, and the PSK itself.
+//
+// https://tools.ietf.org/html/rfc4279#section-2
+func prfPSKPreMasterSecret(psk []byte) []byte {
+	pskLen := uint16(len(psk))
+
+	out := append(make([]byte, 2+pskLen+2), psk...)
+	binary.BigEndian.PutUint16(out, pskLen)
+	binary.BigEndian.PutUint16(out[2+pskLen:], pskLen)
+
+	return out
+}
+
 func prfPreMasterSecret(publicKey, privateKey []byte, curve namedCurve) ([]byte, error) {
 	switch curve {
 	case namedCurveX25519:
diff --git a/server_handlers.go b/server_handlers.go
index c2cb712..5d942e1 100644
--- a/server_handlers.go
+++ b/server_handlers.go
@@ -82,8 +82,6 @@ func serverHandshakeHandler(c *Conn) error {
 			c.state.remoteCertificate = h.certificate
 
 		case *handshakeMessageClientKeyExchange:
-			c.remoteKeypair = &namedCurveKeypair{c.namedCurve, h.publicKey, nil}
-
 			serverRandom, err := c.state.localRandom.Marshal()
 			if err != nil {
 				return err
@@ -93,9 +91,19 @@ func serverHandshakeHandler(c *Conn) error {
 				return err
 			}
 
-			preMasterSecret, err := prfPreMasterSecret(c.remoteKeypair.publicKey, c.localKeypair.privateKey, c.localKeypair.curve)
-			if err != nil {
-				return err
+			var preMasterSecret []byte
+			if c.localPSKCallback != nil {
+				var psk []byte
+				if psk, err = c.localPSKCallback(h.identityHint); err != nil {
+					return err
+				}
+
+				preMasterSecret = prfPSKPreMasterSecret(psk)
+			} else {
+				preMasterSecret, err = prfPreMasterSecret(h.publicKey, c.localKeypair.privateKey, c.localKeypair.curve)
+				if err != nil {
+					return err
+				}
 			}
 
 			c.state.masterSecret, err = prfMasterSecret(preMasterSecret, clientRandom, serverRandom, c.state.cipherSuite.hashFunc())
@@ -106,7 +114,6 @@ func serverHandshakeHandler(c *Conn) error {
 			if err := c.state.cipherSuite.init(c.state.masterSecret, clientRandom, serverRandom /* isClient */, false); err != nil {
 				return err
 			}
-
 		case *handshakeMessageFinished:
 			plainText := c.handshakeCache.pullAndMerge(
 				handshakeCachePullRule{handshakeTypeClientHello, true},
@@ -199,9 +206,14 @@ func serverHandshakeHandler(c *Conn) error {
 			return errClientCertificateNotVerified
 		}
 
-		if c.clientAuth > NoClientCert {
+		switch {
+		case c.localPSKIdentityHint != nil:
+			c.state.localSequenceNumber = 4
+		case c.localPSKCallback != nil:
+			c.state.localSequenceNumber = 3
+		case c.clientAuth > NoClientCert:
 			c.state.localSequenceNumber = 6
-		} else {
+		default:
 			c.state.localSequenceNumber = 5
 		}
 		c.setLocalEpoch(1)
@@ -239,19 +251,22 @@ func serverFlightHandler(c *Conn) (bool, error) {
 		}, false)
 
 	case flight4:
-		extensions := []extension{
-			&extensionSupportedEllipticCurves{
-				ellipticCurves: []namedCurve{namedCurveX25519, namedCurveP256},
-			},
-			&extensionSupportedPointFormats{
-				pointFormats: []ellipticCurvePointFormat{ellipticCurvePointFormatUncompressed},
-			},
-		}
+		extensions := []extension{}
 		if c.state.srtpProtectionProfile != 0 {
 			extensions = append(extensions, &extensionUseSRTP{
 				protectionProfiles: []SRTPProtectionProfile{c.state.srtpProtectionProfile},
 			})
 		}
+		if c.localPSKCallback == nil {
+			extensions = append(extensions, []extension{
+				&extensionSupportedEllipticCurves{
+					ellipticCurves: []namedCurve{namedCurveX25519, namedCurveP256},
+				},
+				&extensionSupportedPointFormats{
+					pointFormats: []ellipticCurvePointFormat{ellipticCurvePointFormatUncompressed},
+				},
+			}...)
+		}
 
 		sequenceNumber := c.state.localSequenceNumber
 		c.internalSend(&recordLayer{
@@ -274,61 +289,40 @@ func serverFlightHandler(c *Conn) (bool, error) {
 		}, false)
 		sequenceNumber++
 
-		c.internalSend(&recordLayer{
-			recordLayerHeader: recordLayerHeader{
-				sequenceNumber:  sequenceNumber,
-				protocolVersion: protocolVersion1_2,
-			},
-			content: &handshake{
-				// sequenceNumber and messageSequence line up, may need to be re-evaluated
-				handshakeHeader: handshakeHeader{
-					messageSequence: uint16(sequenceNumber),
+		if c.localPSKCallback == nil {
+			c.internalSend(&recordLayer{
+				recordLayerHeader: recordLayerHeader{
+					sequenceNumber:  sequenceNumber,
+					protocolVersion: protocolVersion1_2,
 				},
-				handshakeMessage: &handshakeMessageCertificate{
-					certificate: c.localCertificate,
-				}},
-		}, false)
-		sequenceNumber++
+				content: &handshake{
+					// sequenceNumber and messageSequence line up, may need to be re-evaluated
+					handshakeHeader: handshakeHeader{
+						messageSequence: uint16(sequenceNumber),
+					},
+					handshakeMessage: &handshakeMessageCertificate{
+						certificate: c.localCertificate,
+					}},
+			}, false)
+			sequenceNumber++
 
-		if len(c.localKeySignature) == 0 {
-			serverRandom, err := c.state.localRandom.Marshal()
-			if err != nil {
-				return false, err
-			}
-			clientRandom, err := c.state.remoteRandom.Marshal()
-			if err != nil {
-				return false, err
-			}
+			if len(c.localKeySignature) == 0 {
+				serverRandom, err := c.state.localRandom.Marshal()
+				if err != nil {
+					return false, err
+				}
+				clientRandom, err := c.state.remoteRandom.Marshal()
+				if err != nil {
+					return false, err
+				}
 
-			signature, err := generateKeySignature(clientRandom, serverRandom, c.localKeypair.publicKey, c.namedCurve, c.localPrivateKey, HashAlgorithmSHA256)
-			if err != nil {
-				return false, err
+				signature, err := generateKeySignature(clientRandom, serverRandom, c.localKeypair.publicKey, c.namedCurve, c.localPrivateKey, HashAlgorithmSHA256)
+				if err != nil {
+					return false, err
+				}
+				c.localKeySignature = signature
 			}
-			c.localKeySignature = signature
-		}
 
-		c.internalSend(&recordLayer{
-			recordLayerHeader: recordLayerHeader{
-				sequenceNumber:  sequenceNumber,
-				protocolVersion: protocolVersion1_2,
-			},
-			content: &handshake{
-				// sequenceNumber and messageSequence line up, may need to be re-evaluated
-				handshakeHeader: handshakeHeader{
-					messageSequence: uint16(sequenceNumber),
-				},
-				handshakeMessage: &handshakeMessageServerKeyExchange{
-					ellipticCurveType:  ellipticCurveTypeNamedCurve,
-					namedCurve:         c.namedCurve,
-					publicKey:          c.localKeypair.publicKey,
-					hashAlgorithm:      HashAlgorithmSHA256,
-					signatureAlgorithm: signatureAlgorithmECDSA,
-					signature:          c.localKeySignature,
-				}},
-		}, false)
-		sequenceNumber++
-
-		if c.clientAuth > NoClientCert {
 			c.internalSend(&recordLayer{
 				recordLayerHeader: recordLayerHeader{
 					sequenceNumber:  sequenceNumber,
@@ -339,18 +333,63 @@ func serverFlightHandler(c *Conn) (bool, error) {
 					handshakeHeader: handshakeHeader{
 						messageSequence: uint16(sequenceNumber),
 					},
-					handshakeMessage: &handshakeMessageCertificateRequest{
-						certificateTypes: []clientCertificateType{clientCertificateTypeRSASign, clientCertificateTypeECDSASign},
-						signatureHashAlgorithms: []signatureHashAlgorithm{
-							{HashAlgorithmSHA256, signatureAlgorithmRSA},
-							{HashAlgorithmSHA384, signatureAlgorithmRSA},
-							{HashAlgorithmSHA512, signatureAlgorithmRSA},
-							{HashAlgorithmSHA256, signatureAlgorithmECDSA},
-							{HashAlgorithmSHA384, signatureAlgorithmECDSA},
-							{HashAlgorithmSHA512, signatureAlgorithmECDSA},
+					handshakeMessage: &handshakeMessageServerKeyExchange{
+						ellipticCurveType:  ellipticCurveTypeNamedCurve,
+						namedCurve:         c.namedCurve,
+						publicKey:          c.localKeypair.publicKey,
+						hashAlgorithm:      HashAlgorithmSHA256,
+						signatureAlgorithm: signatureAlgorithmECDSA,
+						signature:          c.localKeySignature,
+					}},
+			}, false)
+			sequenceNumber++
+
+			if c.clientAuth > NoClientCert {
+				c.internalSend(&recordLayer{
+					recordLayerHeader: recordLayerHeader{
+						sequenceNumber:  sequenceNumber,
+						protocolVersion: protocolVersion1_2,
+					},
+					content: &handshake{
+						// sequenceNumber and messageSequence line up, may need to be re-evaluated
+						handshakeHeader: handshakeHeader{
+							messageSequence: uint16(sequenceNumber),
+						},
+						handshakeMessage: &handshakeMessageCertificateRequest{
+							certificateTypes: []clientCertificateType{clientCertificateTypeRSASign, clientCertificateTypeECDSASign},
+							signatureHashAlgorithms: []signatureHashAlgorithm{
+								{HashAlgorithmSHA256, signatureAlgorithmRSA},
+								{HashAlgorithmSHA384, signatureAlgorithmRSA},
+								{HashAlgorithmSHA512, signatureAlgorithmRSA},
+								{HashAlgorithmSHA256, signatureAlgorithmECDSA},
+								{HashAlgorithmSHA384, signatureAlgorithmECDSA},
+								{HashAlgorithmSHA512, signatureAlgorithmECDSA},
+							},
 						},
 					},
+				}, false)
+				sequenceNumber++
+			}
+		} else if c.localPSKIdentityHint != nil {
+			/* To help the client in selecting which identity to use, the server
+			*  can provide a "PSK identity hint" in the ServerKeyExchange message.
+			*  If no hint is provided, the ServerKeyExchange message is omitted.
+			*
+			*  https://tools.ietf.org/html/rfc4279#section-2
+			 */
+			c.internalSend(&recordLayer{
+				recordLayerHeader: recordLayerHeader{
+					sequenceNumber:  sequenceNumber,
+					protocolVersion: protocolVersion1_2,
 				},
+				content: &handshake{
+					// sequenceNumber and messageSequence line up, may need to be re-evaluated
+					handshakeHeader: handshakeHeader{
+						messageSequence: uint16(sequenceNumber),
+					},
+					handshakeMessage: &handshakeMessageServerKeyExchange{
+						identityHint: c.localPSKIdentityHint,
+					}},
 			}, false)
 			sequenceNumber++
 		}


Test patch:
diff --git a/conn_test.go b/conn_test.go
index 064af44..e8ea9f1 100644
--- a/conn_test.go
+++ b/conn_test.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"errors"
 	"fmt"
+	"io"
 	"net"
 	"testing"
 	"time"
@@ -176,6 +177,108 @@ func TestExportKeyingMaterial(t *testing.T) {
 	}
 }
 
+func TestPSK(t *testing.T) {
+	// Limit runtime in case of deadlocks
+	lim := test.TimeOut(time.Second * 20)
+	defer lim.Stop()
+
+	for _, test := range []struct {
+		Name           string
+		ServerIdentity []byte
+	}{
+		{
+			Name:           "Server identity specified",
+			ServerIdentity: []byte("Test Identity"),
+		},
+		{
+			Name:           "Server identity nil",
+			ServerIdentity: nil,
+		},
+	} {
+
+		clientIdentity := []byte("Client Identity")
+		clientErr := make(chan error, 1)
+
+		ca, cb := net.Pipe()
+		go func() {
+			conf := &Config{
+				PSK: func(hint []byte) ([]byte, error) {
+					if !bytes.Equal(test.ServerIdentity, hint) { // nolint
+						return nil, fmt.Errorf("TestPSK: Client got invalid identity expected(% 02x) actual(% 02x)", test.ServerIdentity, hint) // nolint
+					}
+
+					return []byte{0xAB, 0xC1, 0x23}, nil
+				},
+				PSKIdentityHint: clientIdentity,
+				CipherSuites:    []CipherSuiteID{TLS_PSK_WITH_AES_128_CCM8},
+			}
+
+			_, err := testClient(ca, conf, false)
+			clientErr <- err
+		}()
+
+		config := &Config{
+			PSK: func(hint []byte) ([]byte, error) {
+				if !bytes.Equal(clientIdentity, hint) {
+					return nil, fmt.Errorf("TestPSK: Server got invalid identity expected(% 02x) actual(% 02x)", clientIdentity, hint)
+				}
+				return []byte{0xAB, 0xC1, 0x23}, nil
+			},
+			PSKIdentityHint: test.ServerIdentity,
+			CipherSuites:    []CipherSuiteID{TLS_PSK_WITH_AES_128_CCM8},
+		}
+
+		if _, err := testServer(cb, config, false); err != nil {
+			t.Fatalf("TestPSK: Server failed(%v)", err)
+		}
+
+		if err := <-clientErr; err != nil {
+			t.Fatal(err)
+		}
+	}
+}
+
+func TestPSKHintFail(t *testing.T) {
+	pskRejected := errors.New("PSK Rejected")
+
+	// Limit runtime in case of deadlocks
+	lim := test.TimeOut(time.Second * 20)
+	defer lim.Stop()
+
+	clientErr := make(chan error, 1)
+
+	ca, cb := net.Pipe()
+	go func() {
+		conf := &Config{
+			PSK: func(hint []byte) ([]byte, error) {
+				return nil, pskRejected
+			},
+			PSKIdentityHint: []byte{},
+			CipherSuites:    []CipherSuiteID{TLS_PSK_WITH_AES_128_CCM8},
+		}
+
+		_, err := testClient(ca, conf, false)
+		clientErr <- err
+	}()
+
+	config := &Config{
+		PSK: func(hint []byte) ([]byte, error) {
+			return nil, pskRejected
+		},
+		PSKIdentityHint: []byte{},
+		CipherSuites:    []CipherSuiteID{TLS_PSK_WITH_AES_128_CCM8},
+	}
+
+	if _, err := testServer(cb, config, false); err != io.EOF && err != io.ErrClosedPipe {
+		t.Fatalf("TestPSK: Server error exp(%v) failed(%v)", io.ErrClosedPipe, err)
+	}
+
+	if err := <-clientErr; err != pskRejected {
+		t.Fatalf("TestPSK: Client error exp(%v) failed(%v)", pskRejected, err)
+	}
+
+}
+
 func TestSRTPConfiguration(t *testing.T) {
 	for _, test := range []struct {
 		Name            string
@@ -206,7 +309,7 @@ func TestSRTPConfiguration(t *testing.T) {
 			ClientSRTP:      []SRTPProtectionProfile{SRTP_AES128_CM_HMAC_SHA1_80},
 			ServerSRTP:      nil,
 			ExpectedProfile: 0,
-			WantClientError: fmt.Errorf("Client requested SRTP but we have no matching profiles"),
+			WantClientError: io.EOF,
 			WantServerError: fmt.Errorf("Client requested SRTP but we have no matching profiles"),
 		},
 		{
@@ -239,8 +342,8 @@ func TestSRTPConfiguration(t *testing.T) {
 
 		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)
+			if !(res.err != nil && test.WantClientError != nil && res.err.Error() == test.WantClientError.Error()) {
+				t.Errorf("TestSRTPConfiguration: Client Error Mismatch '%s': expected(%v) actual(%v)", test.Name, test.WantClientError, res.err)
 			}
 		}
 
@@ -334,7 +437,7 @@ func TestCipherSuiteConfiguration(t *testing.T) {
 			Name:               "CipherSuites mismatch",
 			ClientCipherSuites: []CipherSuiteID{TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256},
 			ServerCipherSuites: []CipherSuiteID{TLS_ECDHE_ECDSA_WITH_AES_256_CBC_SHA},
-			WantClientError:    errCipherSuiteNoIntersection,
+			WantClientError:    io.EOF,
 			WantServerError:    errCipherSuiteNoIntersection,
 		},
 	} {
@@ -359,8 +462,8 @@ func TestCipherSuiteConfiguration(t *testing.T) {
 
 		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)
+			if !(res.err != nil && test.WantClientError != nil && res.err.Error() == test.WantClientError.Error()) {
+				t.Errorf("TestSRTPConfiguration: Client Error Mismatch '%s': expected(%v) actual(%v)", test.Name, test.WantClientError, res.err)
 			}
 		}
 	}
@@ -373,6 +476,8 @@ func TestPSKConfiguration(t *testing.T) {
 		ServerHasCertificate bool
 		ClientPSK            PSKCallback
 		ServerPSK            PSKCallback
+		ClientPSKIdentity    []byte
+		ServerPSKIdentity    []byte
 		WantClientError      error
 		WantServerError      error
 	}{
@@ -382,8 +487,10 @@ func TestPSKConfiguration(t *testing.T) {
 			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)
+			ClientPSKIdentity:    []byte{0x00},
+			ServerPSKIdentity:    []byte{0x00},
+			WantClientError:      errNoAvailableCipherSuites,
+			WantServerError:      errNoAvailableCipherSuites,
 		},
 		{
 			Name:                 "PSK and certificate specified",
@@ -391,9 +498,33 @@ func TestPSKConfiguration(t *testing.T) {
 			ServerHasCertificate: true,
 			ClientPSK:            func([]byte) ([]byte, error) { return []byte{0x00, 0x01, 0x02}, nil },
 			ServerPSK:            func([]byte) ([]byte, error) { return []byte{0x00, 0x01, 0x02}, nil },
+			ClientPSKIdentity:    []byte{0x00},
+			ServerPSKIdentity:    []byte{0x00},
 			WantClientError:      errPSKAndCertificate,
 			WantServerError:      errPSKAndCertificate,
 		},
+		{
+			Name:                 "PSK and no identity 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 },
+			ClientPSKIdentity:    nil,
+			ServerPSKIdentity:    nil,
+			WantClientError:      errPSKAndIdentityMustBeSetForClient,
+			WantServerError:      errNoAvailableCipherSuites,
+		},
+		{
+			Name:                 "No PSK and identity specified",
+			ClientHasCertificate: false,
+			ServerHasCertificate: false,
+			ClientPSK:            nil,
+			ServerPSK:            nil,
+			ClientPSKIdentity:    []byte{0x00},
+			ServerPSKIdentity:    []byte{0x00},
+			WantClientError:      errIdentityNoPSK,
+			WantServerError:      errServerMustHaveCertificate,
+		},
 	} {
 		ca, cb := net.Pipe()
 		type result struct {
@@ -403,21 +534,21 @@ func TestPSKConfiguration(t *testing.T) {
 		c := make(chan result)
 
 		go func() {
-			client, err := testClient(ca, &Config{PSK: test.ClientPSK}, test.ClientHasCertificate)
+			client, err := testClient(ca, &Config{PSK: test.ClientPSK, PSKIdentityHint: test.ClientPSKIdentity}, test.ClientHasCertificate)
 			c <- result{client, err}
 		}()
 
-		_, err := testServer(cb, &Config{PSK: test.ServerPSK}, test.ServerHasCertificate)
+		_, err := testServer(cb, &Config{PSK: test.ServerPSK, PSKIdentityHint: test.ServerPSKIdentity}, 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)
+				t.Fatalf("TestPSKConfiguration: 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)
+			if !(res.err != nil && test.WantClientError != nil && res.err.Error() == test.WantClientError.Error()) {
+				t.Fatalf("TestPSKConfiguration: Client Error Mismatch '%s': expected(%v) actual(%v)", test.Name, test.WantClientError, res.err)
 			}
 		}
 	}
diff --git a/e2e/e2e_test.go b/e2e/e2e_test.go
index 0792f3a..7f825cf 100644
--- a/e2e/e2e_test.go
+++ b/e2e/e2e_test.go
@@ -4,20 +4,36 @@ import (
 	"errors"
 	"io"
 	"net"
-	"os"
-	"runtime"
-	"runtime/pprof"
 	"sync/atomic"
 	"testing"
 	"time"
 
 	"github.com/pion/dtls"
+	"github.com/pion/transport/test"
 )
 
 const testMessage = "Hello World"
 const testTimeLimit = 5 * time.Second
 const messageRetry = 200 * time.Millisecond
 
+func randomPort(t testing.TB) int {
+	t.Helper()
+	conn, err := net.ListenPacket("udp4", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("failed to pickPort: %v", err)
+	}
+	defer func() {
+		_ = conn.Close()
+	}()
+	switch addr := conn.LocalAddr().(type) {
+	case *net.UDPAddr:
+		return addr.Port
+	default:
+		t.Fatalf("unknown addr type %T", addr)
+		return 0
+	}
+}
+
 func simpleReadWrite(errChan chan error, outChan chan string, conn io.ReadWriteCloser, listener io.Closer, messageRecvCount *uint64) {
 	go func() {
 		buffer := make([]byte, 8192)
@@ -55,22 +71,83 @@ func simpleReadWrite(errChan chan error, outChan chan string, conn io.ReadWriteC
 	maybePushError(conn.Close())
 }
 
-func pickPort(t testing.TB) int {
-	t.Helper()
-	conn, err := net.ListenPacket("udp4", "127.0.0.1:0")
-	if err != nil {
-		t.Fatalf("failed to pickPort: %v", err)
-	}
-	defer func() {
-		_ = conn.Close()
+func assertE2ECommunication(clientConfig, serverConfig *dtls.Config, t *testing.T) {
+	errChan := make(chan error)
+	clientChan := make(chan string)
+	serverChan := make(chan string)
+	var messageRecvCount uint64 // Counter to make sure both sides got a message
+
+	serverPort := randomPort(t)
+	serverReady := make(chan struct{})
+
+	// DTLS Client
+	go func() {
+		select {
+		case <-serverReady:
+			// OK
+		case <-time.After(time.Second):
+			errChan <- errors.New("waiting on serverReady err: timeout")
+		}
+		conn, err := dtls.Dial("udp",
+			&net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: serverPort},
+			clientConfig,
+		)
+		if err != nil {
+			errChan <- err
+			return
+		}
+		simpleReadWrite(errChan, clientChan, conn, nil, &messageRecvCount)
+	}()
+
+	// DTLS Server
+	go func() {
+		listener, err := dtls.Listen("udp",
+			&net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: serverPort},
+			serverConfig,
+		)
+		if err != nil {
+			errChan <- err
+			return
+		}
+		serverReady <- struct{}{}
+		conn, err := listener.Accept()
+		if err != nil {
+			errChan <- err
+			return
+		}
+
+		simpleReadWrite(errChan, serverChan, conn, listener, &messageRecvCount)
+	}()
+
+	func() {
+		seenClient, seenServer := false, false
+		for {
+			select {
+			case err := <-errChan:
+				t.Fatal(err)
+			case <-time.After(testTimeLimit):
+				t.Fatalf("Test timeout, seenClient %t seenServer %t", seenClient, seenServer)
+			case clientMsg := <-clientChan:
+				if clientMsg != testMessage {
+					t.Fatalf("clientMsg does not equal test message: %s %s", clientMsg, testMessage)
+				}
+
+				seenClient = true
+				if seenClient && seenServer {
+					return
+				}
+			case serverMsg := <-serverChan:
+				if serverMsg != testMessage {
+					t.Fatalf("serverMsg does not equal test message: %s %s", serverMsg, testMessage)
+				}
+
+				seenServer = true
+				if seenClient && seenServer {
+					return
+				}
+			}
+		}
 	}()
-	switch addr := conn.LocalAddr().(type) {
-	case *net.UDPAddr:
-		return addr.Port
-	default:
-		t.Fatalf("unknown addr type %T", addr)
-		return 0
-	}
 }
 
 /*
@@ -80,103 +157,41 @@ func pickPort(t testing.TB) int {
 	- Assert that no Goroutines are leaked
 */
 func TestPionE2ESimple(t *testing.T) {
-	expectedGoRoutineCount := runtime.NumGoroutine()
+	lim := test.TimeOut(time.Second * 30)
+	defer lim.Stop()
+
+	report := test.CheckRoutines(t)
+	defer report()
+
 	for _, cipherSuite := range []dtls.CipherSuiteID{
 		dtls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256,
 		dtls.TLS_ECDHE_ECDSA_WITH_AES_256_CBC_SHA,
 	} {
-		errChan := make(chan error)
-		clientChan := make(chan string)
-		serverChan := make(chan string)
-		var messageRecvCount uint64 // Counter to make sure both sides got a message
-
-		serverPort := pickPort(t)
-		serverReady := make(chan struct{})
-
 		cert, key, err := dtls.GenerateSelfSigned()
 		if err != nil {
 			t.Fatal(err)
 		}
 
-		// DTLS Client
-		go func() {
-			select {
-			case <-serverReady:
-				// OK
-			case <-time.After(time.Second):
-				errChan <- errors.New("waiting on serverReady err: timeout")
-			}
-			conn, err := dtls.Dial("udp",
-				&net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: serverPort},
-				&dtls.Config{Certificate: cert, PrivateKey: key, CipherSuites: []dtls.CipherSuiteID{cipherSuite}},
-			)
-			if err != nil {
-				errChan <- err
-				return
-			}
-			simpleReadWrite(errChan, clientChan, conn, nil, &messageRecvCount)
-		}()
-
-		// DTLS Server
-		go func() {
-			listener, err := dtls.Listen("udp",
-				&net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: serverPort},
-				&dtls.Config{Certificate: cert, PrivateKey: key, CipherSuites: []dtls.CipherSuiteID{cipherSuite}},
-			)
-			if err != nil {
-				errChan <- err
-				return
-			}
-			serverReady <- struct{}{}
-			conn, err := listener.Accept()
-			if err != nil {
-				errChan <- err
-				return
-			}
+		cfg := &dtls.Config{Certificate: cert, PrivateKey: key, CipherSuites: []dtls.CipherSuiteID{cipherSuite}}
+		assertE2ECommunication(cfg, cfg, t)
 
-			simpleReadWrite(errChan, serverChan, conn, listener, &messageRecvCount)
-		}()
-
-		func() {
-			seenClient, seenServer := false, false
-			for {
-				select {
-				case err := <-errChan:
-					t.Fatal(err)
-				case <-time.After(testTimeLimit):
-					t.Fatalf("Test timeout, seenClient %t seenServer %t", seenClient, seenServer)
-				case clientMsg := <-clientChan:
-					if clientMsg != testMessage {
-						t.Fatalf("clientMsg does not equal test message: %s %s", clientMsg, testMessage)
-					}
-
-					seenClient = true
-					if seenClient && seenServer {
-						return
-					}
-				case serverMsg := <-serverChan:
-					if serverMsg != testMessage {
-						t.Fatalf("serverMsg does not equal test message: %s %s", serverMsg, testMessage)
-					}
-
-					seenServer = true
-					if seenClient && seenServer {
-						return
-					}
-				}
-			}
-		}()
+	}
 
-		// We have seen Client+Server communication! Now count goroutines to make sure we haven't leaked
-		time.Sleep(time.Second) // TODO racey
-		goRoutineCount := runtime.NumGoroutine()
-		if goRoutineCount != expectedGoRoutineCount {
-			if err := pprof.Lookup("goroutine").WriteTo(os.Stderr, 1); err != nil {
-				t.Fatal(err)
-			}
-			t.Fatalf("goRoutineCount != expectedGoRoutineCount, possible leak: %d %d", goRoutineCount, expectedGoRoutineCount)
-		}
+}
 
-	}
+func TestPionE2ESimplePSK(t *testing.T) {
+	lim := test.TimeOut(time.Second * 30)
+	defer lim.Stop()
 
+	report := test.CheckRoutines(t)
+	defer report()
+
+	cfg := &dtls.Config{
+		PSK: func(hint []byte) ([]byte, error) {
+			return []byte{0xAB, 0xC1, 0x23}, nil
+		},
+		PSKIdentityHint: []byte{0x01, 0x02, 0x03, 0x04, 0x05},
+		CipherSuites:    []dtls.CipherSuiteID{dtls.TLS_PSK_WITH_AES_128_CCM8},
+	}
+	assertE2ECommunication(cfg, cfg, t)
 }