;;; jabber-test-omemo-module.el --- Tests for jabber-omemo-module  -*- lexical-binding: t; -*-

;;; Commentary:

;; OMEMO native C module interface.

;;; Code:

(require 'ert)
(require 'jabber-omemo-core)

;; The C module signals `jabber-omemo-error' but doesn't register it
;; with `define-error'.  ERT's :type matching needs the hierarchy.
(define-error 'jabber-omemo-error "OMEMO error")

;;; Group 1: Module loading

(ert-deftest jabber-test-omemo-module-provides-feature ()
  "The module provides the `jabber-omemo-core' feature."
  (should (featurep 'jabber-omemo-core)))

(ert-deftest jabber-test-omemo-module-functions-bound ()
  "All internal functions are bound after loading."
  (should (fboundp 'jabber-omemo--setup-store))
  (should (fboundp 'jabber-omemo--deserialize-store))
  (should (fboundp 'jabber-omemo--serialize-store))
  (should (fboundp 'jabber-omemo--get-bundle))
  (should (fboundp 'jabber-omemo--rotate-signed-pre-key))
  (should (fboundp 'jabber-omemo--refill-pre-keys))
  (should (fboundp 'jabber-omemo--encrypt-message))
  (should (fboundp 'jabber-omemo--decrypt-message))
  (should (fboundp 'jabber-omemo--make-session))
  (should (fboundp 'jabber-omemo--initiate-session))
  (should (fboundp 'jabber-omemo--serialize-session))
  (should (fboundp 'jabber-omemo--deserialize-session))
  (should (fboundp 'jabber-omemo--encrypt-key))
  (should (fboundp 'jabber-omemo--decrypt-key))
  (should (fboundp 'jabber-omemo--session-skipped-keys))
  (should (fboundp 'jabber-omemo--session-set-skipped-keys))
  (should (fboundp 'jabber-omemo--heartbeat))
  (should (fboundp 'jabber-omemo--aesgcm-decrypt))
  (should (fboundp 'jabber-omemo--aesgcm-encrypt)))

;;; Group 2: Store lifecycle

(ert-deftest jabber-test-omemo-module-setup-store-returns-unibyte ()
  "setup-store returns a non-empty unibyte string."
  (let ((blob (jabber-omemo--setup-store)))
    (should (stringp blob))
    (should (not (multibyte-string-p blob)))
    (should (> (length blob) 0))))

(ert-deftest jabber-test-omemo-module-deserialize-store-returns-user-ptr ()
  "deserialize-store returns a user-ptr."
  (let* ((blob (jabber-omemo--setup-store))
         (ptr (jabber-omemo--deserialize-store blob)))
    (should (user-ptrp ptr))))

(ert-deftest jabber-test-omemo-module-store-round-trip ()
  "Serializing a deserialized store produces the same blob."
  (let* ((blob1 (jabber-omemo--setup-store))
         (ptr (jabber-omemo--deserialize-store blob1))
         (blob2 (jabber-omemo--serialize-store ptr)))
    (should (string= blob1 blob2))))

(ert-deftest jabber-test-omemo-module-protobuf-rejects-bad-varints ()
  "Deserialization rejects oversized and unterminated uint32 varints."
  (dolist (blob (list (unibyte-string #x08 #x80 #x80 #x80 #x80 #x80 #x00)
                      (unibyte-string #x08 #x80 #x80 #x80 #x80 #x80)
                      (unibyte-string #x08 #x80 #x80 #x80 #x80 #x10)))
    (should-error (jabber-omemo--deserialize-store blob)
                  :type 'jabber-omemo-error)))

(ert-deftest jabber-test-omemo-module-protobuf-rejects-truncated-field ()
  "Deserialization rejects a length field beyond the remaining input."
  (should-error
   (jabber-omemo--deserialize-store (unibyte-string #x12 #x20))
   :type 'jabber-omemo-error))

(ert-deftest jabber-test-omemo-module-pointer-kinds-are-distinct ()
  "Store and session wrappers reject the other native pointer kind."
  (let ((store (jabber-omemo--deserialize-store (jabber-omemo--setup-store)))
        (session (jabber-omemo--make-session)))
    (dolist (form (list (lambda () (jabber-omemo--serialize-store session))
                        (lambda () (jabber-omemo--get-bundle session))
                        (lambda () (jabber-omemo--serialize-session store))
                        (lambda () (jabber-omemo--used-pre-key-id store))
                        (lambda () (jabber-omemo--encrypt-key store
                                                              (make-string 32 0)))
                        (lambda () (jabber-omemo--decrypt-key
                                    session session nil "bad"))
                        (lambda () (jabber-omemo--heartbeat session session))))
      (should-error (funcall form) :type 'jabber-omemo-error))))

;;; Group 3: Bundle extraction

(ert-deftest jabber-test-omemo-module-get-bundle-plist-keys ()
  "get-bundle returns a plist with expected keys."
  (let* ((blob (jabber-omemo--setup-store))
         (ptr (jabber-omemo--deserialize-store blob))
         (bundle (jabber-omemo--get-bundle ptr)))
    (should (plist-get bundle :identity-key))
    (should (plist-get bundle :signed-pre-key))
    (should (plist-get bundle :signed-pre-key-id))
    (should (plist-get bundle :signature))
    (should (plist-get bundle :pre-keys))))

(ert-deftest jabber-test-omemo-module-identity-key-length ()
  "Identity key is 33 bytes (0x05 prefix + 32-byte key)."
  (let* ((blob (jabber-omemo--setup-store))
         (ptr (jabber-omemo--deserialize-store blob))
         (bundle (jabber-omemo--get-bundle ptr))
         (ik (plist-get bundle :identity-key)))
    (should (= (length ik) 33))
    (should (= (aref ik 0) #x05))))

(ert-deftest jabber-test-omemo-module-signed-pre-key-length ()
  "Signed pre-key is 33 bytes."
  (let* ((blob (jabber-omemo--setup-store))
         (ptr (jabber-omemo--deserialize-store blob))
         (bundle (jabber-omemo--get-bundle ptr))
         (spk (plist-get bundle :signed-pre-key)))
    (should (= (length spk) 33))))

(ert-deftest jabber-test-omemo-module-signature-length ()
  "Signature is 64 bytes."
  (let* ((blob (jabber-omemo--setup-store))
         (ptr (jabber-omemo--deserialize-store blob))
         (bundle (jabber-omemo--get-bundle ptr))
         (sig (plist-get bundle :signature)))
    (should (= (length sig) 64))))

(ert-deftest jabber-test-omemo-module-pre-keys-count ()
  "Pre-keys list has 100 entries."
  (let* ((blob (jabber-omemo--setup-store))
         (ptr (jabber-omemo--deserialize-store blob))
         (bundle (jabber-omemo--get-bundle ptr))
         (pks (plist-get bundle :pre-keys)))
    (should (= (length pks) 100))))

(ert-deftest jabber-test-omemo-module-pre-key-format ()
  "Each pre-key is (id . key) with integer id and 33-byte key."
  (let* ((blob (jabber-omemo--setup-store))
         (ptr (jabber-omemo--deserialize-store blob))
         (bundle (jabber-omemo--get-bundle ptr))
         (pks (plist-get bundle :pre-keys))
         (first-pk (car pks)))
    (should (consp first-pk))
    (should (integerp (car first-pk)))
    (should (= (length (cdr first-pk)) 33))))

;;; Group 4: Key rotation

(ert-deftest jabber-test-omemo-module-rotate-changes-spk-id ()
  "Rotating the signed pre-key changes its ID."
  (let* ((blob (jabber-omemo--setup-store))
         (ptr (jabber-omemo--deserialize-store blob))
         (id-before (plist-get (jabber-omemo--get-bundle ptr)
                               :signed-pre-key-id)))
    (jabber-omemo--rotate-signed-pre-key ptr)
    (let ((id-after (plist-get (jabber-omemo--get-bundle ptr)
                               :signed-pre-key-id)))
      (should-not (= id-before id-after)))))

(ert-deftest jabber-test-omemo-module-refill-pre-keys ()
  "refill-pre-keys does not error on a fresh store."
  (let* ((blob (jabber-omemo--setup-store))
         (ptr (jabber-omemo--deserialize-store blob)))
    (jabber-omemo--refill-pre-keys ptr)
    (should t)))

(ert-deftest jabber-test-omemo-module-refilled-store-varint-round-trip ()
  "Normal pre-key replacement survives the ID 127/128 size boundary."
  (let ((store (jabber-omemo--deserialize-store (jabber-omemo--setup-store))))
    (dotimes (_ 28)
      (let ((id (caar (plist-get (jabber-omemo--get-bundle store) :pre-keys))))
        (should (jabber-omemo--remove-pre-key store id))
        (jabber-omemo--refill-pre-keys store)))
    (let* ((ids (mapcar #'car (plist-get (jabber-omemo--get-bundle store)
                                        :pre-keys)))
           (blob (jabber-omemo--serialize-store store))
           (restored (jabber-omemo--deserialize-store blob)))
      (should (memq 128 ids))
      (should (seq-some (lambda (id) (< id 128)) ids))
      ;; Compare as a boolean so ERT never prints serialized key material.
      (should (eq t (string= blob (jabber-omemo--serialize-store restored)))))))

;;; Group 5: Message encrypt/decrypt

(ert-deftest jabber-test-omemo-module-encrypt-returns-plist ()
  "encrypt-message returns plist with :key, :iv, :ciphertext."
  (let ((result (jabber-omemo--encrypt-message
                 (encode-coding-string "hello" 'utf-8))))
    (should (plist-get result :key))
    (should (plist-get result :iv))
    (should (plist-get result :ciphertext))))

(ert-deftest jabber-test-omemo-module-encrypt-key-length ()
  "Encryption key is 32 bytes (16 AES + 16 auth tag)."
  (let* ((result (jabber-omemo--encrypt-message
                  (encode-coding-string "test" 'utf-8)))
         (key (plist-get result :key)))
    (should (= (length key) 32))))

(ert-deftest jabber-test-omemo-module-encrypt-iv-length ()
  "IV is 12 bytes."
  (let* ((result (jabber-omemo--encrypt-message
                  (encode-coding-string "test" 'utf-8)))
         (iv (plist-get result :iv)))
    (should (= (length iv) 12))))

(ert-deftest jabber-test-omemo-module-decrypt-recovers-plaintext ()
  "Decrypting an encrypted message recovers the original."
  (let* ((msg (encode-coding-string "Hello there!" 'utf-8))
         (enc (jabber-omemo--encrypt-message msg))
         (dec (jabber-omemo--decrypt-message
               (plist-get enc :key)
               (plist-get enc :iv)
               (plist-get enc :ciphertext))))
    (should (string= msg dec))))

(ert-deftest jabber-test-omemo-module-encrypt-decrypt-round-trip-utf8 ()
  "Round-trip works with UTF-8 content."
  (let* ((msg (encode-coding-string "Hallo Welt! \u00e4\u00f6\u00fc" 'utf-8))
         (enc (jabber-omemo--encrypt-message msg))
         (dec (jabber-omemo--decrypt-message
               (plist-get enc :key)
               (plist-get enc :iv)
               (plist-get enc :ciphertext))))
    (should (string= msg dec))))

(ert-deftest jabber-test-omemo-module-decrypt-requires-exact-key-and-iv ()
  "Message decryption rejects short and oversized keys and IVs."
  (let* ((enc (jabber-omemo--encrypt-message "payload"))
         (key (plist-get enc :key))
         (iv (plist-get enc :iv))
         (ciphertext (plist-get enc :ciphertext)))
    (dolist (bad-key (list (substring key 1) (concat key "x")))
      (should-error (jabber-omemo--decrypt-message bad-key iv ciphertext)
                    :type 'jabber-omemo-error))
    (dolist (bad-iv (list (substring iv 1) (concat iv "x")))
      (should-error (jabber-omemo--decrypt-message key bad-iv ciphertext)
                    :type 'jabber-omemo-error))))

(ert-deftest jabber-test-omemo-module-decrypt-wrong-key-signals-error ()
  "Decrypting with a wrong key signals jabber-omemo-error."
  (let* ((msg (encode-coding-string "secret" 'utf-8))
         (enc (jabber-omemo--encrypt-message msg))
         (bad-key (make-string 32 ?x)))
    (should-error
     (jabber-omemo--decrypt-message
      bad-key
      (plist-get enc :iv)
      (plist-get enc :ciphertext))
     :type 'jabber-omemo-error)))

;;; Group 6: Session lifecycle

(ert-deftest jabber-test-omemo-module-initiate-session-returns-user-ptr ()
  "initiate-session returns a user-ptr."
  (let* ((alice-blob (jabber-omemo--setup-store))
         (alice (jabber-omemo--deserialize-store alice-blob))
         (bob-blob (jabber-omemo--setup-store))
         (bob (jabber-omemo--deserialize-store bob-blob))
         (bundle (jabber-omemo--get-bundle bob))
         (pk (car (plist-get bundle :pre-keys)))
         (session (jabber-omemo--initiate-session
                   alice
                   (plist-get bundle :signature)
                   (plist-get bundle :signed-pre-key)
                   (plist-get bundle :identity-key)
                   (cdr pk)
                   (plist-get bundle :signed-pre-key-id)
                   (car pk))))
    (should (user-ptrp session))))

(ert-deftest jabber-test-omemo-module-serialize-session-returns-unibyte ()
  "serialize-session returns a non-empty unibyte string."
  (let* ((alice-blob (jabber-omemo--setup-store))
         (alice (jabber-omemo--deserialize-store alice-blob))
         (bob-blob (jabber-omemo--setup-store))
         (bob (jabber-omemo--deserialize-store bob-blob))
         (bundle (jabber-omemo--get-bundle bob))
         (pk (car (plist-get bundle :pre-keys)))
         (session (jabber-omemo--initiate-session
                   alice
                   (plist-get bundle :signature)
                   (plist-get bundle :signed-pre-key)
                   (plist-get bundle :identity-key)
                   (cdr pk)
                   (plist-get bundle :signed-pre-key-id)
                   (car pk)))
         (blob (jabber-omemo--serialize-session session)))
    (should (stringp blob))
    (should (not (multibyte-string-p blob)))
    (should (> (length blob) 0))))

(ert-deftest jabber-test-omemo-module-session-round-trip ()
  "Serializing a deserialized session preserves the data."
  (let* ((alice-blob (jabber-omemo--setup-store))
         (alice (jabber-omemo--deserialize-store alice-blob))
         (bob-blob (jabber-omemo--setup-store))
         (bob (jabber-omemo--deserialize-store bob-blob))
         (bundle (jabber-omemo--get-bundle bob))
         (pk (car (plist-get bundle :pre-keys)))
         (session (jabber-omemo--initiate-session
                   alice
                   (plist-get bundle :signature)
                   (plist-get bundle :signed-pre-key)
                   (plist-get bundle :identity-key)
                   (cdr pk)
                   (plist-get bundle :signed-pre-key-id)
                   (car pk)))
         (blob1 (jabber-omemo--serialize-session session))
         (session2 (jabber-omemo--deserialize-session blob1))
         (blob2 (jabber-omemo--serialize-session session2)))
    (should (user-ptrp session2))
    (should (string= blob1 blob2))))

(defun jabber-test-omemo-module--set-u32 (string offset value)
  "Store big-endian VALUE in STRING at OFFSET."
  (dotimes (i 4)
    (aset string (+ offset i) (logand 255 (ash value (* -8 (- 3 i))))))
  string)

(ert-deftest jabber-test-omemo-module-session-envelope-validation ()
  "Malformed envelopes, excess keys, and trailing data are rejected."
  (let* ((session (jabber-omemo--make-session))
         (blob (jabber-omemo--serialize-session session))
         (bad-version (copy-sequence blob))
         (too-many (copy-sequence blob)))
    (aset bad-version 11 2)
    (jabber-test-omemo-module--set-u32 too-many 16 1001)
    (dolist (bad (list bad-version too-many (concat blob "x")))
      (should-error (jabber-omemo--deserialize-session bad)
                    :type 'jabber-omemo-error))))

(ert-deftest jabber-test-omemo-module-session-legacy-blob-loads ()
  "The raw picomemo payload remains a valid legacy session blob."
  (let* ((session (jabber-omemo--make-session))
         (blob (jabber-omemo--serialize-session session))
         (raw-size (+ (ash (aref blob 12) 24)
                      (ash (aref blob 13) 16)
                      (ash (aref blob 14) 8)
                      (aref blob 15)))
         (raw (substring blob 20 (+ 20 raw-size))))
    (should-not (jabber-omemo--legacy-session-blob-p blob))
    (should (jabber-omemo--legacy-session-blob-p raw))
    (should (user-ptrp (jabber-omemo--deserialize-session raw)))))

(ert-deftest jabber-test-omemo-module-skipped-key-fifo-limit ()
  "A session retains the newest 1000 imported skipped keys in FIFO order."
  (let* ((session (jabber-omemo--make-session))
         (dh (make-string 32 ?d))
         (keys (cl-loop for i below 1001
                        collect (list i dh (make-string 32 ?m)))))
    (jabber-omemo--session-set-skipped-keys session keys)
    (let ((stored (jabber-omemo--session-skipped-keys session)))
      (should (= 1000 (length stored)))
      (should (= 1 (caar stored)))
      (should (= 1000 (car (car (last stored))))))))

(ert-deftest jabber-test-omemo-module-initiate-session-bad-signature ()
  "initiate-session with a bad signature signals an error."
  (let* ((alice-blob (jabber-omemo--setup-store))
         (alice (jabber-omemo--deserialize-store alice-blob))
         (bob-blob (jabber-omemo--setup-store))
         (bob (jabber-omemo--deserialize-store bob-blob))
         (bundle (jabber-omemo--get-bundle bob))
         (pk (car (plist-get bundle :pre-keys)))
         (bad-sig (make-string 64 0)))
    (should-error
     (jabber-omemo--initiate-session
      alice bad-sig
      (plist-get bundle :signed-pre-key)
      (plist-get bundle :identity-key)
      (cdr pk)
      (plist-get bundle :signed-pre-key-id)
      (car pk))
     :type 'jabber-omemo-error)))

(ert-deftest jabber-test-omemo-module-initiate-requires-exact-inputs ()
  "Session initiation rejects short and oversized cryptographic inputs."
  (let* ((alice (jabber-omemo--deserialize-store (jabber-omemo--setup-store)))
         (bob (jabber-omemo--deserialize-store (jabber-omemo--setup-store)))
         (bundle (jabber-omemo--get-bundle bob))
         (signature (plist-get bundle :signature))
         (spk (plist-get bundle :signed-pre-key))
         (ik (plist-get bundle :identity-key))
         (pk (car (plist-get bundle :pre-keys)))
         (inputs (list signature spk ik (cdr pk))))
    (dotimes (index 4)
      (dolist (bad (list (substring (nth index inputs) 1)
                         (concat (nth index inputs) "x")))
        (let ((changed (copy-sequence inputs)))
          (setf (nth index changed) bad)
          (should-error
           (apply #'jabber-omemo--initiate-session
                  alice (append changed
                                (list (plist-get bundle :signed-pre-key-id)
                                      (car pk))))
           :type 'jabber-omemo-error))))))

;;; Group 7: Key encrypt/decrypt round-trip

(ert-deftest jabber-test-omemo-module-encrypt-key-returns-plist ()
  "encrypt-key returns a plist with :data and :pre-key-p."
  (let* ((alice-blob (jabber-omemo--setup-store))
         (alice (jabber-omemo--deserialize-store alice-blob))
         (bob-blob (jabber-omemo--setup-store))
         (bob (jabber-omemo--deserialize-store bob-blob))
         (bundle (jabber-omemo--get-bundle bob))
         (pk (car (plist-get bundle :pre-keys)))
         (session (jabber-omemo--initiate-session
                   alice
                   (plist-get bundle :signature)
                   (plist-get bundle :signed-pre-key)
                   (plist-get bundle :identity-key)
                   (cdr pk)
                   (plist-get bundle :signed-pre-key-id)
                   (car pk)))
         (test-key (make-string 32 ?A))
         (result (jabber-omemo--encrypt-key session test-key)))
    (should (plist-get result :data))
    (should (plist-member result :pre-key-p))))

(ert-deftest jabber-test-omemo-module-first-message-is-pre-key ()
  "First encrypted key message has :pre-key-p = t."
  (let* ((alice-blob (jabber-omemo--setup-store))
         (alice (jabber-omemo--deserialize-store alice-blob))
         (bob-blob (jabber-omemo--setup-store))
         (bob (jabber-omemo--deserialize-store bob-blob))
         (bundle (jabber-omemo--get-bundle bob))
         (pk (car (plist-get bundle :pre-keys)))
         (session (jabber-omemo--initiate-session
                   alice
                   (plist-get bundle :signature)
                   (plist-get bundle :signed-pre-key)
                   (plist-get bundle :identity-key)
                   (cdr pk)
                   (plist-get bundle :signed-pre-key-id)
                   (car pk)))
         (test-key (make-string 32 ?B))
         (result (jabber-omemo--encrypt-key session test-key)))
    (should (eq t (plist-get result :pre-key-p)))))

(ert-deftest jabber-test-omemo-module-key-encrypt-decrypt-round-trip ()
  "Full Alice/Bob key encrypt/decrypt round-trip."
  (let* ((alice-blob (jabber-omemo--setup-store))
         (alice (jabber-omemo--deserialize-store alice-blob))
         (bob-blob (jabber-omemo--setup-store))
         (bob (jabber-omemo--deserialize-store bob-blob))
         (bundle (jabber-omemo--get-bundle bob))
         (pk (car (plist-get bundle :pre-keys)))
         ;; Alice initiates session with Bob's bundle
         (alice-session (jabber-omemo--initiate-session
                         alice
                         (plist-get bundle :signature)
                         (plist-get bundle :signed-pre-key)
                         (plist-get bundle :identity-key)
                         (cdr pk)
                         (plist-get bundle :signed-pre-key-id)
                         (car pk)))
         ;; Alice encrypts a 32-byte key
         (original-key (make-string 32 ?K))
         (encrypted (jabber-omemo--encrypt-key alice-session original-key))
         (enc-data (plist-get encrypted :data))
         (is-prekey (plist-get encrypted :pre-key-p))
         ;; Bob decrypts with a fresh (empty) session
         (bob-session (jabber-omemo--make-session))
         (decrypted (jabber-omemo--decrypt-key
                     bob-session bob is-prekey enc-data)))
    (should (string= original-key decrypted))))

(ert-deftest jabber-test-omemo-module-decrypt-key-rejects-bad-protobuf ()
  "Key decryption rejects hostile varints and truncated length fields."
  (let ((store (jabber-omemo--deserialize-store (jabber-omemo--setup-store)))
        (session (jabber-omemo--make-session)))
    (dolist (message
             (list (unibyte-string #x33 #x08 #x80 #x80 #x80 #x80 #x80 #x00)
                   (unibyte-string #x33 #x08 #x80 #x80 #x80 #x80 #x80)
                   (unibyte-string #x33 #x08 #x80 #x80 #x80 #x80 #x10)
                   (unibyte-string #x33 #x12 #x20)))
      (should-error
       (jabber-omemo--decrypt-key session store t message)
       :type 'jabber-omemo-error))))

(ert-deftest jabber-test-omemo-module-consecutive-messages-are-pre-key ()
  "Consecutive messages from initiator stay pre-key until reply."
  (let* ((alice-blob (jabber-omemo--setup-store))
         (alice (jabber-omemo--deserialize-store alice-blob))
         (bob-blob (jabber-omemo--setup-store))
         (bob (jabber-omemo--deserialize-store bob-blob))
         (bundle (jabber-omemo--get-bundle bob))
         (pk (car (plist-get bundle :pre-keys)))
         (alice-session (jabber-omemo--initiate-session
                         alice
                         (plist-get bundle :signature)
                         (plist-get bundle :signed-pre-key)
                         (plist-get bundle :identity-key)
                         (cdr pk)
                         (plist-get bundle :signed-pre-key-id)
                         (car pk)))
         (key1 (make-string 32 ?A))
         (enc1 (jabber-omemo--encrypt-key alice-session key1))
         (key2 (make-string 32 ?B))
         (enc2 (jabber-omemo--encrypt-key alice-session key2)))
    ;; Both messages are pre-key since Bob hasn't replied yet
    (should (eq t (plist-get enc1 :pre-key-p)))
    (should (eq t (plist-get enc2 :pre-key-p)))))

;;; Group 8: Heartbeat

(ert-deftest jabber-test-omemo-module-heartbeat-nil-on-fresh-session ()
  "heartbeat returns nil on a fresh session."
  (let* ((alice-blob (jabber-omemo--setup-store))
         (alice (jabber-omemo--deserialize-store alice-blob))
         (bob-blob (jabber-omemo--setup-store))
         (bob (jabber-omemo--deserialize-store bob-blob))
         (bundle (jabber-omemo--get-bundle bob))
         (pk (car (plist-get bundle :pre-keys)))
         (session (jabber-omemo--initiate-session
                   alice
                   (plist-get bundle :signature)
                   (plist-get bundle :signed-pre-key)
                   (plist-get bundle :identity-key)
                   (cdr pk)
                   (plist-get bundle :signed-pre-key-id)
                   (car pk))))
    (should-not (jabber-omemo--heartbeat session alice))))

;;; Group 9: aesgcm encrypt/decrypt (XEP-0454)

(ert-deftest jabber-test-omemo-module-aesgcm-encrypt-returns-plist ()
  "aesgcm-encrypt returns a plist with :key, :iv, :ciphertext."
  (let ((result (jabber-omemo--aesgcm-encrypt "hello")))
    (should (plist-get result :key))
    (should (plist-get result :iv))
    (should (plist-get result :ciphertext))))

(ert-deftest jabber-test-omemo-module-aesgcm-encrypt-key-length ()
  "aesgcm-encrypt returns a 32-byte key."
  (let ((result (jabber-omemo--aesgcm-encrypt "test")))
    (should (= 32 (length (plist-get result :key))))))

(ert-deftest jabber-test-omemo-module-aesgcm-encrypt-iv-length ()
  "aesgcm-encrypt returns a 12-byte IV."
  (let ((result (jabber-omemo--aesgcm-encrypt "test")))
    (should (= 12 (length (plist-get result :iv))))))

(ert-deftest jabber-test-omemo-module-aesgcm-encrypt-ciphertext-length ()
  "Ciphertext is plaintext length + 16 bytes for GCM tag."
  (let* ((plaintext "hello world")
         (result (jabber-omemo--aesgcm-encrypt plaintext)))
    (should (= (+ (length plaintext) 16)
               (length (plist-get result :ciphertext))))))

(ert-deftest jabber-test-omemo-module-aesgcm-encrypt-decrypt-round-trip ()
  "Encrypting then decrypting recovers the original plaintext."
  (let* ((plaintext "The quick brown fox jumps over the lazy dog")
         (enc (jabber-omemo--aesgcm-encrypt plaintext))
         (dec (jabber-omemo--aesgcm-decrypt
               (plist-get enc :key)
               (plist-get enc :iv)
               (plist-get enc :ciphertext))))
    (should (string= plaintext dec))))

(ert-deftest jabber-test-omemo-module-aesgcm-encrypt-binary-round-trip ()
  "Encrypt/decrypt round-trip works for binary data."
  (let* ((plaintext (make-string 256 0))
         (_ (dotimes (i 256) (aset plaintext i i)))
         (enc (jabber-omemo--aesgcm-encrypt plaintext))
         (dec (jabber-omemo--aesgcm-decrypt
               (plist-get enc :key)
               (plist-get enc :iv)
               (plist-get enc :ciphertext))))
    (should (string= plaintext dec))))

(ert-deftest jabber-test-omemo-module-aesgcm-encrypt-unique-keys ()
  "Each call generates different keys."
  (let ((r1 (jabber-omemo--aesgcm-encrypt "test"))
        (r2 (jabber-omemo--aesgcm-encrypt "test")))
    (should-not (string= (plist-get r1 :key) (plist-get r2 :key)))))

;;; Group 10: One-time pre-key removal

(ert-deftest jabber-test-omemo-module-remove-pre-key ()
  "remove-pre-key zeroes exactly the matching slot."
  (let* ((store (jabber-omemo--deserialize-store (jabber-omemo--setup-store)))
         (before (plist-get (jabber-omemo--get-bundle store) :pre-keys))
         (victim (car (nth 3 before))))
    (should (jabber-omemo--remove-pre-key store victim))
    (let ((after (plist-get (jabber-omemo--get-bundle store) :pre-keys)))
      (should (= (1- (length before)) (length after)))
      (should-not (assq victim after)))
    ;; Removing an already-removed id is a no-op.
    (should-not (jabber-omemo--remove-pre-key store victim))))

(ert-deftest jabber-test-omemo-module-remove-pre-key-unknown-id ()
  "remove-pre-key returns nil for an id not in the store."
  (let ((store (jabber-omemo--deserialize-store (jabber-omemo--setup-store))))
    (should-not (jabber-omemo--remove-pre-key store 999999))
    (should-not (jabber-omemo--remove-pre-key store 0))))

(ert-deftest jabber-test-omemo-module-remove-then-refill-uses-fresh-id ()
  "Refilling after removal restores the count with a new id."
  (let* ((store (jabber-omemo--deserialize-store (jabber-omemo--setup-store)))
         (before (plist-get (jabber-omemo--get-bundle store) :pre-keys))
         (victim (car (car before))))
    (jabber-omemo--remove-pre-key store victim)
    (jabber-omemo--refill-pre-keys store)
    (let ((after (plist-get (jabber-omemo--get-bundle store) :pre-keys)))
      (should (= (length before) (length after)))
      (should-not (assq victim after)))))

(ert-deftest jabber-test-omemo-module-used-pre-key-id-fresh-session ()
  "A session that never consumed a pre-key reports id 0."
  (should (zerop (jabber-omemo--used-pre-key-id
                  (jabber-omemo--make-session)))))

(ert-deftest jabber-test-omemo-module-used-pre-key-id-after-decrypt ()
  "After a pre-key decrypt the session reports the consumed id."
  (let* ((alice (jabber-omemo--deserialize-store (jabber-omemo--setup-store)))
         (bob (jabber-omemo--deserialize-store (jabber-omemo--setup-store)))
         (bundle (jabber-omemo--get-bundle bob))
         (pk (car (plist-get bundle :pre-keys)))
         (alice-session (jabber-omemo--initiate-session
                         alice
                         (plist-get bundle :signature)
                         (plist-get bundle :signed-pre-key)
                         (plist-get bundle :identity-key)
                         (cdr pk)
                         (plist-get bundle :signed-pre-key-id)
                         (car pk)))
         (encrypted (jabber-omemo--encrypt-key alice-session
                                               (make-string 32 ?K)))
         (bob-session (jabber-omemo--make-session)))
    (jabber-omemo--decrypt-key bob-session bob
                               (plist-get encrypted :pre-key-p)
                               (plist-get encrypted :data))
    (should (= (car pk) (jabber-omemo--used-pre-key-id bob-session)))))

;;; Group: Skipped message keys

(defmacro jabber-test-omemo-module--with-session-pair (&rest body)
  "Run BODY with `alice', `bob', and `alice-session' bound.
Alice has initiated a session towards Bob's bundle."
  (declare (indent 0) (debug t))
  `(let* ((alice (jabber-omemo--deserialize-store (jabber-omemo--setup-store)))
          (bob (jabber-omemo--deserialize-store (jabber-omemo--setup-store)))
          (bundle (jabber-omemo--get-bundle bob))
          (pk (car (plist-get bundle :pre-keys)))
          (alice-session (jabber-omemo--initiate-session
                          alice
                          (plist-get bundle :signature)
                          (plist-get bundle :signed-pre-key)
                          (plist-get bundle :identity-key)
                          (cdr pk)
                          (plist-get bundle :signed-pre-key-id)
                          (car pk))))
     (ignore alice-session)
     ,@body))

(ert-deftest jabber-test-omemo-module-skipped-keys-accessors-roundtrip ()
  "Skipped keys set on a session read back unchanged."
  (let* ((session (jabber-omemo--make-session))
         (keys (list (list 3 (make-string 32 ?d) (make-string 32 ?m))
                     (list 7 (make-string 32 ?e) (make-string 32 ?n)))))
    (should (null (jabber-omemo--session-skipped-keys session)))
    (jabber-omemo--session-set-skipped-keys session keys)
    (should (equal keys (jabber-omemo--session-skipped-keys session)))
    ;; Replacing clears the previous list.
    (jabber-omemo--session-set-skipped-keys session nil)
    (should (null (jabber-omemo--session-skipped-keys session)))))

(ert-deftest jabber-test-omemo-module-skipped-key-import-is-atomic ()
  "Rejected skipped-key imports preserve the session's original keys."
  (let* ((session (jabber-omemo--make-session))
         (dh (make-string 32 ?d))
         (mk (make-string 32 ?m))
         (original (list (list 7 dh mk)))
         (bad-lists
          (list (cons (list 8 dh mk) 9)
                (list 8 dh)
                (list 8 dh mk 'extra)
                (list (cons 8 (cons dh mk)))
                (list (list 'bad dh mk))
                (list (list -1 dh mk))
                (list (list (1+ #xffffffff) dh mk))
                (list (list 8 'bad mk))
                (list (list 8 dh 'bad))
                (list (list 8 (make-string 31 ?d) mk))
                (list (list 8 (make-string 33 ?d) mk))
                (list (list 8 dh (make-string 31 ?m)))
                (list (list 8 dh (make-string 33 ?m))))))
    (jabber-omemo--session-set-skipped-keys session original)
    (dolist (bad bad-lists)
      (should-error (jabber-omemo--session-set-skipped-keys session bad)
                    :type 'jabber-omemo-error)
      (should (equal original
                     (jabber-omemo--session-skipped-keys session))))))

(ert-deftest jabber-test-omemo-module-skipped-keys-survive-finalizers ()
  "Enumerating skipped keys is safe while other sessions are finalized."
  (let* ((session (jabber-omemo--make-session))
         (keys (list (list 3 (make-string 32 ?d) (make-string 32 ?m))
                     (list 7 (make-string 32 ?e) (make-string 32 ?n)))))
    (jabber-omemo--session-set-skipped-keys session keys)
    (dotimes (i 200)
      (let ((disposable (jabber-omemo--make-session)))
        (jabber-omemo--session-set-skipped-keys
         disposable
         (list (list i
                     (make-string 32 (+ ?a (% i 26)))
                     (make-string 32 (+ ?A (% i 26))))))))
    (let ((gc-cons-threshold 1))
      (should (equal keys (jabber-omemo--session-skipped-keys session))))
    (garbage-collect)
    (should (equal keys (jabber-omemo--session-skipped-keys session)))))

(ert-deftest jabber-test-omemo-module-out-of-order-decrypt ()
  "A message skipped over in the ratchet still decrypts afterwards."
  (jabber-test-omemo-module--with-session-pair
    (let* ((k1 (make-string 32 ?1))
           (k2 (make-string 32 ?2))
           (k3 (make-string 32 ?3))
           (m1 (jabber-omemo--encrypt-key alice-session k1))
           (m2 (jabber-omemo--encrypt-key alice-session k2))
           (m3 (jabber-omemo--encrypt-key alice-session k3))
           (bob-session (jabber-omemo--make-session)))
      ;; Deliver message 3 first: keys for 1 and 2 must be skipped.
      (should (string= k3 (jabber-omemo--decrypt-key
                           bob-session bob
                           (plist-get m3 :pre-key-p) (plist-get m3 :data))))
      (should (= 2 (length (jabber-omemo--session-skipped-keys bob-session))))
      ;; The late arrivals decrypt from the skipped keys, single-use.
      (should (string= k1 (jabber-omemo--decrypt-key
                           bob-session bob
                           (plist-get m1 :pre-key-p) (plist-get m1 :data))))
      (should (string= k2 (jabber-omemo--decrypt-key
                           bob-session bob
                           (plist-get m2 :pre-key-p) (plist-get m2 :data))))
      (should (null (jabber-omemo--session-skipped-keys bob-session))))))

(ert-deftest jabber-test-omemo-module-failed-decrypt-keeps-skipped-key ()
  "A corrupted late message does not consume its skipped key."
  (jabber-test-omemo-module--with-session-pair
    (let* ((k1 (make-string 32 ?1))
           (k2 (make-string 32 ?2))
           (m1 (jabber-omemo--encrypt-key alice-session k1))
           (m2 (jabber-omemo--encrypt-key alice-session k2))
           (bob-session (jabber-omemo--make-session)))
      (jabber-omemo--decrypt-key
       bob-session bob (plist-get m2 :pre-key-p) (plist-get m2 :data))
      (let* ((corrupt (copy-sequence (plist-get m1 :data)))
             (last (1- (length corrupt))))
        (aset corrupt last (logxor 1 (aref corrupt last)))
        (should-error
         (jabber-omemo--decrypt-key
          bob-session bob (plist-get m1 :pre-key-p) corrupt)
         :type 'jabber-omemo-error))
      (should (= 1 (length (jabber-omemo--session-skipped-keys bob-session))))
      (should (string= k1 (jabber-omemo--decrypt-key
                           bob-session bob
                           (plist-get m1 :pre-key-p)
                           (plist-get m1 :data))))
      (should (null (jabber-omemo--session-skipped-keys bob-session)))
      (should-error
       (jabber-omemo--decrypt-key
        bob-session bob (plist-get m1 :pre-key-p) (plist-get m1 :data))
       :type 'jabber-omemo-error))))

(ert-deftest jabber-test-omemo-module-skipped-keys-survive-reserialization ()
  "Skipped keys carried over to a reloaded session still decrypt."
  (jabber-test-omemo-module--with-session-pair
    (let* ((k1 (make-string 32 ?1))
           (k2 (make-string 32 ?2))
           (m1 (jabber-omemo--encrypt-key alice-session k1))
           (m2 (jabber-omemo--encrypt-key alice-session k2))
           (bob-session (jabber-omemo--make-session)))
      ;; Deliver message 2 first; message 1's key is skipped.
      (should (string= k2 (jabber-omemo--decrypt-key
                           bob-session bob
                           (plist-get m2 :pre-key-p) (plist-get m2 :data))))
      ;; Simulate a restart.  The session blob owns its skipped keys.
      (let* ((blob (jabber-omemo--serialize-session bob-session))
             (restored (jabber-omemo--deserialize-session blob)))
        (should (= 1 (length
                      (jabber-omemo--session-skipped-keys restored))))
        (should (string= k1 (jabber-omemo--decrypt-key
                             restored bob
                             (plist-get m1 :pre-key-p)
                             (plist-get m1 :data))))))))

(provide 'jabber-test-omemo-module)
;;; jabber-test-omemo-module.el ends here
