Skip to content

Commit 29d8055

Browse files
tulirtulir
committed
appstate: expand LID hack to snapshots
1 parent 8e1e589 commit 29d8055

5 files changed

Lines changed: 86 additions & 5 deletions

File tree

‎appstate/decode.go‎

Lines changed: 60 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -215,6 +215,7 @@ func (proc *Processor) decodeMutations(
215215
out.AddMAC(indexMAC, valueMAC)
216216
}
217217
out.Mutations = append(out.Mutations, Mutation{
218+
KeyID: mutation.GetRecord().GetKeyID().GetID(),
218219
Operation: mutation.GetOperation(),
219220
Action: syncAction.GetValue(),
220221
Version: syncAction.GetVersion(),
@@ -257,7 +258,14 @@ func (proc *Processor) validateSnapshotMAC(ctx context.Context, name WAPatchName
257258
return
258259
}
259260

260-
func (proc *Processor) decodeSnapshot(ctx context.Context, name WAPatchName, ss *waServerSync.SyncdSnapshot, initialState HashState, validateMACs bool, newMutationsInput []Mutation) (newMutations []Mutation, currentState HashState, err error) {
261+
func (proc *Processor) decodeSnapshot(
262+
ctx context.Context,
263+
name WAPatchName,
264+
ss *waServerSync.SyncdSnapshot,
265+
initialState HashState,
266+
validateMACs bool,
267+
newMutationsInput []Mutation,
268+
) (newMutations []Mutation, currentState HashState, err error) {
261269
currentState = initialState
262270
currentState.Version = ss.GetVersion().GetVersion()
263271

@@ -269,9 +277,17 @@ func (proc *Processor) decodeSnapshot(ctx context.Context, name WAPatchName, ss
269277
}
270278
}
271279

280+
var fakeIndexesToRemove map[[32]byte][]byte
272281
var warn []error
273282
warn, err = currentState.updateHash(encryptedMutations, func(indexMAC []byte, maxIndex int) ([]byte, error) {
274-
return nil, nil
283+
vm, newIndexMAC, err := proc.evilHackForLIDMutation(ctx, name, indexMAC, encryptedMutations[maxIndex], maxIndex, encryptedMutations, false)
284+
if vm != nil && newIndexMAC != nil && len(indexMAC) == 32 {
285+
if fakeIndexesToRemove == nil {
286+
fakeIndexesToRemove = make(map[[32]byte][]byte)
287+
}
288+
fakeIndexesToRemove[indexMACToArray(indexMAC)] = newIndexMAC
289+
}
290+
return vm, err
275291
})
276292
if len(warn) > 0 {
277293
proc.Log.Warnf("Warnings while updating hash for %s: %+v", name, warn)
@@ -291,7 +307,7 @@ func (proc *Processor) decodeSnapshot(ctx context.Context, name WAPatchName, ss
291307

292308
var out patchOutput
293309
out.Mutations = newMutationsInput
294-
err = proc.decodeMutations(ctx, encryptedMutations, &out, validateMACs, currentState.Version, nil)
310+
err = proc.decodeMutations(ctx, encryptedMutations, &out, validateMACs, currentState.Version, fakeIndexesToRemove)
295311
if err != nil {
296312
err = fmt.Errorf("failed to decode snapshot of v%d: %w", currentState.Version, err)
297313
return
@@ -313,6 +329,8 @@ func (proc *Processor) evilHackForLIDMutation(
313329
oldIndexMAC []byte,
314330
mutation *waServerSync.SyncdMutation,
315331
mutationNum int,
332+
prevMutations []*waServerSync.SyncdMutation,
333+
checkDatabase bool,
316334
) (newValueMAC, newIndexMAC []byte, err error) {
317335
_, _, index, _, keys, err := proc.decodeMutation(ctx, mutation, mutationNum, true)
318336
if err != nil {
@@ -351,7 +369,44 @@ func (proc *Processor) evilHackForLIDMutation(
351369
return nil, nil, fmt.Errorf("failed to marshal modified index for LID hack: %w", err)
352370
}
353371
newIndexMAC = concatAndHMAC(sha256.New, keys.Index, indexBytes)
354-
newValueMAC, err = proc.Store.AppState.GetAppStateMutationMAC(ctx, string(patchName), newIndexMAC)
372+
currentKeyID := mutation.GetRecord().GetKeyID().GetID()
373+
for i := mutationNum - 1; i >= 0; i-- {
374+
newKeyID := prevMutations[i].GetRecord().GetKeyID().GetID()
375+
if !bytes.Equal(currentKeyID, newKeyID) {
376+
keys, err = proc.getAppStateKey(ctx, newKeyID)
377+
if err != nil {
378+
return nil, nil, fmt.Errorf("failed to get key %X to decode mutation for LID hack: %w", newKeyID, err)
379+
}
380+
currentKeyID = newKeyID
381+
newIndexMAC = concatAndHMAC(sha256.New, keys.Index, indexBytes)
382+
}
383+
if bytes.Equal(prevMutations[i].GetRecord().GetIndex().GetBlob(), newIndexMAC) {
384+
if prevMutations[i].GetOperation() == waServerSync.SyncdMutation_SET {
385+
value := prevMutations[i].GetRecord().GetValue().GetBlob()
386+
newValueMAC = value[len(value)-32:]
387+
} else {
388+
// Found a REMOVE operation, no previous value
389+
return nil, nil, nil
390+
}
391+
}
392+
}
393+
if newValueMAC == nil && checkDatabase {
394+
newValueMAC, err = proc.Store.AppState.GetAppStateMutationMAC(ctx, string(patchName), newIndexMAC)
395+
if err == nil && newValueMAC == nil {
396+
var allKeys []*store.AppStateSyncKey
397+
allKeys, err = proc.Store.AppStateKeys.GetAllAppStateSyncKeys(ctx)
398+
for _, key := range allKeys {
399+
altIndexMAC := concatAndHMAC(sha256.New, expandAppStateKeys(key.Data).Index, indexBytes)
400+
newValueMAC, err = proc.Store.AppState.GetAppStateMutationMAC(ctx, string(patchName), altIndexMAC)
401+
if newValueMAC != nil {
402+
newIndexMAC = altIndexMAC
403+
}
404+
if err != nil || newValueMAC != nil {
405+
break
406+
}
407+
}
408+
}
409+
}
355410
if err != nil {
356411
// explosions (return is below)
357412
} else if newValueMAC == nil {
@@ -402,7 +457,7 @@ func (proc *Processor) validatePatch(
402457
if vm != nil || err != nil || !allowEvilLIDHack {
403458
return vm, err
404459
}
405-
vm, newIndexMAC, err := proc.evilHackForLIDMutation(ctx, patchName, indexMAC, patch.Mutations[maxIndex], maxIndex)
460+
vm, newIndexMAC, err := proc.evilHackForLIDMutation(ctx, patchName, indexMAC, patch.Mutations[maxIndex], maxIndex, patch.Mutations, true)
406461
if vm != nil && newIndexMAC != nil && len(indexMAC) == 32 {
407462
if fakeIndexesToRemove == nil {
408463
fakeIndexesToRemove = make(map[[32]byte][]byte)

‎appstate/hash.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import (
2020
)
2121

2222
type Mutation struct {
23+
KeyID []byte
2324
Operation waServerSync.SyncdMutation_SyncdOperation
2425
Action *waSyncAction.SyncActionValue
2526
Version int32

‎store/noop.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -136,6 +136,10 @@ func (n *NoopStore) GetLatestAppStateSyncKeyID(ctx context.Context) ([]byte, err
136136
return nil, n.Error
137137
}
138138

139+
func (n *NoopStore) GetAllAppStateSyncKeys(ctx context.Context) ([]*AppStateSyncKey, error) {
140+
return nil, nil
141+
}
142+
139143
func (n *NoopStore) PutAppStateVersion(ctx context.Context, name string, version uint64, hash [128]byte) error {
140144
return n.Error
141145
}

‎store/sqlstore/store.go‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -442,6 +442,7 @@ const (
442442
SET key_data=excluded.key_data, timestamp=excluded.timestamp, fingerprint=excluded.fingerprint
443443
WHERE excluded.timestamp > whatsmeow_app_state_sync_keys.timestamp
444444
`
445+
getAllAppStateSyncKeysQuery = `SELECT key_data, timestamp, fingerprint FROM whatsmeow_app_state_sync_keys WHERE jid=$1 ORDER BY timestamp DESC`
445446
getAppStateSyncKeyQuery = `SELECT key_data, timestamp, fingerprint FROM whatsmeow_app_state_sync_keys WHERE jid=$1 AND key_id=$2`
446447
getLatestAppStateSyncKeyIDQuery = `SELECT key_id FROM whatsmeow_app_state_sync_keys WHERE jid=$1 ORDER BY timestamp DESC LIMIT 1`
447448
)
@@ -451,6 +452,25 @@ func (s *SQLStore) PutAppStateSyncKey(ctx context.Context, id []byte, key store.
451452
return err
452453
}
453454

455+
func (s *SQLStore) GetAllAppStateSyncKeys(ctx context.Context) ([]*store.AppStateSyncKey, error) {
456+
rows, err := s.db.Query(ctx, getAllAppStateSyncKeysQuery, s.JID)
457+
if err != nil {
458+
return nil, err
459+
}
460+
var out []*store.AppStateSyncKey
461+
for rows.Next() {
462+
var item store.AppStateSyncKey
463+
err = rows.Scan(&item.Data, &item.Timestamp, &item.Fingerprint)
464+
if err != nil {
465+
return nil, err
466+
}
467+
if len(item.Data) > 0 {
468+
out = append(out, &item)
469+
}
470+
}
471+
return out, rows.Close()
472+
}
473+
454474
func (s *SQLStore) GetAppStateSyncKey(ctx context.Context, id []byte) (*store.AppStateSyncKey, error) {
455475
var key store.AppStateSyncKey
456476
err := s.db.QueryRow(ctx, getAppStateSyncKeyQuery, s.JID, id).Scan(&key.Data, &key.Timestamp, &key.Fingerprint)

‎store/store.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@ type AppStateSyncKeyStore interface {
6161
PutAppStateSyncKey(ctx context.Context, id []byte, key AppStateSyncKey) error
6262
GetAppStateSyncKey(ctx context.Context, id []byte) (*AppStateSyncKey, error)
6363
GetLatestAppStateSyncKeyID(ctx context.Context) ([]byte, error)
64+
GetAllAppStateSyncKeys(ctx context.Context) ([]*AppStateSyncKey, error)
6465
}
6566

6667
type AppStateMutationMAC struct {

0 commit comments

Comments
 (0)