Support for DTLS session authentication context updates via outbound dgram (#63)

* Support for DTLS session authentication context updates via outbound datagram context

* Address comments

* Remove hacky ways of setrting context

Author: akolosov-n

kotlin-mbedtls-netty/src/main/kotlin/org/opencoap/ssl/netty/DtlsChannelHandler.kt

@@ -21,7 +21,6 @@ import io.netty.channel.ChannelHandlerContext
 import io.netty.channel.ChannelPromise
 import io.netty.channel.socket.DatagramPacket
 import org.opencoap.ssl.SslConfig
-import org.opencoap.ssl.SslException
 import org.opencoap.ssl.transport.ByteBufferPacket
 import org.opencoap.ssl.transport.DtlsServer
 import org.opencoap.ssl.transport.DtlsSessionLifecycleCallbacks
@@ -92,23 +91,9 @@ class DtlsChannelHandler @JvmOverloads constructor(
         when (msg) {
             is DatagramPacketWithContext -> {
                 write(msg, promise, ctx)
-                if (msg.sessionContext.sessionExpirationHint) {
-                    promise.toCompletableFuture().thenAccept {
-                        dtlsServer.closeSession(msg.recipient())
-                    }
-                }
+                dtlsServer.handleOutboundDtlsSessionContext(msg.recipient(), msg.sessionContext, promise.toCompletableFuture())
             }
             is DatagramPacket -> write(msg, promise, ctx)
-            is SessionAuthenticationContext -> {
-                msg.map.forEach { (key, value) ->
-                    if (!dtlsServer.putSessionAuthenticationContext(msg.adr, key, value)) {
-                        promise.setFailure(SslException("Session does not exists"))
-                    }
-                }
-                if (!promise.isDone) {
-                    promise.setSuccess()
-                }
-            }
 
             else -> ctx.write(msg, promise)
         }

kotlin-mbedtls-netty/src/main/kotlin/org/opencoap/ssl/netty/SessionAuthenticationContext.kt

@@ -32,6 +32,9 @@ class EchoHandler : ChannelInboundHandlerAdapter() {
         val authContext = (sessionContext.authenticationContext["AUTH"] ?: "")
         val dgramContent = dgram.content().toByteArray()
         val goToSleep = dgramContent.toString(Charset.defaultCharset()).endsWith(":sleep")
+        val newAuthContext = dgramContent.toString(Charset.defaultCharset())
+            .takeIf { it.startsWith("auth:") }
+            ?.substringAfter(":")
 
         val reply = ctx.alloc().buffer(dgramContent.size + 20)
         reply.writeBytes(echoPrefix)
@@ -43,7 +46,10 @@ class EchoHandler : ChannelInboundHandlerAdapter() {
                 reply,
                 dgram.sender(),
                 null,
-                sessionContext.copy(sessionExpirationHint = goToSleep)
+                sessionContext.copy(
+                    authenticationContext = newAuthContext?.let { mapOf("AUTH" to it) } ?: emptyMap(),
+                    sessionSuspensionHint = goToSleep
+                )
             )
         )
     }

kotlin-mbedtls-netty/src/test/kotlin/org/opencoap/ssl/netty/EchoHandler.kt

@@ -36,7 +36,6 @@ import org.opencoap.ssl.EmptyCidSupplier
 import org.opencoap.ssl.PskAuth
 import org.opencoap.ssl.RandomCidSupplier
 import org.opencoap.ssl.SslConfig
-import org.opencoap.ssl.SslException
 import org.opencoap.ssl.netty.NettyHelpers.createBootstrap
 import org.opencoap.ssl.transport.DtlsServer
 import org.opencoap.ssl.transport.HashMapSessionStore
@@ -173,15 +172,16 @@ class NettyTest {
     }
 
     @Test
-    fun `should forward authentication context`() {
+    fun `should forward authentication context passed inside outbound datagram`() {
         // connect and handshake
         val client = NettyTransportAdapter.connect(clientConf, srvAddress).mapToString()
 
         assertTrue(client.send("hi").await())
         assertEquals("ECHO:hi", client.receive(5.seconds).await())
 
         // when
-        srvChannel.writeAndFlush(SessionAuthenticationContext(client.localAddress(), mapOf("AUTH" to "007:"))).get()
+        assertTrue(client.send("auth:007:").await())
+        assertEquals("ECHO:auth:007:", client.receive(5.seconds).await())
 
         // then
         assertTrue(client.send("hi").await())
@@ -190,13 +190,6 @@ class NettyTest {
         client.close()
     }
 
-    @Test
-    fun `should fail to forward authentication context for non existing client`() {
-        assertThatThrownBy {
-            srvChannel.writeAndFlush(SessionAuthenticationContext(localAddress(1), mapOf("AUTH" to "007:"))).get()
-        }.hasRootCause(SslException("Session does not exists"))
-    }
-
     @Test
     fun `server should load session from store`() {
         sessionStore.write(StoredSessionPair.cid, SessionWithContext(StoredSessionPair.srvSession, mapOf(), Instant.ofEpochSecond(123456789)))

kotlin-mbedtls-netty/src/test/kotlin/org/opencoap/ssl/netty/NettyTest.kt

@@ -94,18 +94,22 @@ class DtlsServer(
         return (sessions[peerAddress] as? DtlsSession)?.encrypt(plainPacket)
     }
 
-    fun putSessionAuthenticationContext(adr: InetSocketAddress, key: String, value: String?): Boolean {
-        return when (val s = sessions[adr]) {
-            is DtlsSession -> {
-                if (value != null) {
-                    s.authenticationContext += (key to value)
-                } else {
-                    s.authenticationContext -= key
+    private fun updateSessionAuthenticationContext(adr: InetSocketAddress, authCtxUpdate: Map<String, String?>): Boolean {
+        if (authCtxUpdate.isEmpty()) return true
+
+        return when (val s = sessions[adr] as? DtlsSession) {
+            null -> false
+
+            else -> {
+                authCtxUpdate.forEach { (key, value) ->
+                    if (value != null) {
+                        s.authenticationContext += (key to value)
+                    } else {
+                        s.authenticationContext -= key
+                    }
                 }
                 true
             }
-
-            else -> false
         }
     }
 
@@ -118,13 +122,22 @@ class DtlsServer(
         }
     }
 
-    fun closeSession(addr: InetSocketAddress) {
+    private fun closeSession(addr: InetSocketAddress) {
         sessions.remove(addr)?.apply {
             storeAndClose()
             logger.info("[{}] [CID:{}] DTLS session was stored", peerAddress, (this as? DtlsSession)?.sessionContext?.cid?.toHex() ?: "na")
         }
     }
 
+    fun handleOutboundDtlsSessionContext(adr: InetSocketAddress, ctx: DtlsSessionContext, writeFuture: CompletableFuture<Boolean>) {
+        if (ctx.sessionSuspensionHint) {
+            writeFuture.thenAccept {
+                closeSession(adr)
+            }
+        }
+        updateSessionAuthenticationContext(adr, ctx.authenticationContext)
+    }
+
     fun loadSession(sessBuf: SessionWithContext?, adr: InetSocketAddress, cid: ByteArray): Boolean {
         return try {
             if (sessBuf == null) {

kotlin-mbedtls/src/main/kotlin/org/opencoap/ssl/transport/DtlsServer.kt

@@ -97,15 +97,11 @@ class DtlsServerTransport private constructor(
 
         when {
             encPacket == null -> completedFuture(false)
-            packet.sessionContext.sessionExpirationHint -> {
-                transport.send(encPacket).thenApply { isSuccess ->
-                    if (isSuccess) {
-                        dtlsServer.closeSession(packet.peerAddress)
-                    }
-                    isSuccess
+            else -> {
+                transport.send(encPacket).also {
+                    dtlsServer.handleOutboundDtlsSessionContext(packet.peerAddress, packet.sessionContext, it)
                 }
             }
-            else -> transport.send(encPacket)
         }
     }.thenCompose(Function.identity())
 
@@ -118,9 +114,4 @@ class DtlsServerTransport private constructor(
         }.get(30, TimeUnit.SECONDS)
         executor.shutdown()
     }
-
-    fun putSessionAuthenticationContext(adr: InetSocketAddress, key: String, value: String?): CompletableFuture<Boolean> =
-        executor.supply {
-            dtlsServer.putSessionAuthenticationContext(adr, key, value)
-        }
 }

kotlin-mbedtls/src/main/kotlin/org/opencoap/ssl/transport/DtlsServerTransport.kt

@@ -25,7 +25,7 @@ data class DtlsSessionContext @JvmOverloads constructor(
     val peerCertificateSubject: String? = null,
     val cid: ByteArray? = null,
     val sessionStartTimestamp: Instant? = null,
-    val sessionExpirationHint: Boolean = false
+    val sessionSuspensionHint: Boolean = false
 ) {
     companion object {
         @JvmField
@@ -45,7 +45,7 @@ data class DtlsSessionContext @JvmOverloads constructor(
             if (!cid.contentEquals(other.cid)) return false
         } else if (other.cid != null) return false
         if (sessionStartTimestamp != other.sessionStartTimestamp) return false
-        if (sessionExpirationHint != other.sessionExpirationHint) return false
+        if (sessionSuspensionHint != other.sessionSuspensionHint) return false
 
         return true
     }
@@ -55,7 +55,7 @@ data class DtlsSessionContext @JvmOverloads constructor(
         result = 31 * result + (peerCertificateSubject?.hashCode() ?: 0)
         result = 31 * result + (cid?.contentHashCode() ?: 0)
         result = 31 * result + (sessionStartTimestamp?.hashCode() ?: 0)
-        result = 31 * result + (sessionExpirationHint.hashCode())
+        result = 31 * result + (sessionSuspensionHint.hashCode())
         return result
     }
 }

kotlin-mbedtls/src/main/kotlin/org/opencoap/ssl/transport/DtlsSessionContext.kt

@@ -70,8 +70,13 @@ class DtlsServerTransportTest {
         if (msg == "error") {
             throw Exception("error")
         } else if (msg.startsWith("Authenticate:")) {
-            server.putSessionAuthenticationContext(packet.peerAddress, "auth", msg.substring(12))
-            server.send(Packet("OK".toByteBuffer(), packet.peerAddress))
+            server.send(
+                Packet(
+                    "OK".toByteBuffer(),
+                    packet.peerAddress,
+                    DtlsSessionContext(authenticationContext = mapOf("auth" to msg.substring(12)))
+                )
+            )
         } else {
             val ctx = (packet.sessionContext.authenticationContext["auth"] ?: "")
             server.send(packet.map { "$msg:resp$ctx".toByteBuffer() })
@@ -456,25 +461,14 @@ class DtlsServerTransportTest {
     }
 
     @Test
-    fun `should set and use session context`() {
-        // given
-        server = DtlsServerTransport.create(conf, sessionStore = sessionStore)
-        val serverReceived = server.receive(1.seconds)
-        // and, client connected
+    fun `should set and use session context passed inside outbound datagram`() {
+        server = DtlsServerTransport.create(conf, expireAfter = 100.millis, sessionStore = sessionStore, lifecycleCallbacks = sslLifecycleCallbacks).listen(echoHandler)
+        // client connected
         val client = DtlsTransmitter.connect(server, clientConfig).await()
-        client.send("hello!")
-        assertEquals("hello!", serverReceived.await().buffer.decodeToString())
-
-        // when, session context is set
-        assertTrue(server.putSessionAuthenticationContext(serverReceived.await().peerAddress, "auth", "id:dev-007").await())
-
-        // and, client sends messages
-        client.send("msg1")
-        client.send("msg2")
-
-        // then
-        assertEquals(mapOf("auth" to "id:dev-007"), server.receive(1.seconds).await().sessionContext.authenticationContext)
-        assertEquals(mapOf("auth" to "id:dev-007"), server.receive(1.seconds).await().sessionContext.authenticationContext)
+        client.send("Authenticate:dev-007")
+        assertEquals("OK", client.receiveString())
+        client.send("hi")
+        assertEquals("hi:resp:dev-007", client.receiveString())
 
         client.close()
     }
@@ -491,7 +485,7 @@ class DtlsServerTransportTest {
         assertEquals("dupa", client.receive(1.seconds).await())
 
         client.send("sleep")
-        server.send(Packet("sleep".toByteBuffer(), serverReceived.await().peerAddress, sessionContext = DtlsSessionContext(sessionExpirationHint = true)))
+        server.send(Packet("sleep".toByteBuffer(), serverReceived.await().peerAddress, sessionContext = DtlsSessionContext(sessionSuspensionHint = true)))
         assertEquals("sleep", client.receive(1.seconds).await())
 
         await.atMost(5.seconds).untilAsserted {