fix

FerretDB · May 5, 2022 · 30ae377 · 30ae377
1 parent 74324ca
commit 30ae377
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/integration/basic_test.go b/integration/basic_test.go
@@ -99,7 +99,7 @@ func TestFindCommentMethod(t *testing.T) {
 	ctx, collection := setup(t, shareddata.Scalars)
 
 	var doc bson.D
-	opts := options.FindOne().SetComment("some test text")
+	opts := options.FindOne().SetComment("*/ 1; DROP DATABASE postgres --")
 	err := collection.FindOne(ctx, bson.D{{"_id", "string"}}, opts).Decode(&doc)
 	require.NoError(t, err)
 }
@@ -109,6 +109,6 @@ func TestFindCommentQuery(t *testing.T) {
 	ctx, collection := setup(t, shareddata.Scalars)
 
 	var doc bson.D
-	err := collection.FindOne(ctx, bson.M{"_id": "string", "$comment": "some test text"}).Decode(&doc)
+	err := collection.FindOne(ctx, bson.M{"_id": "string", "$comment": "*/ 1; DROP DATABASE postgres --"}).Decode(&doc)
 	require.NoError(t, err)
 }
diff --git a/internal/handlers/pg/fetch.go b/internal/handlers/pg/fetch.go
@@ -18,6 +18,7 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"strings"
 
 	"github.com/jackc/pgx/v4"
 
@@ -77,7 +78,10 @@ func nextRow(rows pgx.Rows) (*types.Document, error) {
 func (h *Handler) protoFetch(ctx context.Context, param fetchParam) ([]*types.Document, error) {
 	var sql string
 	if param.comment != "" {
+		param.comment = strings.ReplaceAll(param.comment, "/*", "/ *")
+		param.comment = strings.ReplaceAll(param.comment, "*/", "* /")
 		param.comment = fmt.Sprintf("/* %s */", param.comment)
+
 		sql = fmt.Sprintf(`SELECT %s _jsonb FROM %s`, param.comment, pgx.Identifier{param.db, param.collection}.Sanitize())
 	} else {
 		sql = fmt.Sprintf(`SELECT _jsonb FROM %s`, pgx.Identifier{param.db, param.collection}.Sanitize())