fix(mongo): escape string value startsWith $ in set/upsert (#57)

cordiverse · Dec 2, 2023 · a2bd92f · a2bd92f
1 parent b7ae5b7
commit a2bd92f
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 6 deletions.
diff --git 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
@@ -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
@@ -129,16 +129,18 @@ 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', {
         $or: [
           { 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') },
       ])