e2ee tests

Makefile

@@ -15,7 +15,7 @@ all:
 	go build -ldflags "-X main.commit=${COMMIT}" -tags "${GOTAGS}" -o release/telegabber
 
 test:
-	go test -tags "${GOTAGS}" -v ./config ./ ./telegram ./xmpp ./xmpp/gateway ./persistence ./telegram/formatter ./badger
+	go test -tags "${GOTAGS}" -v ./config ./ ./telegram ./xmpp ./xmpp/gateway ./persistence ./telegram/formatter ./badger ./e2ee/...
 
 lint:
 	$(GOPATH)/bin/golint ./...

e2ee/omemo/libsignal/address.go

@@ -1,14 +1,12 @@
 package libsignal
 
 /*
+#include <stdlib.h>
 #include <signal_protocol.h>
 */
 import "C"
 
-import (
-	"runtime"
-	"unsafe"
-)
+import "unsafe"
 
 // Address identifies a Signal Protocol session peer: a logical recipient
 // name (for OMEMO, a bare JID) plus a numeric device ID.
@@ -17,23 +15,35 @@ type Address struct {
 	DeviceID uint32
 }
 
-// withCAddress converts addr to a C signal_protocol_address and invokes fn
-// with a pointer to it. The pointer - and the memory its name field borrows
-// from addr.Name - is only valid for the duration of fn; it must not be
-// retained past fn's return.
-func withCAddress(addr Address, fn func(*C.signal_protocol_address)) {
-	nameBytes := []byte(addr.Name)
-	var namePtr *C.char
-	if len(nameBytes) > 0 {
-		namePtr = (*C.char)(unsafe.Pointer(&nameBytes[0]))
+// cAddress is a C-heap-allocated signal_protocol_address (struct and name
+// string both), for the specific cases - session_builder_create and
+// session_cipher_create - where the C library stores the *pointer* we pass
+// it and keeps dereferencing it for the object's entire lifetime (a
+// shallow assignment, e.g. `result_cipher->remote_address = remote_address;`
+// in session_cipher.c, confirmed by reading the source - not a defensive
+// copy). A transient, freed-after-the-call address (as a short-lived cgo
+// helper would naturally produce) becomes a use-after-free the moment the
+// constructor returns; this type's lifetime must instead match its owning
+// SessionBuilder/SessionCipher, freed only in their Close methods.
+type cAddress struct {
+	ptr *C.signal_protocol_address
 }
-	cAddr := C.signal_protocol_address{
-		name:      namePtr,
-		name_len:  C.size_t(len(nameBytes)),
-		device_id: C.int32_t(addr.DeviceID),
+
+func newCAddress(addr Address) *cAddress {
+	namePtr := C.CString(addr.Name)
+	c := (*C.signal_protocol_address)(C.malloc(C.size_t(unsafe.Sizeof(C.signal_protocol_address{}))))
+	c.name = namePtr
+	c.name_len = C.size_t(len(addr.Name))
+	c.device_id = C.int32_t(addr.DeviceID)
+	return &cAddress{ptr: c}
+}
+
+func (a *cAddress) free() {
+	if a.ptr != nil {
+		C.free(unsafe.Pointer(a.ptr.name))
+		C.free(unsafe.Pointer(a.ptr))
+		a.ptr = nil
 	}
-	fn(&cAddr)
-	runtime.KeepAlive(nameBytes)
 }
 
 // addressFromC copies a C signal_protocol_address into a Go Address.

e2ee/omemo/libsignal/cipher.go

@@ -35,6 +35,7 @@ type CiphertextMessage struct {
 // SessionCipher performs encrypt/decrypt for an already-established session.
 type SessionCipher struct {
 	raw  *C.session_cipher
+	addr *cAddress
 }
 
 // NewSessionCipher creates a session cipher for storeCtx's local identity
@@ -42,19 +43,19 @@ type SessionCipher struct {
 // cipher *produces* (ProtocolVersionLegacy or ProtocolVersionModern);
 // decrypting adapts to whatever version the incoming message declares.
 func NewSessionCipher(ctx *Context, storeCtx *StoreContext, remote Address, version int) (*SessionCipher, error) {
+	addr := newCAddress(remote)
+
 	var raw *C.session_cipher
-	var code C.int
-	withCAddress(remote, func(cAddr *C.signal_protocol_address) {
-		code = C.session_cipher_create(&raw, storeCtx.raw, cAddr, ctx.raw)
-	})
-	if code != C.SG_SUCCESS {
+	if code := C.session_cipher_create(&raw, storeCtx.raw, addr.ptr, ctx.raw); code != C.SG_SUCCESS {
+		addr.free()
 		return nil, newError("session_cipher_create", int(code))
 	}
 	if err := setCipherVersion(raw, version); err != nil {
 		C.session_cipher_free(raw)
+		addr.free()
 		return nil, err
 	}
-	return &SessionCipher{raw: raw}, nil
+	return &SessionCipher{raw: raw, addr: addr}, nil
 }
 
 // Close frees the session cipher. Do not use it afterward.
@@ -63,6 +64,7 @@ func (c *SessionCipher) Close() {
 		C.session_cipher_free(c.raw)
 		c.raw = nil
 	}
+	c.addr.free()
 }
 
 // Encrypt encrypts plaintext for this cipher's session, in whichever

e2ee/omemo/libsignal/keys.go

@@ -236,6 +236,16 @@ func deserializeSignedPreKey(ctx *Context, record []byte) (*C.session_signed_pre
 // prekey: its numeric id, raw public key, and (both) signature forms the
 // library maintains - a legacy signature and an OMEMO signature (libomemo-c
 // computes both regardless of which protocol version ends up using it).
+//
+// When publishing a bundle (or constructing a RemoteBundle from one),
+// picking the wrong one of these two for the target's protocol version
+// fails verification with SG_ERR_INVALID_KEY, not a version-mismatch
+// error: session_builder_process_pre_key_bundle re-serializes the signed
+// prekey's public key via ec_public_key_serialize (version < 4, the
+// 33-byte DJB_TYPE-prefixed legacy form) or ec_public_key_serialize_omemo
+// (version >= 4, the raw 32-byte Montgomery form) and checks it against
+// whichever signature is supplied - so use Signature for
+// ProtocolVersionLegacy and SignatureOMEMO for ProtocolVersionModern.
 type SignedPreKeyInfo struct {
 	ID             uint32
 	PublicKey      []byte

e2ee/omemo/libsignal/memstore_test.go

@@ -0,0 +1,192 @@
+package libsignal_test
+
+import (
+	"bytes"
+	"fmt"
+	"strconv"
+	"strings"
+	"sync"
+
+	"dev.narayana.im/narayana/telegabber/e2ee/omemo/libsignal"
+)
+
+// memStore is a minimal in-memory libsignal.Store used only to give the
+// round-trip tests in this package somewhere to persist state during a
+// session. It is not the real Store implementation - that's badgerstore,
+// which has its own dedicated tests.
+type memStore struct {
+	mu sync.Mutex
+
+	identityPublic  []byte
+	identityPrivate []byte
+	registrationID  uint32
+
+	preKeys          map[uint32][]byte
+	signedPreKeys    map[uint32][]byte
+	remoteIdentities map[string][]byte
+	sessions         map[string][]byte
+}
+
+func newMemStore(public, private []byte, registrationID uint32) *memStore {
+	return &memStore{
+		identityPublic:   public,
+		identityPrivate:  private,
+		registrationID:   registrationID,
+		preKeys:          map[uint32][]byte{},
+		signedPreKeys:    map[uint32][]byte{},
+		remoteIdentities: map[string][]byte{},
+		sessions:         map[string][]byte{},
+	}
+}
+
+func addrKey(addr libsignal.Address) string {
+	return fmt.Sprintf("%s:%d", addr.Name, addr.DeviceID)
+}
+
+func (m *memStore) GetIdentityKeyPair() ([]byte, []byte, error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	return m.identityPublic, m.identityPrivate, nil
+}
+
+func (m *memStore) GetLocalRegistrationID() (uint32, error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	return m.registrationID, nil
+}
+
+func (m *memStore) SaveIdentity(addr libsignal.Address, key []byte) error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	if key == nil {
+		delete(m.remoteIdentities, addrKey(addr))
+		return nil
+	}
+	m.remoteIdentities[addrKey(addr)] = append([]byte{}, key...)
+	return nil
+}
+
+func (m *memStore) IsTrustedIdentity(addr libsignal.Address, key []byte) (bool, error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	stored, ok := m.remoteIdentities[addrKey(addr)]
+	if !ok {
+		return true, nil
+	}
+	return bytes.Equal(stored, key), nil
+}
+
+func (m *memStore) LoadPreKey(id uint32) ([]byte, bool, error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	v, ok := m.preKeys[id]
+	return v, ok, nil
+}
+
+func (m *memStore) StorePreKey(id uint32, record []byte) error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.preKeys[id] = append([]byte{}, record...)
+	return nil
+}
+
+func (m *memStore) ContainsPreKey(id uint32) bool {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	_, ok := m.preKeys[id]
+	return ok
+}
+
+func (m *memStore) RemovePreKey(id uint32) error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	delete(m.preKeys, id)
+	return nil
+}
+
+func (m *memStore) LoadSignedPreKey(id uint32) ([]byte, bool, error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	v, ok := m.signedPreKeys[id]
+	return v, ok, nil
+}
+
+func (m *memStore) StoreSignedPreKey(id uint32, record []byte) error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.signedPreKeys[id] = append([]byte{}, record...)
+	return nil
+}
+
+func (m *memStore) ContainsSignedPreKey(id uint32) bool {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	_, ok := m.signedPreKeys[id]
+	return ok
+}
+
+func (m *memStore) RemoveSignedPreKey(id uint32) error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	delete(m.signedPreKeys, id)
+	return nil
+}
+
+func (m *memStore) LoadSession(addr libsignal.Address) ([]byte, bool, error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	v, ok := m.sessions[addrKey(addr)]
+	return v, ok, nil
+}
+
+func (m *memStore) GetSubDeviceSessions(name string) ([]uint32, error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	var ids []uint32
+	prefix := name + ":"
+	for k := range m.sessions {
+		if strings.HasPrefix(k, prefix) {
+			if id, err := strconv.ParseUint(strings.TrimPrefix(k, prefix), 10, 32); err == nil {
+				ids = append(ids, uint32(id))
+			}
+		}
+	}
+	return ids, nil
+}
+
+func (m *memStore) StoreSession(addr libsignal.Address, record []byte) error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.sessions[addrKey(addr)] = append([]byte{}, record...)
+	return nil
+}
+
+func (m *memStore) ContainsSession(addr libsignal.Address) bool {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	_, ok := m.sessions[addrKey(addr)]
+	return ok
+}
+
+func (m *memStore) DeleteSession(addr libsignal.Address) error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	delete(m.sessions, addrKey(addr))
+	return nil
+}
+
+func (m *memStore) DeleteAllSessions(name string) (int, error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	prefix := name + ":"
+	count := 0
+	for k := range m.sessions {
+		if strings.HasPrefix(k, prefix) {
+			delete(m.sessions, k)
+			count++
+		}
+	}
+	return count, nil
+}
+
+var _ libsignal.Store = (*memStore)(nil)

e2ee/omemo/libsignal/session.go

@@ -33,7 +33,7 @@ type RemoteBundle struct {
 	PreKeyPublic          []byte // nil if PreKeyID == 0
 	SignedPreKeyID        uint32
 	SignedPreKeyPublic    []byte
-	SignedPreKeySignature []byte
+	SignedPreKeySignature []byte // see SignedPreKeyInfo: pick Signature or SignatureOMEMO to match the version this bundle is processed with
 	IdentityKeyPublic     []byte
 }
 
@@ -41,6 +41,7 @@ type RemoteBundle struct {
 // one remote device.
 type SessionBuilder struct {
 	raw  *C.session_builder
+	addr *cAddress
 }
 
 // NewSessionBuilder creates a session builder for storeCtx's local identity
@@ -48,19 +49,19 @@ type SessionBuilder struct {
 // built by ProcessPreKeyBundle will use for encryption going forward
 // (ProtocolVersionLegacy or ProtocolVersionModern).
 func NewSessionBuilder(ctx *Context, storeCtx *StoreContext, remote Address, version int) (*SessionBuilder, error) {
+	addr := newCAddress(remote)
+
 	var raw *C.session_builder
-	var code C.int
-	withCAddress(remote, func(cAddr *C.signal_protocol_address) {
-		code = C.session_builder_create(&raw, storeCtx.raw, cAddr, ctx.raw)
-	})
-	if code != C.SG_SUCCESS {
+	if code := C.session_builder_create(&raw, storeCtx.raw, addr.ptr, ctx.raw); code != C.SG_SUCCESS {
+		addr.free()
 		return nil, newError("session_builder_create", int(code))
 	}
 	if err := setBuilderVersion(raw, version); err != nil {
 		C.session_builder_free(raw)
+		addr.free()
 		return nil, err
 	}
-	return &SessionBuilder{raw: raw}, nil
+	return &SessionBuilder{raw: raw, addr: addr}, nil
 }
 
 // Close frees the session builder. Do not use it afterward.
@@ -69,6 +70,7 @@ func (b *SessionBuilder) Close() {
 		C.session_builder_free(b.raw)
 		b.raw = nil
 	}
+	b.addr.free()
 }
 
 // ProcessPreKeyBundle establishes a session from a remote peer's fetched

e2ee/omemo/libsignal/session_test.go

@@ -0,0 +1,221 @@
+package libsignal_test
+
+// This is a self-consistency round trip, not a transcription of published
+// X3DH/Double Ratchet test vectors: two simulated parties generate real
+// keys through the cgo-linked library, establish a session via X3DH, and
+// exchange messages, verifying the plaintext survives encrypt-on-one-side
+// / decrypt-on-the-other for both protocol versions. That's the property
+// that actually matters for this binding (it interoperates with the real
+// library correctly), and avoids the risk of mistranscribing external
+// vectors for a library/wire-format combination with no independently
+// published vectors readily available.
+
+import (
+	"bytes"
+	"testing"
+
+	"dev.narayana.im/narayana/telegabber/e2ee/omemo/libsignal"
+)
+
+// identity bundles up one simulated party's generated keys and store.
+type identity struct {
+	ctx      *libsignal.Context
+	storeCtx *libsignal.StoreContext
+
+	registrationID     uint32
+	identityPublicKey  []byte
+	signedPreKeyID     uint32
+	signedPreKeyRecord []byte
+	preKeyID           uint32
+	preKeyRecord       []byte
+}
+
+func newIdentity(t *testing.T) *identity {
+	t.Helper()
+
+	ctx, err := libsignal.NewContext()
+	if err != nil {
+		t.Fatalf("NewContext: %v", err)
+	}
+	t.Cleanup(ctx.Close)
+
+	idKeyPair, err := libsignal.GenerateIdentityKeyPair(ctx)
+	if err != nil {
+		t.Fatalf("GenerateIdentityKeyPair: %v", err)
+	}
+	pub, priv, err := libsignal.SplitIdentityKeyPair(ctx, idKeyPair.Record)
+	if err != nil {
+		t.Fatalf("SplitIdentityKeyPair: %v", err)
+	}
+	regID, err := libsignal.GenerateRegistrationID(ctx)
+	if err != nil {
+		t.Fatalf("GenerateRegistrationID: %v", err)
+	}
+
+	store := newMemStore(pub, priv, regID)
+
+	signedPreKey, err := libsignal.GenerateSignedPreKey(ctx, idKeyPair, 1)
+	if err != nil {
+		t.Fatalf("GenerateSignedPreKey: %v", err)
+	}
+	if err := store.StoreSignedPreKey(signedPreKey.ID, signedPreKey.Record); err != nil {
+		t.Fatalf("StoreSignedPreKey: %v", err)
+	}
+
+	preKeys, err := libsignal.GeneratePreKeys(ctx, 1, 1)
+	if err != nil {
+		t.Fatalf("GeneratePreKeys: %v", err)
+	}
+	if len(preKeys) != 1 {
+		t.Fatalf("expected 1 prekey, got %d", len(preKeys))
+	}
+	if err := store.StorePreKey(preKeys[0].ID, preKeys[0].Record); err != nil {
+		t.Fatalf("StorePreKey: %v", err)
+	}
+
+	storeCtx, err := libsignal.NewStoreContext(ctx, store)
+	if err != nil {
+		t.Fatalf("NewStoreContext: %v", err)
+	}
+	t.Cleanup(storeCtx.Close)
+
+	return &identity{
+		ctx:                ctx,
+		storeCtx:           storeCtx,
+		registrationID:     regID,
+		identityPublicKey:  pub,
+		signedPreKeyID:     signedPreKey.ID,
+		signedPreKeyRecord: signedPreKey.Record,
+		preKeyID:           preKeys[0].ID,
+		preKeyRecord:       preKeys[0].Record,
+	}
+}
+
+// bundle returns id's own device bundle, as it would be published in an
+// XEP-0384 bundle - i.e. what a remote peer's SessionBuilder needs to
+// start a session with id via X3DH.
+//
+// Which signature to include depends on the recipient's protocol version:
+// session_builder_process_pre_key_bundle re-serializes the signed prekey's
+// public key via ec_public_key_serialize (version < 4, the 33-byte
+// DJB_TYPE-prefixed legacy form) or ec_public_key_serialize_omemo
+// (version >= 4, the raw 32-byte Montgomery form) before verifying it
+// against whichever signature was supplied - confirmed by reading
+// session_builder.c. Supplying the wrong one of the two signatures
+// libomemo-c's signed-prekey generation always computes fails with
+// SG_ERR_INVALID_KEY, not a version-related error, which is what actually
+// surfaced this while writing this test.
+func (id *identity) bundle(t *testing.T, deviceID uint32, version int) libsignal.RemoteBundle {
+	t.Helper()
+
+	spkInfo, err := libsignal.DecodeSignedPreKey(id.ctx, id.signedPreKeyRecord)
+	if err != nil {
+		t.Fatalf("DecodeSignedPreKey: %v", err)
+	}
+	pkInfo, err := libsignal.DecodePreKey(id.ctx, id.preKeyRecord)
+	if err != nil {
+		t.Fatalf("DecodePreKey: %v", err)
+	}
+
+	signature := spkInfo.Signature
+	if version >= libsignal.ProtocolVersionModern {
+		signature = spkInfo.SignatureOMEMO
+	}
+
+	return libsignal.RemoteBundle{
+		RegistrationID:        id.registrationID,
+		DeviceID:              deviceID,
+		PreKeyID:              pkInfo.ID,
+		PreKeyPublic:          pkInfo.PublicKey,
+		SignedPreKeyID:        spkInfo.ID,
+		SignedPreKeyPublic:    spkInfo.PublicKey,
+		SignedPreKeySignature: signature,
+		IdentityKeyPublic:     id.identityPublicKey,
+	}
+}
+
+func TestSessionRoundTrip(t *testing.T) {
+	for _, tc := range []struct {
+		name    string
+		version int
+		omemo   bool
+	}{
+		{"legacy", libsignal.ProtocolVersionLegacy, false},
+		{"modern", libsignal.ProtocolVersionModern, true},
+	} {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			if tc.version == libsignal.ProtocolVersionModern && !libsignal.ModernOMEMOSupported {
+				t.Skip("modern OMEMO not supported by this build (signal_legacy tag)")
+			}
+
+			alice := newIdentity(t)
+			bob := newIdentity(t)
+
+			aliceAddr := libsignal.Address{Name: "alice@example.com", DeviceID: 1}
+			bobAddr := libsignal.Address{Name: "bob@example.com", DeviceID: 1}
+
+			// Alice fetches Bob's bundle and establishes a session (X3DH).
+			builder, err := libsignal.NewSessionBuilder(alice.ctx, alice.storeCtx, bobAddr, tc.version)
+			if err != nil {
+				t.Fatalf("NewSessionBuilder: %v", err)
+			}
+			defer builder.Close()
+			if err := builder.ProcessPreKeyBundle(alice.ctx, bob.bundle(t, bobAddr.DeviceID, tc.version)); err != nil {
+				t.Fatalf("ProcessPreKeyBundle: %v", err)
+			}
+
+			aliceCipher, err := libsignal.NewSessionCipher(alice.ctx, alice.storeCtx, bobAddr, tc.version)
+			if err != nil {
+				t.Fatalf("NewSessionCipher (alice): %v", err)
+			}
+			defer aliceCipher.Close()
+
+			plaintext1 := []byte("hello bob")
+			ct1, err := aliceCipher.Encrypt(plaintext1)
+			if err != nil {
+				t.Fatalf("Encrypt (alice->bob): %v", err)
+			}
+			if ct1.Type != libsignal.CiphertextPreKeyType {
+				t.Fatalf("expected first message to be a prekey message, got type %d", ct1.Type)
+			}
+
+			// Bob decrypts Alice's first message - this establishes Bob's
+			// side of the session as a side effect; no ProcessPreKeyBundle
+			// needed on his end, only a bare SessionCipher.
+			bobCipher, err := libsignal.NewSessionCipher(bob.ctx, bob.storeCtx, aliceAddr, tc.version)
+			if err != nil {
+				t.Fatalf("NewSessionCipher (bob): %v", err)
+			}
+			defer bobCipher.Close()
+
+			got1, err := bobCipher.Decrypt(bob.ctx, *ct1, tc.omemo, bob.registrationID)
+			if err != nil {
+				t.Fatalf("Decrypt (bob<-alice): %v", err)
+			}
+			if !bytes.Equal(got1, plaintext1) {
+				t.Fatalf("round trip mismatch: got %q, want %q", got1, plaintext1)
+			}
+
+			// Bob replies; Alice decrypts using her already-established
+			// session (an ordinary ratchet message this time, not a
+			// prekey message).
+			plaintext2 := []byte("hi alice")
+			ct2, err := bobCipher.Encrypt(plaintext2)
+			if err != nil {
+				t.Fatalf("Encrypt (bob->alice): %v", err)
+			}
+			if ct2.Type != libsignal.CiphertextSignalType {
+				t.Fatalf("expected reply to be an ordinary ratchet message, got type %d", ct2.Type)
+			}
+
+			got2, err := aliceCipher.Decrypt(alice.ctx, *ct2, tc.omemo, alice.registrationID)
+			if err != nil {
+				t.Fatalf("Decrypt (alice<-bob): %v", err)
+			}
+			if !bytes.Equal(got2, plaintext2) {
+				t.Fatalf("round trip mismatch: got %q, want %q", got2, plaintext2)
+			}
+		})
+	}
+}

e2ee/store/badgerstore/store_test.go

@@ -0,0 +1,259 @@
+package badgerstore_test
+
+import (
+	"testing"
+
+	"dev.narayana.im/narayana/telegabber/e2ee/omemo/libsignal"
+	"dev.narayana.im/narayana/telegabber/e2ee/store/badgerstore"
+)
+
+func openTestStore(t *testing.T) *badgerstore.Store {
+	t.Helper()
+
+	db, err := badgerstore.Open(t.TempDir(), nil)
+	if err != nil {
+		t.Fatalf("Open: %v", err)
+	}
+	t.Cleanup(func() { db.Close() })
+
+	ctx, err := libsignal.NewContext()
+	if err != nil {
+		t.Fatalf("NewContext: %v", err)
+	}
+	t.Cleanup(ctx.Close)
+
+	return db.Store(ctx, "testlogin", "1234@example.com")
+}
+
+func TestIdentityKeyPairRoundTrip(t *testing.T) {
+	store := openTestStore(t)
+
+	if has, err := store.HasIdentityKeyPair(); err != nil || has {
+		t.Fatalf("expected no identity key pair yet, has=%v err=%v", has, err)
+	}
+
+	// A fresh context is fine for generation - the record format doesn't
+	// depend on which *Context instance produced it.
+	ctx, err := libsignal.NewContext()
+	if err != nil {
+		t.Fatalf("NewContext: %v", err)
+	}
+	defer ctx.Close()
+
+	idKeyPair, err := libsignal.GenerateIdentityKeyPair(ctx)
+	if err != nil {
+		t.Fatalf("GenerateIdentityKeyPair: %v", err)
+	}
+	wantPublic, _, err := libsignal.SplitIdentityKeyPair(ctx, idKeyPair.Record)
+	if err != nil {
+		t.Fatalf("SplitIdentityKeyPair: %v", err)
+	}
+
+	if err := store.SaveIdentityKeyPair(idKeyPair.Record); err != nil {
+		t.Fatalf("SaveIdentityKeyPair: %v", err)
+	}
+	if has, err := store.HasIdentityKeyPair(); err != nil || !has {
+		t.Fatalf("expected identity key pair to be saved, has=%v err=%v", has, err)
+	}
+
+	gotPublic, _, err := store.GetIdentityKeyPair()
+	if err != nil {
+		t.Fatalf("GetIdentityKeyPair: %v", err)
+	}
+	if string(gotPublic) != string(wantPublic) {
+		t.Fatalf("public key mismatch after round trip")
+	}
+}
+
+func TestRegistrationIDRoundTrip(t *testing.T) {
+	store := openTestStore(t)
+
+	if err := store.SaveRegistrationID(4242); err != nil {
+		t.Fatalf("SaveRegistrationID: %v", err)
+	}
+	got, err := store.GetLocalRegistrationID()
+	if err != nil {
+		t.Fatalf("GetLocalRegistrationID: %v", err)
+	}
+	if got != 4242 {
+		t.Fatalf("got %d, want 4242", got)
+	}
+}
+
+func TestPreKeyRoundTrip(t *testing.T) {
+	store := openTestStore(t)
+
+	if store.ContainsPreKey(7) {
+		t.Fatalf("expected prekey 7 to not exist yet")
+	}
+	if err := store.StorePreKey(7, []byte("prekey-record")); err != nil {
+		t.Fatalf("StorePreKey: %v", err)
+	}
+	if !store.ContainsPreKey(7) {
+		t.Fatalf("expected prekey 7 to exist")
+	}
+	record, found, err := store.LoadPreKey(7)
+	if err != nil || !found {
+		t.Fatalf("LoadPreKey: found=%v err=%v", found, err)
+	}
+	if string(record) != "prekey-record" {
+		t.Fatalf("got %q", record)
+	}
+	if err := store.RemovePreKey(7); err != nil {
+		t.Fatalf("RemovePreKey: %v", err)
+	}
+	if store.ContainsPreKey(7) {
+		t.Fatalf("expected prekey 7 to be removed")
+	}
+}
+
+func TestSignedPreKeyRoundTrip(t *testing.T) {
+	store := openTestStore(t)
+
+	if err := store.StoreSignedPreKey(3, []byte("signed-record")); err != nil {
+		t.Fatalf("StoreSignedPreKey: %v", err)
+	}
+	if !store.ContainsSignedPreKey(3) {
+		t.Fatalf("expected signed prekey 3 to exist")
+	}
+	record, found, err := store.LoadSignedPreKey(3)
+	if err != nil || !found || string(record) != "signed-record" {
+		t.Fatalf("LoadSignedPreKey: record=%q found=%v err=%v", record, found, err)
+	}
+	if err := store.RemoveSignedPreKey(3); err != nil {
+		t.Fatalf("RemoveSignedPreKey: %v", err)
+	}
+	if store.ContainsSignedPreKey(3) {
+		t.Fatalf("expected signed prekey 3 to be removed")
+	}
+}
+
+func TestNextPreKeyIDs(t *testing.T) {
+	store := openTestStore(t)
+
+	start1, err := store.NextPreKeyIDs(10)
+	if err != nil {
+		t.Fatalf("NextPreKeyIDs: %v", err)
+	}
+	if start1 != 0 {
+		t.Fatalf("expected first block to start at 0, got %d", start1)
+	}
+	start2, err := store.NextPreKeyIDs(5)
+	if err != nil {
+		t.Fatalf("NextPreKeyIDs: %v", err)
+	}
+	if start2 != 10 {
+		t.Fatalf("expected second block to start at 10, got %d", start2)
+	}
+}
+
+func TestSessionRoundTrip(t *testing.T) {
+	store := openTestStore(t)
+	addr := libsignal.Address{Name: "bob@example.com", DeviceID: 1}
+
+	if store.ContainsSession(addr) {
+		t.Fatalf("expected no session yet")
+	}
+	if err := store.StoreSession(addr, []byte("session-record")); err != nil {
+		t.Fatalf("StoreSession: %v", err)
+	}
+	if !store.ContainsSession(addr) {
+		t.Fatalf("expected session to exist")
+	}
+
+	record, found, err := store.LoadSession(addr)
+	if err != nil || !found || string(record) != "session-record" {
+		t.Fatalf("LoadSession: record=%q found=%v err=%v", record, found, err)
+	}
+
+	otherAddr := libsignal.Address{Name: "bob@example.com", DeviceID: 2}
+	if err := store.StoreSession(otherAddr, []byte("session-record-2")); err != nil {
+		t.Fatalf("StoreSession (device 2): %v", err)
+	}
+
+	ids, err := store.GetSubDeviceSessions("bob@example.com")
+	if err != nil {
+		t.Fatalf("GetSubDeviceSessions: %v", err)
+	}
+	if len(ids) != 2 {
+		t.Fatalf("expected 2 device sessions, got %d (%v)", len(ids), ids)
+	}
+
+	count, err := store.DeleteAllSessions("bob@example.com")
+	if err != nil {
+		t.Fatalf("DeleteAllSessions: %v", err)
+	}
+	if count != 2 {
+		t.Fatalf("expected to delete 2 sessions, deleted %d", count)
+	}
+	if store.ContainsSession(addr) || store.ContainsSession(otherAddr) {
+		t.Fatalf("expected both sessions to be gone")
+	}
+}
+
+func TestIsTrustedIdentityTOFU(t *testing.T) {
+	store := openTestStore(t)
+	addr := libsignal.Address{Name: "bob@example.com", DeviceID: 1}
+
+	key1 := []byte("identity-key-v1")
+	trusted, err := store.IsTrustedIdentity(addr, key1)
+	if err != nil {
+		t.Fatalf("IsTrustedIdentity (unseen): %v", err)
+	}
+	if !trusted {
+		t.Fatalf("expected an unseen identity to be trusted (TOFU)")
+	}
+
+	if err := store.SaveIdentity(addr, key1); err != nil {
+		t.Fatalf("SaveIdentity: %v", err)
+	}
+
+	trusted, err = store.IsTrustedIdentity(addr, key1)
+	if err != nil || !trusted {
+		t.Fatalf("expected the same key to remain trusted: trusted=%v err=%v", trusted, err)
+	}
+
+	key2 := []byte("identity-key-v2-different")
+	trusted, err = store.IsTrustedIdentity(addr, key2)
+	if err != nil {
+		t.Fatalf("IsTrustedIdentity (mismatched): %v", err)
+	}
+	if trusted {
+		t.Fatalf("expected a mismatched key to be untrusted, not silently re-trusted")
+	}
+}
+
+func TestEnabledFlag(t *testing.T) {
+	store := openTestStore(t)
+
+	enabled, err := store.Enabled()
+	if err != nil {
+		t.Fatalf("Enabled (default): %v", err)
+	}
+	if enabled {
+		t.Fatalf("expected OMEMO to be off by default")
+	}
+
+	if err := store.SetEnabled(true); err != nil {
+		t.Fatalf("SetEnabled: %v", err)
+	}
+	enabled, err = store.Enabled()
+	if err != nil || !enabled {
+		t.Fatalf("expected OMEMO to be enabled after SetEnabled(true): enabled=%v err=%v", enabled, err)
+	}
+}
+
+func TestRemoteDeviceListCache(t *testing.T) {
+	store := openTestStore(t)
+
+	if _, found, err := store.RemoteDeviceListCache("bob@example.com"); err != nil || found {
+		t.Fatalf("expected no cached device list yet: found=%v err=%v", found, err)
+	}
+	if err := store.SaveRemoteDeviceListCache("bob@example.com", []byte("<devices/>")); err != nil {
+		t.Fatalf("SaveRemoteDeviceListCache: %v", err)
+	}
+	doc, found, err := store.RemoteDeviceListCache("bob@example.com")
+	if err != nil || !found || string(doc) != "<devices/>" {
+		t.Fatalf("RemoteDeviceListCache: doc=%q found=%v err=%v", doc, found, err)
+	}
+}