diff --git a/src/codec.ts b/src/codec.ts
index cac699b..3a71021 100644
--- a/src/codec.ts
+++ b/src/codec.ts
@@ -96,7 +96,7 @@ export class NoiseSecureTransferEncoder implements Encoder {
       return;
     }
 
-    const preparedPayload = this.hsResult.writeMessage(message.payload, this.hsResult.nametagsOutbound);
+    const preparedPayload = this.hsResult.writeMessage(message.payload);
 
     const payload = preparedPayload.serialize();
 
@@ -136,7 +136,7 @@ export class NoiseSecureTransferDecoder implements Decoder<NoiseSecureMessage> {
 
     const payloadV2 = PayloadV2.deserialize(proto.payload);
 
-    const decryptedPayload = this.hsResult.readMessage(payloadV2, this.hsResult.nametagsInbound);
+    const decryptedPayload = this.hsResult.readMessage(payloadV2);
 
     return new NoiseSecureMessage(proto, decryptedPayload);
   }
diff --git a/src/crypto.ts b/src/crypto.ts
index e4bd455..5183674 100644
--- a/src/crypto.ts
+++ b/src/crypto.ts
@@ -24,10 +24,7 @@ export function intoCurve25519Key(s: Uint8Array): bytes32 {
 }
 
 export function getHKDF(ck: bytes32, ikm: Uint8Array): Hkdf {
-  const hkdf = new HKDF(SHA256, ikm, ck);
-  const okmU8Array = hkdf.expand(96);
-  const okm = okmU8Array;
-
+  const okm = getHKDFRaw(ck, ikm, 96);
   const k1 = okm.subarray(0, 32);
   const k2 = okm.subarray(32, 64);
   const k3 = okm.subarray(64, 96);
diff --git a/src/handshake.ts b/src/handshake.ts
index 69770c5..3cfd11a 100644
--- a/src/handshake.ts
+++ b/src/handshake.ts
@@ -55,11 +55,14 @@ export class HandshakeResult {
   // due to nonce exhaustion, then the application must delete the CipherState and terminate the session.
 
   // Writes an encrypted message using the proper Cipher State
-  writeMessage(transportMessage: Uint8Array, outboundMessageNametagBuffer: MessageNametagBuffer): PayloadV2 {
+  writeMessage(
+    transportMessage: Uint8Array,
+    outboundMessageNametagBuffer: MessageNametagBuffer | undefined = undefined
+  ): PayloadV2 {
     const payload2 = new PayloadV2();
 
     // We set the message nametag using the input buffer
-    payload2.messageNametag = outboundMessageNametagBuffer.pop();
+    payload2.messageNametag = (outboundMessageNametagBuffer ?? this.nametagsOutbound).pop();
 
     // According to 35/WAKU2-NOISE RFC, no Handshake protocol information is sent when exchanging messages
     // This correspond to setting protocol-id to 0
@@ -74,13 +77,16 @@ export class HandshakeResult {
 
   // Reads an encrypted message using the proper Cipher State
   // Decryption is attempted only if the input PayloadV2 has a messageNametag equal to the one expected
-  readMessage(readPayload2: PayloadV2, inboundMessageNametagBuffer: MessageNametagBuffer): Uint8Array {
+  readMessage(
+    readPayload2: PayloadV2,
+    inboundMessageNametagBuffer: MessageNametagBuffer | undefined = undefined
+  ): Uint8Array {
     // The output decrypted message
     let message = new Uint8Array();
 
     // If the message nametag does not correspond to the nametag expected in the inbound message nametag buffer
     // an error is raised (to be handled externally, i.e. re-request lost messages, discard, etc.)
-    const nametagIsOk = inboundMessageNametagBuffer.checkNametag(readPayload2.messageNametag);
+    const nametagIsOk = (inboundMessageNametagBuffer ?? this.nametagsInbound).checkNametag(readPayload2.messageNametag);
     if (!nametagIsOk) {
       throw new Error("nametag is not ok");
     }
@@ -95,7 +101,7 @@ export class HandshakeResult {
         // We unpad the decrypted message
         message = pkcs7.unpad(paddedMessage);
         // The message successfully decrypted, we can delete the first element of the inbound Message Nametag Buffer
-        inboundMessageNametagBuffer.delete(1);
+        this.nametagsInbound.delete(1);
       } catch (err) {
         console.debug("A read message failed decryption. Returning empty message as plaintext.");
         message = new Uint8Array();
diff --git a/src/waku-noise-pairing.spec.ts b/src/waku-noise-pairing.spec.ts
index 0f8cd49..88e48ce 100644
--- a/src/waku-noise-pairing.spec.ts
+++ b/src/waku-noise-pairing.spec.ts
@@ -246,51 +246,40 @@ describe("Waku Noise Sessions", () => {
       expect(uint8ArrayEquals(message, readMessage!.payload)).to.be.true;
     }
 
-    // TODO
-    // TODO
-    // TODO
-    // TODO
-    // TODO
-    /*
     // We test how nametag buffers help in detecting lost messages
     // Alice writes two messages to Bob, but only the second is received
+    let message = randomBytes(32, rng);
+    let payload2 = aliceHSResult.writeMessage(message);
+    message = randomBytes(32, rng);
+    payload2 = aliceHSResult.writeMessage(message);
     try {
-      const message = randomBytes(32, rng);
-      payload2 = aliceHSResult.writeMessage(message, aliceHSResult.nametagsOutbound);
-      message = randomBytes(32, rng);
-      payload2 = aliceHSResult.writeMessage(aliceHSResult.nametagsOutbound);
-    } catch (NoiseSomeMessagesWereLost) {
-      let readMessage = readMessage(
-        bobHSResult,
-        payload2,
-        (inboundMessageNametagBuffer = bobHSResult.nametagsInbound)
-      ).get();
+      bobHSResult.readMessage(payload2);
+      expect(false, "should not reach here").to.be.true;
+    } catch (err) {
+      let message;
+      if (err instanceof Error) message = err.message;
+      else message = String(err);
+      expect(message).to.be.equals("nametag is not ok");
     }
 
     // We adjust bob nametag buffer for next test (i.e. the missed message is correctly recovered)
     bobHSResult.nametagsInbound.delete(2);
-    let message = randomBytes(32, rng);
-    payload2 = writeMessage(bobHSResult, message, (outboundMessageNametagBuffer = bobHSResult.nametagsOutbound));
-    readMessage = readMessage(
-      aliceHSResult,
-      payload2,
-      (inboundMessageNametagBuffer = aliceHSResult.nametagsInbound)
-    ).get();
-
-    expect(uint8ArrayEquals(message, readMessage!.payload)).to.be.true;
+    message = randomBytes(32, rng);
+    payload2 = bobHSResult.writeMessage(message);
+    const readMessage = aliceHSResult.readMessage(payload2);
+    expect(uint8ArrayEquals(message, readMessage)).to.be.true;
 
     // We test if a missing nametag is correctly detected
+    message = randomBytes(32, rng);
+    payload2 = aliceHSResult.writeMessage(message);
+    bobHSResult.nametagsInbound.delete(1);
     try {
-      const message = randomBytes(32, rng);
-      const payload2 = aliceHSResult.writeMessage(message, aliceHSResult.nametagsOutbound);
-      bobHSResult.nametagsInbound.delete(1);
-    } catch (NoiseMessageNametagError) {
-      let readMessage = readMessage(
-        bobHSResult,
-        payload2,
-        (inboundMessageNametagBuffer = bobHSResult.nametagsInbound)
-      ).get();
+      bobHSResult.readMessage(payload2);
+    } catch (err) {
+      let message;
+      if (err instanceof Error) message = err.message;
+      else message = String(err);
+      expect(message).to.be.equals("nametag is not ok");
     }
-    */
   });
 });