slidge-whatsapp 0.2.5__cp311-cp311-manylinux_2_36_aarch64.whl → 0.3.0__cp311-cp311-manylinux_2_36_aarch64.whl

slidge_whatsapp/vendor/github.com/vektah/gqlparser/v2/ast/value.go

@@ -0,0 +1,122 @@
+package ast
+
+import (
+	"fmt"
+	"strconv"
+	"strings"
+)
+
+type ValueKind int
+
+const (
+	Variable ValueKind = iota
+	IntValue
+	FloatValue
+	StringValue
+	BlockValue
+	BooleanValue
+	NullValue
+	EnumValue
+	ListValue
+	ObjectValue
+)
+
+type Value struct {
+	Raw      string
+	Children ChildValueList
+	Kind     ValueKind
+	Position *Position `dump:"-" json:"-"`
+	Comment  *CommentGroup
+
+	// Require validation
+	Definition         *Definition
+	VariableDefinition *VariableDefinition
+	ExpectedType       *Type
+}
+
+type ChildValue struct {
+	Name     string
+	Value    *Value
+	Position *Position `dump:"-" json:"-"`
+	Comment  *CommentGroup
+}
+
+func (v *Value) Value(vars map[string]interface{}) (interface{}, error) {
+	if v == nil {
+		return nil, nil
+	}
+	switch v.Kind {
+	case Variable:
+		if value, ok := vars[v.Raw]; ok {
+			return value, nil
+		}
+		if v.VariableDefinition != nil && v.VariableDefinition.DefaultValue != nil {
+			return v.VariableDefinition.DefaultValue.Value(vars)
+		}
+		return nil, nil
+	case IntValue:
+		return strconv.ParseInt(v.Raw, 10, 64)
+	case FloatValue:
+		return strconv.ParseFloat(v.Raw, 64)
+	case StringValue, BlockValue, EnumValue:
+		return v.Raw, nil
+	case BooleanValue:
+		return strconv.ParseBool(v.Raw)
+	case NullValue:
+		return nil, nil
+	case ListValue:
+		var val []interface{}
+		for _, elem := range v.Children {
+			elemVal, err := elem.Value.Value(vars)
+			if err != nil {
+				return val, err
+			}
+			val = append(val, elemVal)
+		}
+		return val, nil
+	case ObjectValue:
+		val := map[string]interface{}{}
+		for _, elem := range v.Children {
+			elemVal, err := elem.Value.Value(vars)
+			if err != nil {
+				return val, err
+			}
+			val[elem.Name] = elemVal
+		}
+		return val, nil
+	default:
+		panic(fmt.Errorf("unknown value kind %d", v.Kind))
+	}
+}
+
+func (v *Value) String() string {
+	if v == nil {
+		return "<nil>"
+	}
+	switch v.Kind {
+	case Variable:
+		return "$" + v.Raw
+	case IntValue, FloatValue, EnumValue, BooleanValue, NullValue:
+		return v.Raw
+	case StringValue, BlockValue:
+		return strconv.Quote(v.Raw)
+	case ListValue:
+		var val []string
+		for _, elem := range v.Children {
+			val = append(val, elem.Value.String())
+		}
+		return "[" + strings.Join(val, ",") + "]"
+	case ObjectValue:
+		var val []string
+		for _, elem := range v.Children {
+			val = append(val, elem.Name+":"+elem.Value.String())
+		}
+		return "{" + strings.Join(val, ",") + "}"
+	default:
+		panic(fmt.Errorf("unknown value kind %d", v.Kind))
+	}
+}
+
+func (v *Value) Dump() string {
+	return v.String()
+}

slidge_whatsapp/vendor/go.mau.fi/libsignal/groups/GroupCipher.go

@@ -1,6 +1,7 @@
 package groups
 
 import (
+	"context"
 	"fmt"
 
 	"go.mau.fi/libsignal/cipher"
@@ -34,9 +35,12 @@ type GroupCipher struct {
 }
 
 // Encrypt will take the given message in bytes and return encrypted bytes.
-func (c *GroupCipher) Encrypt(plaintext []byte) (protocol.GroupCiphertextMessage, error) {
+func (c *GroupCipher) Encrypt(ctx context.Context, plaintext []byte) (protocol.GroupCiphertextMessage, error) {
 	// Load the sender key based on id from our store.
-	keyRecord := c.senderKeyStore.LoadSenderKey(c.senderKeyID)
+	keyRecord, err := c.senderKeyStore.LoadSenderKey(ctx, c.senderKeyID)
+	if err != nil {
+		return nil, err
+	}
 	senderKeyState, err := keyRecord.SenderKeyState()
 	if err != nil {
 		return nil, err
@@ -63,15 +67,20 @@ func (c *GroupCipher) Encrypt(plaintext []byte) (protocol.GroupCiphertextMessage
 	)
 
 	senderKeyState.SetSenderChainKey(senderKeyState.SenderChainKey().Next())
-	c.senderKeyStore.StoreSenderKey(c.senderKeyID, keyRecord)
+	if err := c.senderKeyStore.StoreSenderKey(ctx, c.senderKeyID, keyRecord); err != nil {
+		return nil, err
+	}
 
 	return senderKeyMessage, nil
 }
 
 // Decrypt decrypts the given message using an existing session that
 // is stored in the senderKey store.
-func (c *GroupCipher) Decrypt(senderKeyMessage *protocol.SenderKeyMessage) ([]byte, error) {
-	keyRecord := c.senderKeyStore.LoadSenderKey(c.senderKeyID)
+func (c *GroupCipher) Decrypt(ctx context.Context, senderKeyMessage *protocol.SenderKeyMessage) ([]byte, error) {
+	keyRecord, err := c.senderKeyStore.LoadSenderKey(ctx, c.senderKeyID)
+	if err != nil {
+		return nil, err
+	}
 
 	if keyRecord.IsEmpty() {
 		return nil, fmt.Errorf("%w for %s in %s", signalerror.ErrNoSenderKeyForUser, c.senderKeyID.Sender().String(), c.senderKeyID.GroupID())
@@ -101,7 +110,9 @@ func (c *GroupCipher) Decrypt(senderKeyMessage *protocol.SenderKeyMessage) ([]by
 	}
 
 	// Store the sender key by id.
-	c.senderKeyStore.StoreSenderKey(c.senderKeyID, keyRecord)
+	if err := c.senderKeyStore.StoreSenderKey(ctx, c.senderKeyID, keyRecord); err != nil {
+		return nil, err
+	}
 
 	return plaintext, nil
 }

slidge_whatsapp/vendor/go.mau.fi/libsignal/groups/GroupSessionBuilder.go

@@ -9,6 +9,8 @@
 package groups
 
 import (
+	"context"
+
 	"go.mau.fi/libsignal/groups/state/record"
 	"go.mau.fi/libsignal/groups/state/store"
 	"go.mau.fi/libsignal/protocol"
@@ -34,21 +36,27 @@ type SessionBuilder struct {
 
 // Process will process an incoming group message and set up the corresponding
 // session for it.
-func (b *SessionBuilder) Process(senderKeyName *protocol.SenderKeyName,
-	msg *protocol.SenderKeyDistributionMessage) {
+func (b *SessionBuilder) Process(ctx context.Context, senderKeyName *protocol.SenderKeyName,
+	msg *protocol.SenderKeyDistributionMessage) error {
 
-	senderKeyRecord := b.senderKeyStore.LoadSenderKey(senderKeyName)
+	senderKeyRecord, err := b.senderKeyStore.LoadSenderKey(ctx, senderKeyName)
+	if err != nil {
+		return err
+	}
 	if senderKeyRecord == nil {
 		senderKeyRecord = record.NewSenderKey(b.serializer.SenderKeyRecord, b.serializer.SenderKeyState)
 	}
 	senderKeyRecord.AddSenderKeyState(msg.ID(), msg.Iteration(), msg.ChainKey(), msg.SignatureKey())
-	b.senderKeyStore.StoreSenderKey(senderKeyName, senderKeyRecord)
+	return b.senderKeyStore.StoreSenderKey(ctx, senderKeyName, senderKeyRecord)
 }
 
 // Create will create a new group session for the given name.
-func (b *SessionBuilder) Create(senderKeyName *protocol.SenderKeyName) (*protocol.SenderKeyDistributionMessage, error) {
+func (b *SessionBuilder) Create(ctx context.Context, senderKeyName *protocol.SenderKeyName) (*protocol.SenderKeyDistributionMessage, error) {
 	// Load the senderkey by name
-	senderKeyRecord := b.senderKeyStore.LoadSenderKey(senderKeyName)
+	senderKeyRecord, err := b.senderKeyStore.LoadSenderKey(ctx, senderKeyName)
+	if err != nil {
+		return nil, err
+	}
 
 	// If the record is empty, generate new keys.
 	if senderKeyRecord == nil || senderKeyRecord.IsEmpty() {
@@ -62,7 +70,9 @@ func (b *SessionBuilder) Create(senderKeyName *protocol.SenderKeyName) (*protoco
 			keyhelper.GenerateSenderKey(),
 			signingKey,
 		)
-		b.senderKeyStore.StoreSenderKey(senderKeyName, senderKeyRecord)
+		if err := b.senderKeyStore.StoreSenderKey(ctx, senderKeyName, senderKeyRecord); err != nil {
+			return nil, err
+		}
 	}
 
 	// Get the senderkey state.

slidge_whatsapp/vendor/go.mau.fi/libsignal/groups/state/store/SenderKeyStore.go

@@ -1,11 +1,13 @@
 package store
 
 import (
+	"context"
+
 	"go.mau.fi/libsignal/groups/state/record"
 	"go.mau.fi/libsignal/protocol"
 )
 
 type SenderKey interface {
-	StoreSenderKey(senderKeyName *protocol.SenderKeyName, keyRecord *record.SenderKey)
-	LoadSenderKey(senderKeyName *protocol.SenderKeyName) *record.SenderKey
+	StoreSenderKey(ctx context.Context, senderKeyName *protocol.SenderKeyName, keyRecord *record.SenderKey) error
+	LoadSenderKey(ctx context.Context, senderKeyName *protocol.SenderKeyName) (*record.SenderKey, error)
 }

slidge_whatsapp/vendor/go.mau.fi/libsignal/keys/chain/ChainKey.go

@@ -4,6 +4,7 @@ package chain
 import (
 	"crypto/hmac"
 	"crypto/sha256"
+
 	"go.mau.fi/libsignal/kdf"
 	"go.mau.fi/libsignal/keys/message"
 )

slidge_whatsapp/vendor/go.mau.fi/libsignal/keys/identity/IdentityKey.go

@@ -4,6 +4,7 @@ package identity
 
 import (
 	"encoding/hex"
+
 	"go.mau.fi/libsignal/ecc"
 )
 

slidge_whatsapp/vendor/go.mau.fi/libsignal/logger/DefaultLogger.go

@@ -77,9 +77,9 @@ func (d *defaultLogger) Error(caller, msg string) {
 // allows granular logging of different go files.
 //
 // Example:
-//   logger.Configure("RootKey.go,Curve.go")
-//   logger.Configure("all")
 //
+//	logger.Configure("RootKey.go,Curve.go")
+//	logger.Configure("all")
 func (d *defaultLogger) Configure(settings string) {
 	d.namespaces = strings.Split(settings, ",")
 }

slidge_whatsapp/vendor/go.mau.fi/libsignal/serialize/FingerprintProtocol.pb.go

@@ -9,10 +9,11 @@
 package serialize
 
 import (
-	protoreflect "google.golang.org/protobuf/reflect/protoreflect"
-	protoimpl "google.golang.org/protobuf/runtime/protoimpl"
 	reflect "reflect"
 	sync "sync"
+
+	protoreflect "google.golang.org/protobuf/reflect/protoreflect"
+	protoimpl "google.golang.org/protobuf/runtime/protoimpl"
 )
 
 const (

slidge_whatsapp/vendor/go.mau.fi/libsignal/serialize/LocalStorageProtocol.pb.go

@@ -9,10 +9,11 @@
 package serialize
 
 import (
-	protoreflect "google.golang.org/protobuf/reflect/protoreflect"
-	protoimpl "google.golang.org/protobuf/runtime/protoimpl"
 	reflect "reflect"
 	sync "sync"
+
+	protoreflect "google.golang.org/protobuf/reflect/protoreflect"
+	protoimpl "google.golang.org/protobuf/runtime/protoimpl"
 )
 
 const (

slidge_whatsapp/vendor/go.mau.fi/libsignal/serialize/WhisperTextProtocol.pb.go

@@ -9,10 +9,11 @@
 package serialize
 
 import (
-	protoreflect "google.golang.org/protobuf/reflect/protoreflect"
-	protoimpl "google.golang.org/protobuf/runtime/protoimpl"
 	reflect "reflect"
 	sync "sync"
+
+	protoreflect "google.golang.org/protobuf/reflect/protoreflect"
+	protoimpl "google.golang.org/protobuf/runtime/protoimpl"
 )
 
 const (

slidge_whatsapp/vendor/go.mau.fi/libsignal/session/Session.go

@@ -2,6 +2,7 @@
 package session
 
 import (
+	"context"
 	"fmt"
 
 	"go.mau.fi/libsignal/ecc"
@@ -56,9 +57,9 @@ func NewBuilderFromSignal(signalStore store.SignalProtocol,
 // used to encrypt/decrypt messages in that session.
 //
 // Sessions are built from one of three different vectors:
-//   * PreKeyBundle retrieved from a server.
-//   * PreKeySignalMessage received from a client.
-//   * KeyExchangeMessage sent to or received from a client.
+//   - PreKeyBundle retrieved from a server.
+//   - PreKeySignalMessage received from a client.
+//   - KeyExchangeMessage sent to or received from a client.
 //
 // Sessions are constructed per recipientId + deviceId tuple.
 // Remote logical users are identified by their recipientId,
@@ -75,22 +76,28 @@ type Builder struct {
 
 // Process builds a new session from a session record and pre
 // key signal message.
-func (b *Builder) Process(sessionRecord *record.Session, message *protocol.PreKeySignalMessage) (unsignedPreKeyID *optional.Uint32, err error) {
+func (b *Builder) Process(ctx context.Context, sessionRecord *record.Session, message *protocol.PreKeySignalMessage) (unsignedPreKeyID *optional.Uint32, err error) {
 
 	// Check to see if the keys are trusted.
 	theirIdentityKey := message.IdentityKey()
-	if !(b.identityKeyStore.IsTrustedIdentity(b.remoteAddress, theirIdentityKey)) {
+	trusted, err := b.identityKeyStore.IsTrustedIdentity(ctx, b.remoteAddress, theirIdentityKey)
+	if err != nil {
+		return nil, err
+	}
+	if !trusted {
 		return nil, signalerror.ErrUntrustedIdentity
 	}
 
 	// Use version 3 of the signal/axolotl protocol.
-	unsignedPreKeyID, err = b.processV3(sessionRecord, message)
+	unsignedPreKeyID, err = b.processV3(ctx, sessionRecord, message)
 	if err != nil {
 		return nil, err
 	}
 
 	// Save the identity key to our identity store.
-	b.identityKeyStore.SaveIdentity(b.remoteAddress, theirIdentityKey)
+	if err := b.identityKeyStore.SaveIdentity(ctx, b.remoteAddress, theirIdentityKey); err != nil {
+		return nil, err
+	}
 
 	// Return the unsignedPreKeyID
 	return unsignedPreKeyID, nil
@@ -99,7 +106,7 @@ func (b *Builder) Process(sessionRecord *record.Session, message *protocol.PreKe
 // ProcessV3 builds a new session from a session record and pre key
 // signal message. After a session is constructed in this way, the embedded
 // SignalMessage can be decrypted.
-func (b *Builder) processV3(sessionRecord *record.Session,
+func (b *Builder) processV3(ctx context.Context, sessionRecord *record.Session,
 	message *protocol.PreKeySignalMessage) (unsignedPreKeyID *optional.Uint32, err error) {
 
 	logger.Debug("Processing message with PreKeyID: ", message.PreKeyID())
@@ -114,7 +121,10 @@ func (b *Builder) processV3(sessionRecord *record.Session,
 	}
 
 	// Load our signed prekey from our signed prekey store.
-	ourSignedPreKeyRecord := b.signedPreKeyStore.LoadSignedPreKey(message.SignedPreKeyID())
+	ourSignedPreKeyRecord, err := b.signedPreKeyStore.LoadSignedPreKey(ctx, message.SignedPreKeyID())
+	if err != nil {
+		return nil, err
+	}
 	if ourSignedPreKeyRecord == nil {
 		return nil, fmt.Errorf("%w with ID %d", signalerror.ErrNoSignedPreKey, message.SignedPreKeyID())
 	}
@@ -131,7 +141,10 @@ func (b *Builder) processV3(sessionRecord *record.Session,
 	// Set our one time pre key with the one from our prekey store
 	// if the message contains a valid pre key id
 	if !message.PreKeyID().IsEmpty {
-		oneTimePreKey := b.preKeyStore.LoadPreKey(message.PreKeyID().Value)
+		oneTimePreKey, err := b.preKeyStore.LoadPreKey(ctx, message.PreKeyID().Value)
+		if err != nil {
+			return nil, err
+		}
 		if oneTimePreKey == nil {
 			return nil, fmt.Errorf("%w with ID %d", signalerror.ErrNoOneTimeKeyFound, message.PreKeyID().Value)
 		}
@@ -158,7 +171,7 @@ func (b *Builder) processV3(sessionRecord *record.Session,
 	sessionState.SetRootKey(derivedKeys.RootKey)
 
 	// Set the session's registration ids and base key
-	sessionState.SetLocalRegistrationID(b.identityKeyStore.GetLocalRegistrationId())
+	sessionState.SetLocalRegistrationID(b.identityKeyStore.GetLocalRegistrationID())
 	sessionState.SetRemoteRegistrationID(message.RegistrationID())
 	sessionState.SetSenderBaseKey(message.BaseKey().Serialize())
 
@@ -171,9 +184,13 @@ func (b *Builder) processV3(sessionRecord *record.Session,
 
 // ProcessBundle builds a new session from a PreKeyBundle retrieved
 // from a server.
-func (b *Builder) ProcessBundle(preKey *prekey.Bundle) error {
+func (b *Builder) ProcessBundle(ctx context.Context, preKey *prekey.Bundle) error {
 	// Check to see if the keys are trusted.
-	if !(b.identityKeyStore.IsTrustedIdentity(b.remoteAddress, preKey.IdentityKey())) {
+	trusted, err := b.identityKeyStore.IsTrustedIdentity(ctx, b.remoteAddress, preKey.IdentityKey())
+	if err != nil {
+		return err
+	}
+	if !trusted {
 		return signalerror.ErrUntrustedIdentity
 	}
 
@@ -191,7 +208,10 @@ func (b *Builder) ProcessBundle(preKey *prekey.Bundle) error {
 	}
 
 	// Load our session and generate keys.
-	sessionRecord := b.sessionStore.LoadSession(b.remoteAddress)
+	sessionRecord, err := b.sessionStore.LoadSession(ctx, b.remoteAddress)
+	if err != nil {
+		return err
+	}
 	if sessionRecord == nil {
 		return fmt.Errorf("LoadSession returned nil")
 	}
@@ -254,7 +274,7 @@ func (b *Builder) ProcessBundle(preKey *prekey.Bundle) error {
 
 	// Set the local registration ID based on the registration id in our identity key store.
 	sessionState.SetLocalRegistrationID(
-		b.identityKeyStore.GetLocalRegistrationId(),
+		b.identityKeyStore.GetLocalRegistrationID(),
 	)
 
 	// Set the remote registration ID based on the given prekey bundle registrationID.
@@ -268,8 +288,12 @@ func (b *Builder) ProcessBundle(preKey *prekey.Bundle) error {
 	)
 
 	// Store the session in our session store and save the identity in our identity store.
-	b.sessionStore.StoreSession(b.remoteAddress, sessionRecord)
-	b.identityKeyStore.SaveIdentity(b.remoteAddress, preKey.IdentityKey())
+	if err := b.sessionStore.StoreSession(ctx, b.remoteAddress, sessionRecord); err != nil {
+		return err
+	}
+	if err := b.identityKeyStore.SaveIdentity(ctx, b.remoteAddress, preKey.IdentityKey()); err != nil {
+		return err
+	}
 
 	return nil
 }

slidge_whatsapp/vendor/go.mau.fi/libsignal/session/SessionCipher.go

@@ -1,6 +1,7 @@
 package session
 
 import (
+	"context"
 	"fmt"
 
 	"go.mau.fi/libsignal/cipher"
@@ -65,8 +66,11 @@ type Cipher struct {
 
 // Encrypt will take the given message in bytes and return an object that follows
 // the CiphertextMessage interface.
-func (d *Cipher) Encrypt(plaintext []byte) (protocol.CiphertextMessage, error) {
-	sessionRecord := d.sessionStore.LoadSession(d.remoteAddress)
+func (d *Cipher) Encrypt(ctx context.Context, plaintext []byte) (protocol.CiphertextMessage, error) {
+	sessionRecord, err := d.sessionStore.LoadSession(ctx, d.remoteAddress)
+	if err != nil {
+		return nil, err
+	}
 	if sessionRecord == nil {
 		return nil, fmt.Errorf("LoadSession returned nil")
 	}
@@ -125,79 +129,109 @@ func (d *Cipher) Encrypt(plaintext []byte) (protocol.CiphertextMessage, error) {
 	}
 
 	sessionState.SetSenderChainKey(chainKey.NextKey())
-	if !d.identityKeyStore.IsTrustedIdentity(d.remoteAddress, sessionState.RemoteIdentityKey()) {
-		// return err
+	trusted, err := d.identityKeyStore.IsTrustedIdentity(ctx, d.remoteAddress, sessionState.RemoteIdentityKey())
+	if err != nil {
+		return nil, err
+	}
+	if !trusted {
+		return nil, signalerror.ErrUntrustedIdentity
+	}
+	if err := d.identityKeyStore.SaveIdentity(ctx, d.remoteAddress, sessionState.RemoteIdentityKey()); err != nil {
+		return nil, err
+	}
+	if err := d.sessionStore.StoreSession(ctx, d.remoteAddress, sessionRecord); err != nil {
+		return nil, err
 	}
-	d.identityKeyStore.SaveIdentity(d.remoteAddress, sessionState.RemoteIdentityKey())
-	d.sessionStore.StoreSession(d.remoteAddress, sessionRecord)
 	return ciphertextMessage, nil
 }
 
 // Decrypt decrypts the given message using an existing session that
 // is stored in the session store.
-func (d *Cipher) Decrypt(ciphertextMessage *protocol.SignalMessage) ([]byte, error) {
-	plaintext, _, err := d.DecryptAndGetKey(ciphertextMessage)
+func (d *Cipher) Decrypt(ctx context.Context, ciphertextMessage *protocol.SignalMessage) ([]byte, error) {
+	plaintext, _, err := d.DecryptAndGetKey(ctx, ciphertextMessage)
 
 	return plaintext, err
 }
 
 // DecryptAndGetKey decrypts the given message using an existing session that
 // is stored in the session store and returns the message keys used for encryption.
-func (d *Cipher) DecryptAndGetKey(ciphertextMessage *protocol.SignalMessage) ([]byte, *message.Keys, error) {
-	if !d.sessionStore.ContainsSession(d.remoteAddress) {
+func (d *Cipher) DecryptAndGetKey(ctx context.Context, ciphertextMessage *protocol.SignalMessage) ([]byte, *message.Keys, error) {
+	contains, err := d.sessionStore.ContainsSession(ctx, d.remoteAddress)
+	if err != nil {
+		return nil, nil, err
+	}
+	if !contains {
 		return nil, nil, fmt.Errorf("%w %s", signalerror.ErrNoSessionForUser, d.remoteAddress.String())
 	}
 
 	// Load the session record from our session store and decrypt the message.
-	sessionRecord := d.sessionStore.LoadSession(d.remoteAddress)
+	sessionRecord, err := d.sessionStore.LoadSession(ctx, d.remoteAddress)
+	if err != nil {
+		return nil, nil, err
+	}
 	if sessionRecord == nil {
 		return nil, nil, fmt.Errorf("LoadSession returned nil")
 	}
-	plaintext, messageKeys, err := d.DecryptWithRecord(sessionRecord, ciphertextMessage)
+	plaintext, messageKeys, err := d.DecryptWithRecord(ctx, sessionRecord, ciphertextMessage)
 	if err != nil {
 		return nil, nil, err
 	}
 
-	if !d.identityKeyStore.IsTrustedIdentity(d.remoteAddress, sessionRecord.SessionState().RemoteIdentityKey()) {
-		// return err
+	trusted, err := d.identityKeyStore.IsTrustedIdentity(ctx, d.remoteAddress, sessionRecord.SessionState().RemoteIdentityKey())
+	if err != nil {
+		return nil, nil, err
+	}
+	if !trusted {
+		return nil, nil, signalerror.ErrUntrustedIdentity
+	}
+	if err := d.identityKeyStore.SaveIdentity(ctx, d.remoteAddress, sessionRecord.SessionState().RemoteIdentityKey()); err != nil {
+		return nil, nil, err
 	}
-	d.identityKeyStore.SaveIdentity(d.remoteAddress, sessionRecord.SessionState().RemoteIdentityKey())
 
 	// Store the session record in our session store.
-	d.sessionStore.StoreSession(d.remoteAddress, sessionRecord)
+	if err := d.sessionStore.StoreSession(ctx, d.remoteAddress, sessionRecord); err != nil {
+		return nil, nil, err
+	}
 	return plaintext, messageKeys, nil
 }
 
-func (d *Cipher) DecryptMessage(ciphertextMessage *protocol.PreKeySignalMessage) ([]byte, error) {
-	plaintext, _, err := d.DecryptMessageReturnKey(ciphertextMessage)
+func (d *Cipher) DecryptMessage(ctx context.Context, ciphertextMessage *protocol.PreKeySignalMessage) ([]byte, error) {
+	plaintext, _, err := d.DecryptMessageReturnKey(ctx, ciphertextMessage)
 	return plaintext, err
 }
 
-func (d *Cipher) DecryptMessageReturnKey(ciphertextMessage *protocol.PreKeySignalMessage) ([]byte, *message.Keys, error) {
+func (d *Cipher) DecryptMessageReturnKey(ctx context.Context, ciphertextMessage *protocol.PreKeySignalMessage) ([]byte, *message.Keys, error) {
 	// Load or create session record for this session.
-	sessionRecord := d.sessionStore.LoadSession(d.remoteAddress)
+	sessionRecord, err := d.sessionStore.LoadSession(ctx, d.remoteAddress)
+	if err != nil {
+		return nil, nil, err
+	}
 	if sessionRecord == nil {
 		return nil, nil, fmt.Errorf("LoadSession returned nil")
 	}
-	unsignedPreKeyID, err := d.builder.Process(sessionRecord, ciphertextMessage)
+	unsignedPreKeyID, err := d.builder.Process(ctx, sessionRecord, ciphertextMessage)
 	if err != nil {
 		return nil, nil, err
 	}
-	plaintext, keys, err := d.DecryptWithRecord(sessionRecord, ciphertextMessage.WhisperMessage())
+	plaintext, keys, err := d.DecryptWithRecord(ctx, sessionRecord, ciphertextMessage.WhisperMessage())
 	if err != nil {
 		return nil, nil, err
 	}
 	// Store the session record in our session store.
-	d.sessionStore.StoreSession(d.remoteAddress, sessionRecord)
+	if err := d.sessionStore.StoreSession(ctx, d.remoteAddress, sessionRecord); err != nil {
+		return nil, nil, err
+	}
 	if !unsignedPreKeyID.IsEmpty {
-		d.preKeyStore.RemovePreKey(unsignedPreKeyID.Value)
+		if err := d.preKeyStore.RemovePreKey(ctx, unsignedPreKeyID.Value); err != nil {
+			return nil, nil, err
+		}
 	}
 	return plaintext, keys, nil
 }
 
 // DecryptWithKey will decrypt the given message using the given symmetric key. This
 // can be used when decrypting messages at a later time if the message key was saved.
-func (d *Cipher) DecryptWithKey(ciphertextMessage *protocol.SignalMessage, key *message.Keys) ([]byte, error) {
+func (d *Cipher) DecryptWithKey(ctx context.Context, ciphertextMessage *protocol.SignalMessage, key *message.Keys) ([]byte, error) {
 	logger.Debug("Decrypting ciphertext body: ", ciphertextMessage.Body())
 	plaintext, err := decrypt(key, ciphertextMessage.Body())
 	if err != nil {
@@ -209,13 +243,13 @@ func (d *Cipher) DecryptWithKey(ciphertextMessage *protocol.SignalMessage, key *
 }
 
 // DecryptWithRecord decrypts the given message using the given session record.
-func (d *Cipher) DecryptWithRecord(sessionRecord *record.Session, ciphertext *protocol.SignalMessage) ([]byte, *message.Keys, error) {
+func (d *Cipher) DecryptWithRecord(ctx context.Context, sessionRecord *record.Session, ciphertext *protocol.SignalMessage) ([]byte, *message.Keys, error) {
 	logger.Debug("Decrypting ciphertext with record: ", sessionRecord)
 	previousStates := sessionRecord.PreviousSessionStates()
 	sessionState := sessionRecord.SessionState()
 
 	// Try and decrypt the message with the current session state.
-	plaintext, messageKeys, err := d.DecryptWithState(sessionState, ciphertext)
+	plaintext, messageKeys, err := d.DecryptWithState(ctx, sessionState, ciphertext)
 
 	// If we received an error using the current session state, loop
 	// through all previous states.
@@ -223,7 +257,7 @@ func (d *Cipher) DecryptWithRecord(sessionRecord *record.Session, ciphertext *pr
 		logger.Warning(err)
 		for i, state := range previousStates {
 			// Try decrypting the message with previous states
-			plaintext, messageKeys, err = d.DecryptWithState(state, ciphertext)
+			plaintext, messageKeys, err = d.DecryptWithState(ctx, state, ciphertext)
 			if err != nil {
 				continue
 			}
@@ -245,7 +279,7 @@ func (d *Cipher) DecryptWithRecord(sessionRecord *record.Session, ciphertext *pr
 }
 
 // DecryptWithState decrypts the given message with the given session state.
-func (d *Cipher) DecryptWithState(sessionState *record.State, ciphertextMessage *protocol.SignalMessage) ([]byte, *message.Keys, error) {
+func (d *Cipher) DecryptWithState(ctx context.Context, sessionState *record.State, ciphertextMessage *protocol.SignalMessage) ([]byte, *message.Keys, error) {
 	logger.Debug("Decrypting ciphertext with session state: ", sessionState)
 	if !sessionState.HasSenderChain() {
 		logger.Error("Unable to decrypt message with state: ", signalerror.ErrUninitializedSession)
@@ -278,7 +312,7 @@ func (d *Cipher) DecryptWithState(sessionState *record.State, ciphertextMessage
 		return nil, nil, fmt.Errorf("failed to verify ciphertext MAC: %w", err)
 	}
 
-	plaintext, err := d.DecryptWithKey(ciphertextMessage, messageKeys)
+	plaintext, err := d.DecryptWithKey(ctx, ciphertextMessage, messageKeys)
 	if err != nil {
 		return nil, nil, err
 	}