Forbid short document keys like $k

Closes #181.
FerretDB · Jan 6, 2022 · c39fb58 · c39fb58
1 parent 9ccdcd4
commit c39fb58
Show file tree

Hide file tree

Showing 5 changed files with 38 additions and 26 deletions.
diff --git a/internal/bson/bson_test.go b/internal/bson/bson_test.go
@@ -179,7 +179,7 @@ func benchmark(b *testing.B, testCases []testCase, newFunc func() bsontype) {
 				for i := 0; i < b.N; i++ {
 					v = newFunc()
 					readErr = v.ReadFrom(bufio.NewReader(br))
-					_, seekErr = br.Seek(io.SeekStart, 0)
+					_, seekErr = br.Seek(0, io.SeekStart)
 				}
 
 				b.StopTimer()

diff --git a/internal/types/document.go b/internal/types/document.go
@@ -26,6 +26,11 @@ func isValidKey(key string) bool {
 		return false
 	}
 
+	// forbid $k, but allow $db
+	if key[0] == '$' && len(key) <= 2 {
+		return false
+	}
+
 	return utf8.ValidString(key)
 }
 

diff --git a/internal/types/document_test.go b/internal/types/document_test.go
@@ -58,6 +58,13 @@ func TestValidate(t *testing.T) {
 			m:    map[string]any{"0": "foo", "1": "bar"},
 		},
 		err: fmt.Errorf(`types.Document.validate: duplicate key: "0"`),
+	}, {
+		name: "fjson keys",
+		doc: Document{
+			keys: []string{"$k"},
+			m:    map[string]any{"$k": "foo"},
+		},
+		err: fmt.Errorf(`types.Document.validate: invalid key: "$k"`),
 	}} {
 		tc := tc
 		t.Run(tc.name, func(t *testing.T) {

diff --git a/internal/wire/op_msg_test.go b/internal/wire/op_msg_test.go
@@ -142,26 +142,7 @@ var msgTestCases = []testCase{{
 }, {
 	name:      "dollar_dot",
 	expectedB: testutil.MustParseDumpFile("testdata", "dollar_dot.hex"),
-	msgHeader: &MsgHeader{
-		MessageLength: 113,
-		RequestID:     35,
-		OpCode:        OP_MSG,
-	},
-	msgBody: &OpMsg{
-		sections: []OpMsgSection{{
-			Documents: []types.Document{types.MustMakeDocument(
-				"insert", "test",
-				"documents", types.MustNewArray(
-					types.MustMakeDocument(
-						"$.", true,
-						"_id", types.ObjectID{0x61, 0xaf, 0x7c, 0x02, 0x75, 0x48, 0x20, 0xa5, 0x92, 0x3e, 0xa4, 0x97},
-					),
-				),
-				"ordered", true,
-				"$db", "test",
-			)},
-		}},
-	},
+	err:       `types.Document.validate: invalid key: "$."`,
 }}
 
 func TestMsg(t *testing.T) {

diff --git a/internal/wire/wire_test.go b/internal/wire/wire_test.go
@@ -17,6 +17,7 @@ package wire
 import (
 	"bufio"
 	"bytes"
+	"errors"
 	"testing"
 	"time"
 
@@ -33,6 +34,7 @@ type testCase struct {
 	expectedB []byte
 	msgHeader *MsgHeader
 	msgBody   MsgBody
+	err       string // unwrapped
 }
 
 func testMessages(t *testing.T, testCases []testCase) {
@@ -63,14 +65,31 @@ func testMessages(t *testing.T, testCases []testCase) {
 				br := bytes.NewReader(tc.expectedB)
 				bufr := bufio.NewReader(br)
 				msgHeader, msgBody, err := ReadMessage(bufr)
-				require.NoError(t, err)
-				assert.Equal(t, tc.msgHeader, msgHeader)
-				assert.Equal(t, tc.msgBody, msgBody)
-				assert.Zero(t, br.Len(), "not all br bytes were consumed")
-				assert.Zero(t, bufr.Buffered(), "not all bufr bytes were consumed")
+				if tc.err == "" {
+					assert.NoError(t, err)
+					assert.Equal(t, tc.msgHeader, msgHeader)
+					assert.Equal(t, tc.msgBody, msgBody)
+					assert.Zero(t, br.Len(), "not all br bytes were consumed")
+					assert.Zero(t, bufr.Buffered(), "not all bufr bytes were consumed")
+					return
+				}
+
+				require.Error(t, err)
+				for {
+					e := errors.Unwrap(err)
+					if e == nil {
+						break
+					}
+					err = e
+				}
+				require.Equal(t, tc.err, err.Error())
 			})
 
 			t.Run("WriteMessage", func(t *testing.T) {
+				if tc.msgHeader == nil {
+					t.Skip("msgHeader is nil")
+				}
+
 				t.Parallel()
 
 				var buf bytes.Buffer