From 16aebac0c7e738e9e2921824192f1e4b49f7dee7 Mon Sep 17 00:00:00 2001
From: Jazz Turner-Baggs <473256+jazzz@users.noreply.github.com>
Date: Thu, 9 Oct 2025 18:30:56 -0700
Subject: [PATCH] Integrate doubleratchet encryption to PrivateV1

---
 protos/encryption.proto                   | 12 +++--
 src/chat_sdk/client.nim                   |  7 ++-
 src/chat_sdk/conversations/private_v1.nim | 64 ++++++++++++++++++-----
 src/chat_sdk/errors.nim                   | 14 +++++
 src/chat_sdk/inbox.nim                    |  6 ++-
 src/chat_sdk/types.nim                    |  3 --
 6 files changed, 83 insertions(+), 23 deletions(-)
 create mode 100644 src/chat_sdk/errors.nim

diff --git a/protos/encryption.proto b/protos/encryption.proto
index 99708fb..cda0fb0 100644
--- a/protos/encryption.proto
+++ b/protos/encryption.proto
@@ -9,7 +9,7 @@ message EncryptedPayload {
 
     oneof encryption {
         encryption.Plaintext plaintext = 1;
-		encryption.Ecies ecies = 2;
+		encryption.Doubleratchet doubleratchet = 2;
     }
 }
 
@@ -17,8 +17,10 @@ message Plaintext {
     bytes payload=1;
 }
 
-message Ecies {
-    bytes encrypted_bytes=1;
-    bytes ephemeral_pubkey = 2;
-    bytes tag = 3;
+message Doubleratchet {
+    bytes dh = 1;       // 32 byte array
+    uint32 msgNum = 2;
+    uint32 prevChainLen = 3;
+    bytes ciphertext = 4;
+    string aux = 5;
 }
diff --git a/src/chat_sdk/client.nim b/src/chat_sdk/client.nim
index cf09226..6d53c4d 100644
--- a/src/chat_sdk/client.nim
+++ b/src/chat_sdk/client.nim
@@ -21,6 +21,7 @@ import #local
   conversations/convo_impl,
   crypto,
   delivery/waku_client,
+  errors,
   identity,
   inbox,
   proto_types,
@@ -215,7 +216,11 @@ proc newPrivateConversation*(client: Client,
       msgId: string): Future[void] {.async.} =
     client.notifyDeliveryAck(conversation, msgId)
 
-  let convo = initPrivateV1(client.identity(), destPubkey, "default", deliveryAckCb)
+  # TODO: remove placeholder key
+  var key : array[32, byte]
+  key[2]=2
+
+  var convo = initPrivateV1Sender(client.identity(), destPubkey, key, deliveryAckCb)
   client.addConversation(convo)
 
   # TODO: Subscribe to new content topic
diff --git a/src/chat_sdk/conversations/private_v1.nim b/src/chat_sdk/conversations/private_v1.nim
index d64ce99..4f1f327 100644
--- a/src/chat_sdk/conversations/private_v1.nim
+++ b/src/chat_sdk/conversations/private_v1.nim
@@ -14,12 +14,15 @@ import ../delivery/waku_client
 
 import ../[
   identity,
+  errors,
   proto_types,
   types,
   utils
 ]
 import convo_type
 
+import ../../naxolotl as nax
+ 
 
 type
   PrivateV1* = ref object of Conversation
@@ -29,6 +32,7 @@ type
     topic: string
     participants: seq[PublicKey]
     discriminator: string
+    doubleratchet: naxolotl.Doubleratchet
 
 proc getTopic*(self: PrivateV1): string =
   ## Returns the topic for the PrivateV1 conversation.
@@ -56,11 +60,32 @@ proc calcMsgId(self: PrivateV1, msgBytes: seq[byte]): string =
   result = getBlake2b(s, 16, "")
 
 
-proc encrypt*(convo: PrivateV1, frame: PrivateV1Frame): EncryptedPayload =
-  result = EncryptedPayload(plaintext: Plaintext(payload: encode(frame)))
+proc encrypt*(convo: PrivateV1, plaintext: var seq[byte]): EncryptedPayload =
+
+  let (header, ciphertext) = convo.doubleratchet.encrypt(plaintext) #TODO: Associated Data
+
+  result = EncryptedPayload(doubleratchet: proto_types.DoubleRatchet(
+    dh: toSeq(header.dhPublic),
+    msgNum: header.msgNumber,
+    prevChainLen: header.prevChainLen,
+    ciphertext: ciphertext)
+  )
+
+proc decrypt*(convo: PrivateV1, enc: EncryptedPayload): Result[seq[byte], ChatError]  =
+  # Ensure correct type as received
+  if enc.doubleratchet.ciphertext == @[]:
+    return err(ChatError(code: errTypeError, context: "Expected doubleratchet encrypted payload got ???"))
+
+  let dr = enc.doubleratchet
+
+  var header = DrHeader(
+    msgNumber: dr.msgNum,
+    prevChainLen: dr.prevChainLen
+  )
+  copyMem(addr header.dhPublic[0], unsafeAddr dr.dh[0], dr.dh.len) # TODO: Avoid this copy
+
+  convo.doubleratchet.decrypt(header, dr.ciphertext, @[]).mapErr(proc(e: NaxolotlError): ChatError = ChatError(code: errWrapped, context: repr(e) ))
 
-proc decrypt*(convo: PrivateV1, enc: EncryptedPayload): PrivateV1Frame =
-  result = decode(enc.plaintext.payload, PrivateV1Frame).get()
 
 
 proc wireCallbacks(convo: PrivateV1, deliveryAckCb: proc(
@@ -91,8 +116,8 @@ proc wireCallbacks(convo: PrivateV1, deliveryAckCb: proc(
 
 
 
-proc initPrivateV1*(owner: Identity, participant: PublicKey,
-        discriminator: string = "default", deliveryAckCb: proc(
+proc initPrivateV1*(owner: Identity, participant: PublicKey, seedKey: array[32, byte],
+        discriminator: string = "default", isSender: bool, deliveryAckCb: proc(
         conversation: Conversation,
       msgId: string): Future[void] {.async.} = nil):
           PrivateV1 =
@@ -107,7 +132,8 @@ proc initPrivateV1*(owner: Identity, participant: PublicKey,
     owner: owner,
     topic: derive_topic(participants, discriminator),
     participants: participants,
-    discriminator: discriminator
+    discriminator: discriminator,
+    doubleratchet: initDoubleratchet(seedKey, owner.privateKey.bytes, participant.bytes, isSender)
   )
 
   result.wireCallbacks(deliveryAckCb)
@@ -115,19 +141,28 @@ proc initPrivateV1*(owner: Identity, participant: PublicKey,
   result.sdsClient.ensureChannel(result.getConvoId()).isOkOr:
     raise newException(ValueError, "bad sds channel")
 
+
+proc initPrivateV1Sender*(owner:Identity, participant: PublicKey, seedKey: array[32, byte], deliveryAckCb: proc(
+        conversation: Conversation, msgId: string): Future[void] {.async.} = nil): PrivateV1 =
+        initPrivateV1(owner, participant, seedKey, "default", true, deliveryAckCb)
+
+proc initPrivateV1Recipient*(owner:Identity, participant: PublicKey, seedKey: array[32, byte], deliveryAckCb: proc(
+        conversation: Conversation, msgId: string): Future[void] {.async.} = nil): PrivateV1 =
+        initPrivateV1(owner, participant, seedKey, "default", false, deliveryAckCb)
+
+
 proc sendFrame(self: PrivateV1, ds: WakuClient,
     msg: PrivateV1Frame):  Future[MessageId]{.async.} =
 
   let frameBytes = encode(msg)
   let msgId = self.calcMsgId(frameBytes)
-  let sdsPayload = self.sdsClient.wrapOutgoingMessage(frameBytes, msgId,
+  var sdsPayload = self.sdsClient.wrapOutgoingMessage(frameBytes, msgId,
       self.getConvoId()).valueOr:
     raise newException(ValueError, fmt"sds wrapOutgoingMessage failed: {repr(error)}")
 
-  let encryptedBytes = EncryptedPayload(plaintext: Plaintext(
-      payload: sdsPayload))
+  let encryptedPayload = self.encrypt(sdsPayload)
 
-  discard ds.sendPayload(self.getTopic(), encryptedBytes.toEnvelope(
+  discard ds.sendPayload(self.getTopic(), encryptedPayload.toEnvelope(
       self.getConvoId()))
 
   result = msgId
@@ -144,9 +179,12 @@ proc handleFrame*[T: ConversationStore](convo: PrivateV1, client: T,
   let enc = decode(bytes, EncryptedPayload).valueOr:
     raise newException(ValueError, fmt"Failed to decode EncryptedPayload: {repr(error)}")
 
-  # TODO: Decrypt the payload
+  let plaintext = convo.decrypt(enc).valueOr:
+    error "decryption failed", error = error
+    return
+
   let (frameData, missingDeps, channelId) = convo.sdsClient.unwrapReceivedMessage(
-      enc.plaintext.payload).valueOr:
+      plaintext).valueOr:
     raise newException(ValueError, fmt"Failed to unwrap SDS message:{repr(error)}")
 
   debug "sds unwrap", convo = convo.id(), missingDeps = missingDeps,
diff --git a/src/chat_sdk/errors.nim b/src/chat_sdk/errors.nim
new file mode 100644
index 0000000..e7b8008
--- /dev/null
+++ b/src/chat_sdk/errors.nim
@@ -0,0 +1,14 @@
+import strformat
+
+type
+  ChatError* = object of CatchableError
+    code*: ErrorCode
+    context*: string
+  
+  ErrorCode* = enum
+    errTypeError
+    errWrapped
+
+
+proc `$`*(x: ChatError): string =
+  fmt"ChatError(code={$x.code}, context: {x.context})"
diff --git a/src/chat_sdk/inbox.nim b/src/chat_sdk/inbox.nim
index d31dac5..6ddff39 100644
--- a/src/chat_sdk/inbox.nim
+++ b/src/chat_sdk/inbox.nim
@@ -78,7 +78,11 @@ proc createPrivateV1FromInvite*[T: ConversationStore](client: T,
       msgId: string): Future[void] {.async.} =
     client.notifyDeliveryAck(conversation, msgId)
 
-  let convo = initPrivateV1(client.identity(), destPubkey, "default", deliveryAckCb)
+  # TODO: remove placeholder key
+  var key : array[32, byte]
+  key[2]=2
+
+  let convo = initPrivateV1Recipient(client.identity(), destPubkey, key, deliveryAckCb)
   notice "Creating PrivateV1 conversation", client = client.getId(),
       topic = convo.getConvoId()
   client.addConversation(convo)
diff --git a/src/chat_sdk/types.nim b/src/chat_sdk/types.nim
index 6f44be7..d8acc92 100644
--- a/src/chat_sdk/types.nim
+++ b/src/chat_sdk/types.nim
@@ -1,4 +1 @@
-type ChatError* = string
-
-
 type MessageId* = string