franpcode/Baileys
1
0
Fork 0
mirror of https://github.com/FranP-code/Baileys.git synced 2025-10-13 00:32:22 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor: upload more pre-keys + use pure function

Browse Source
This commit is contained in:
Adhiraj Singh
2022-04-13 00:02:10 +05:30
parent 08fa212db8
commit 02265fd7f9
4 changed files with 121 additions and 110 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
src/Defaults/index.ts
View File

@@ -61,4 +61,6 @@ export const MEDIA_KEYS = Object.keys(MEDIA_PATH_MAP) as MediaType[]
export const KEY_BUNDLE_TYPE = Buffer.from([5])
export const MIN_PREKEY_COUNT = 5
export const MIN_PREKEY_COUNT = 5
export const INITIAL_PREKEY_COUNT = 50

115
src/Socket/messages-recv.ts
View File

@@ -2,7 +2,7 @@
import { proto } from '../../WAProto'
import { KEY_BUNDLE_TYPE, MIN_PREKEY_COUNT } from '../Defaults'
import { BaileysEventMap, MessageReceiptType, MessageUserReceipt, SocketConfig, WAMessageStubType } from '../Types'
import { debouncedTimeout, decodeMessageStanza, delay, encodeBigEndian, generateSignalPubKey, getStatusFromReceiptType, normalizeMessageContent, xmppPreKey, xmppSignedPreKey } from '../Utils'
import { debouncedTimeout, decodeMessageStanza, delay, encodeBigEndian, generateSignalPubKey, getNextPreKeys, getStatusFromReceiptType, normalizeMessageContent, xmppPreKey, xmppSignedPreKey } from '../Utils'
import { makeKeyedMutex, makeMutex } from '../Utils/make-mutex'
import processMessage from '../Utils/process-message'
import { areJidsSameUser, BinaryNode, BinaryNodeAttributes, getAllBinaryNodeChildren, getBinaryNodeChild, getBinaryNodeChildren, isJidGroup, isJidUser, jidDecode, jidEncode, jidNormalizedUser, S_WHATSAPP_NET } from '../WABinary'
@@ -22,7 +22,6 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => {
ws,
onUnexpectedError,
assertSessions,
assertingPreKeys,
sendNode,
relayMessage,
sendReceipt,
@@ -78,64 +77,70 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => {
const { account, signedPreKey, signedIdentityKey: identityKey } = authState.creds
const deviceIdentity = proto.ADVSignedDeviceIdentity.encode(account).finish()
await assertingPreKeys(1, async preKeys => {
const [keyId] = Object.keys(preKeys)
const key = preKeys[+keyId]
await authState.keys.transaction(
async() => {
const { update, preKeys } = await getNextPreKeys(authState, 1)
const decFrom = node.attrs.from ? jidDecode(node.attrs.from) : undefined
const receipt: BinaryNode = {
tag: 'receipt',
attrs: {
id: msgId,
type: 'retry',
to: isGroup ? node.attrs.from : jidEncode(decFrom!.user, 's.whatsapp.net', decFrom!.device, 0)
},
content: [
{
tag: 'retry',
attrs: {
count: retryCount.toString(),
id: node.attrs.id,
t: node.attrs.t,
v: '1'
}
const [keyId] = Object.keys(preKeys)
const key = preKeys[+keyId]
const decFrom = node.attrs.from ? jidDecode(node.attrs.from) : undefined
const receipt: BinaryNode = {
tag: 'receipt',
attrs: {
id: msgId,
type: 'retry',
to: isGroup ? node.attrs.from : jidEncode(decFrom!.user, 's.whatsapp.net', decFrom!.device, 0)
},
{
tag: 'registration',
attrs: { },
content: encodeBigEndian(authState.creds.registrationId)
}
]
}
if(node.attrs.recipient) {
receipt.attrs.recipient = node.attrs.recipient
}
if(node.attrs.participant) {
receipt.attrs.participant = node.attrs.participant
}
if(retryCount > 1) {
const exec = generateSignalPubKey(Buffer.from(KEY_BUNDLE_TYPE)).slice(0, 1)
const content = receipt.content! as BinaryNode[]
content.push({
tag: 'keys',
attrs: { },
content: [
{ tag: 'type', attrs: { }, content: exec },
{ tag: 'identity', attrs: { }, content: identityKey.public },
xmppPreKey(key, +keyId),
xmppSignedPreKey(signedPreKey),
{ tag: 'device-identity', attrs: { }, content: deviceIdentity }
{
tag: 'retry',
attrs: {
count: retryCount.toString(),
id: node.attrs.id,
t: node.attrs.t,
v: '1'
}
},
{
tag: 'registration',
attrs: { },
content: encodeBigEndian(authState.creds.registrationId)
}
]
})
}
if(node.attrs.recipient) {
receipt.attrs.recipient = node.attrs.recipient
}
if(node.attrs.participant) {
receipt.attrs.participant = node.attrs.participant
}
if(retryCount > 1) {
const exec = generateSignalPubKey(Buffer.from(KEY_BUNDLE_TYPE)).slice(0, 1)
const content = receipt.content! as BinaryNode[]
content.push({
tag: 'keys',
attrs: { },
content: [
{ tag: 'type', attrs: { }, content: exec },
{ tag: 'identity', attrs: { }, content: identityKey.public },
xmppPreKey(key, +keyId),
xmppSignedPreKey(signedPreKey),
{ tag: 'device-identity', attrs: { }, content: deviceIdentity }
]
})
}
await sendNode(receipt)
logger.info({ msgAttrs: node.attrs, retryCount }, 'sent retry receipt')
ev.emit('creds.update', update)
}
await sendNode(receipt)
logger.info({ msgAttrs: node.attrs, retryCount }, 'sent retry receipt')
})
)
}
const processMessageLocal = async(msg: proto.IWebMessageInfo) => {

64
src/Socket/socket.ts
View File

@@ -3,13 +3,11 @@ import EventEmitter from 'events'
import { promisify } from 'util'
import WebSocket from 'ws'
import { proto } from '../../WAProto'
import { DEF_CALLBACK_PREFIX, DEF_TAG_PREFIX, DEFAULT_ORIGIN, KEY_BUNDLE_TYPE, MIN_PREKEY_COUNT } from '../Defaults'
import { DEF_CALLBACK_PREFIX, DEF_TAG_PREFIX, DEFAULT_ORIGIN, INITIAL_PREKEY_COUNT, MIN_PREKEY_COUNT } from '../Defaults'
import { AuthenticationCreds, BaileysEventEmitter, BaileysEventMap, DisconnectReason, SocketConfig } from '../Types'
import { addTransactionCapability, bindWaitForConnectionUpdate, configureSuccessfulPairing, Curve, encodeBigEndian, generateLoginNode, generateMdTagPrefix, generateOrGetPreKeys, generateRegistrationNode, getPreKeys, makeNoiseHandler, printQRIfNecessaryListener, promiseTimeout, useSingleFileAuthState, xmppPreKey, xmppSignedPreKey } from '../Utils'
import { addTransactionCapability, bindWaitForConnectionUpdate, BufferJSON, configureSuccessfulPairing, Curve, generateLoginNode, generateMdTagPrefix, generateRegistrationNode, getNextPreKeysNode, makeNoiseHandler, printQRIfNecessaryListener, promiseTimeout, useSingleFileAuthState } from '../Utils'
import { assertNodeErrorFree, BinaryNode, encodeBinaryNode, getBinaryNodeChild, S_WHATSAPP_NET } from '../WABinary'
const INITIAL_PREKEY_COUNT = 30
/**
* Connects to WA servers and performs:
* - simple queries (no retry mechanism, wait for connection establishment)
@@ -219,31 +217,6 @@ export const makeSocket = ({
startKeepAliveRequest()
}
/**
* get some pre-keys and do something with them
* @param range how many pre-keys to get
* @param execute what to do with them
*/
const assertingPreKeys = async(range: number, execute: (keys: { [_: number]: any }) => Promise<void>) => {
const { newPreKeys, lastPreKeyId, preKeysRange } = generateOrGetPreKeys(authState.creds, range)
const update: Partial<AuthenticationCreds> = {
nextPreKeyId: Math.max(lastPreKeyId + 1, creds.nextPreKeyId),
firstUnuploadedPreKeyId: Math.max(creds.firstUnuploadedPreKeyId, lastPreKeyId + 1)
}
await keys.transaction(
async() => {
await keys.set({ 'pre-key': newPreKeys })
const preKeys = await getPreKeys(keys, preKeysRange[0], preKeysRange[0] + preKeysRange[1])
await execute(preKeys)
}
)
ev.emit('creds.update', update)
}
const getAvailablePreKeysOnServer = async() => {
const result = await query({
tag: 'iq',
@@ -263,29 +236,19 @@ export const makeSocket = ({
/** generates and uploads a set of pre-keys to the server */
const uploadPreKeys = async(count = INITIAL_PREKEY_COUNT) => {
await assertingPreKeys(count, async preKeys => {
logger.info('uploading pre-keys')
await keys.transaction(
async() => {
logger.info({ count }, 'uploading pre-keys')
const { update, node } = await getNextPreKeysNode({ creds, keys }, count)
const node: BinaryNode = {
tag: 'iq',
attrs: {
id: generateMessageTag(),
xmlns: 'encrypt',
type: 'set',
to: S_WHATSAPP_NET,
},
content: [
{ tag: 'registration', attrs: { }, content: encodeBigEndian(creds.registrationId) },
{ tag: 'type', attrs: { }, content: KEY_BUNDLE_TYPE },
{ tag: 'identity', attrs: { }, content: creds.signedIdentityKey.public },
{ tag: 'list', attrs: { }, content: Object.keys(preKeys).map(k => xmppPreKey(preKeys[+k], +k)) },
xmppSignedPreKey(creds.signedPreKey)
]
console.log(JSON.stringify(node, BufferJSON.replacer, 2))
await query(node)
ev.emit('creds.update', update)
logger.info({ count }, 'uploaded pre-keys')
}
await query(node)
logger.info('uploaded pre-keys')
})
)
}
const uploadPreKeysToServerIfRequired = async() => {
@@ -616,7 +579,6 @@ export const makeSocket = ({
return authState.creds.me
},
emitEventsFromMap,
assertingPreKeys,
generateMessageTag,
query,
waitForMessage,

48
src/Utils/signal.ts
View File

@@ -1,8 +1,9 @@
import * as libsignal from 'libsignal'
import { proto } from '../../WAProto'
import { GroupCipher, GroupSessionBuilder, SenderKeyDistributionMessage, SenderKeyName, SenderKeyRecord } from '../../WASignalGroup'
import { AuthenticationCreds, KeyPair, SignalAuthState, SignalIdentity, SignalKeyStore, SignedKeyPair } from '../Types/Auth'
import { assertNodeErrorFree, BinaryNode, getBinaryNodeChild, getBinaryNodeChildBuffer, getBinaryNodeChildren, getBinaryNodeChildUInt, jidDecode, JidWithDevice } from '../WABinary'
import { KEY_BUNDLE_TYPE } from '../Defaults'
import { AuthenticationCreds, AuthenticationState, KeyPair, SignalAuthState, SignalIdentity, SignalKeyStore, SignedKeyPair } from '../Types/Auth'
import { assertNodeErrorFree, BinaryNode, getBinaryNodeChild, getBinaryNodeChildBuffer, getBinaryNodeChildren, getBinaryNodeChildUInt, jidDecode, JidWithDevice, S_WHATSAPP_NET } from '../WABinary'
import { Curve } from './crypto'
import { encodeBigEndian } from './generics'
@@ -60,7 +61,6 @@ export const generateOrGetPreKeys = (creds: AuthenticationCreds, range: number)
}
}
export const xmppSignedPreKey = (key: SignedKeyPair): BinaryNode => (
{
tag: 'skey',
@@ -273,4 +273,46 @@ export const extractDeviceJids = (result: BinaryNode, myJid: string, excludeZero
}
return extracted
}
/**
* get the next N keys for upload or processing
* @param count number of pre-keys to get or generate
*/
export const getNextPreKeys = async({ creds, keys }: AuthenticationState, count: number) => {
const { newPreKeys, lastPreKeyId, preKeysRange } = generateOrGetPreKeys(creds, count)
const update: Partial<AuthenticationCreds> = {
nextPreKeyId: Math.max(lastPreKeyId + 1, creds.nextPreKeyId),
firstUnuploadedPreKeyId: Math.max(creds.firstUnuploadedPreKeyId, lastPreKeyId + 1)
}
await keys.set({ 'pre-key': newPreKeys })
const preKeys = await getPreKeys(keys, preKeysRange[0], preKeysRange[0] + preKeysRange[1])
return { update, preKeys }
}
export const getNextPreKeysNode = async(state: AuthenticationState, count: number) => {
const { creds } = state
const { update, preKeys } = await getNextPreKeys(state, count)
const node: BinaryNode = {
tag: 'iq',
attrs: {
xmlns: 'encrypt',
type: 'set',
to: S_WHATSAPP_NET,
},
content: [
{ tag: 'registration', attrs: { }, content: encodeBigEndian(creds.registrationId) },
{ tag: 'type', attrs: { }, content: KEY_BUNDLE_TYPE },
{ tag: 'identity', attrs: { }, content: creds.signedIdentityKey.public },
{ tag: 'list', attrs: { }, content: Object.keys(preKeys).map(k => xmppPreKey(preKeys[+k], +k)) },
xmppSignedPreKey(creds.signedPreKey)
]
}
return { update, node }
}