From a2bd92ffdfb4a12e8e12c088f155c254f3b219b6 Mon Sep 17 00:00:00 2001
From: Hieuzest <girkirin@hotmail.com>
Date: Sat, 2 Dec 2023 23:12:56 +0800
Subject: [PATCH] fix(mongo): escape string value startsWith $ in
 `set`/`upsert` (#57)

---
 packages/core/src/driver.ts  | 1 +
 packages/mongo/src/index.ts  | 6 +++---
 packages/tests/src/update.ts | 8 +++++---
 3 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/packages/core/src/driver.ts b/packages/core/src/driver.ts
index f1d7ac30..e78cdcf4 100644
--- a/packages/core/src/driver.ts
+++ b/packages/core/src/driver.ts
@@ -219,6 +219,7 @@ export class Database<S = any> {
   }
 
   async stats() {
+    await this.prepared()
     const stats: Driver.Stats = { size: 0, tables: {} }
     await Promise.all(Object.values(this.drivers).map(async (driver) => {
       const { size = 0, tables } = await driver.stats()
diff --git a/packages/mongo/src/index.ts b/packages/mongo/src/index.ts
index 0886ffc4..8279be17 100644
--- a/packages/mongo/src/index.ts
+++ b/packages/mongo/src/index.ts
@@ -1,5 +1,5 @@
 import { BSONType, ClientSession, Collection, Db, IndexDescription, MongoClient, MongoClientOptions, MongoError } from 'mongodb'
-import { Dict, isNullable, makeArray, noop, omit, pick } from 'cosmokit'
+import { Dict, isNullable, makeArray, mapValues, noop, omit, pick } from 'cosmokit'
 import { Database, Driver, Eval, executeEval, executeUpdate, Query, RuntimeError, Selection } from '@minatojs/core'
 import { URLSearchParams } from 'url'
 import { Transformer } from './utils'
@@ -445,7 +445,7 @@ export class MongoDriver extends Driver {
     const coll = this.db.collection(table)
 
     const transformer = new Transformer(this.getVirtualKey(table), undefined, '$' + tempKey + '.')
-    const $set = transformer.eval(update)
+    const $set = transformer.eval(mapValues(update, (value: any) => typeof value === 'string' && value.startsWith('$') ? { $literal: value } : value))
     const $unset = Object.entries($set)
       .filter(([_, value]) => typeof value === 'object')
       .map(([key, _]) => key)
@@ -563,7 +563,7 @@ export class MongoDriver extends Driver {
       for (const update of data) {
         const query = this.transformQuery(pick(update, keys), table)!
         const transformer = new Transformer(this.getVirtualKey(table), undefined, '$' + tempKey + '.')
-        const $set = transformer.eval(update)
+        const $set = transformer.eval(mapValues(update, (value: any) => typeof value === 'string' && value.startsWith('$') ? { $literal: value } : value))
         const $unset = Object.entries($set)
           .filter(([_, value]) => typeof value === 'object')
           .map(([key, _]) => key)
diff --git a/packages/tests/src/update.ts b/packages/tests/src/update.ts
index 8bb67d50..e3c5e93a 100644
--- a/packages/tests/src/update.ts
+++ b/packages/tests/src/update.ts
@@ -129,8 +129,10 @@ namespace OrmOperations {
       const table = await setup(database, 'temp2', barTable)
       const data = table.find(bar => bar.timestamp)!
       data.list = ['2', '3', '3']
+      data.text = `$'"%~\``
       const magicIds = table.slice(2, 4).map((data) => {
         data.list = ['2', '3', '3']
+        data.text = `$'"%~\``
         return data.id
       })
       await expect(database.set('temp2', {
@@ -138,7 +140,7 @@ namespace OrmOperations {
           { id: magicIds },
           { timestamp: magicBorn },
         ],
-      }, { list: ['2', '3', '3'] })).to.eventually.have.shape({ matched: 3 })
+      }, { list: ['2', '3', '3'], text: `$'"%~\`` })).to.eventually.have.shape({ matched: 3 })
       await expect(database.get('temp2', {})).to.eventually.have.shape(table)
     })
 
@@ -218,11 +220,11 @@ namespace OrmOperations {
 
     it('multi condition on composite primary', async () => {
       const table = await setup(database, 'temp3', bazTable)
-      table[1].value = 'bb'
+      table[1].value = `$'"%~\``
       table[2].value = 'cc'
       table.push({ ida: 114, idb: '514', value: 'baz' })
       await database.upsert('temp3', row => [
-        { ida: 2, idb: 'a', value: 'bb' },
+        { ida: 2, idb: 'a', value: `$'"%~\`` },
         { ida: 1, idb: 'b', value: 'cc' },
         { ida: 114, idb: '514', value: $.concat(row.value, 'baz') },
       ])