TOOLS-123 validate --db and --collection inputs in mongorestore

Former-commit-id: ce3410831670cad1e9680edc155046cab528f22f

3 files changed, 80 insertions, 6 deletions

diff --git a/common/util/mongo.go b/common/util/mongo.go
index 266ce3d35a9..fcd67692708 100644
--- a/common/util/mongo.go
+++ b/common/util/mongo.go
@@ -142,8 +142,10 @@ func ValidateDBName(database string) error {
 	}
 
 	// check for illegal characters
-	if strings.ContainsAny(database, InvalidDBChars) {
-		return fmt.Errorf("illegal character found in '%v'", database)
+	for _, illegalRune := range InvalidDBChars {
+		if strings.ContainsRune(database, illegalRune) {
+			return fmt.Errorf("illegal character '%c' found in '%v'", illegalRune, database)
+		}
 	}
 
 	// db name is valid
@@ -154,6 +156,11 @@ func ValidateDBName(database string) error {
 // collection. An error is returned if it is not valid.
 func ValidateCollectionName(collection string) error {
 
+	// collection names cannot be empty
+	if len(collection) == 0 {
+		return fmt.Errorf("collection name cannot be an empty string")
+	}
+
 	// collection names cannot begin with 'system.'
 	if strings.HasPrefix(collection, "system.") {
 		return fmt.Errorf("collection name '%v' is not allowed to begin with"+
@@ -161,8 +168,10 @@ func ValidateCollectionName(collection string) error {
 	}
 
 	// check for illegal characters
-	if strings.ContainsAny(collection, InvalidCollectionChars) {
-		return fmt.Errorf("illegal character found in '%v'", collection)
+	for _, illegalRune := range InvalidCollectionChars {
+		if strings.ContainsRune(collection, illegalRune) {
+			return fmt.Errorf("illegal character '%c' found in '%v'", illegalRune, collection)
+		}
 	}
 
 	// collection name is valid
diff --git a/common/util/mongo_test.go b/common/util/mongo_test.go
index 748d9ed338e..369f8431907 100644
--- a/common/util/mongo_test.go
+++ b/common/util/mongo_test.go
@@ -65,3 +65,44 @@ func TestParseHost(t *testing.T) {
 
 	})
 }
+
+func TestInvalidNames(t *testing.T) {
+
+	Convey("Checking some invalid collection names, ", t, func() {
+		Convey("test.col$ is invalid", func() {
+			So(ValidateDBName("test"), ShouldBeNil)
+			So(ValidateCollectionName("col$"), ShouldNotBeNil)
+			So(ValidateFullNamespace("test.col$"), ShouldNotBeNil)
+		})
+		Convey("db/aaa.col is invalid", func() {
+			So(ValidateDBName("db/aaa"), ShouldNotBeNil)
+			So(ValidateCollectionName("col"), ShouldBeNil)
+			So(ValidateFullNamespace("db/aaa.col"), ShouldNotBeNil)
+		})
+		Convey("db. is invalid", func() {
+			So(ValidateDBName("db"), ShouldBeNil)
+			So(ValidateCollectionName(""), ShouldNotBeNil)
+			So(ValidateFullNamespace("db."), ShouldNotBeNil)
+		})
+		Convey("db space.col is invalid", func() {
+			So(ValidateDBName("db space"), ShouldNotBeNil)
+			So(ValidateCollectionName("col"), ShouldBeNil)
+			So(ValidateFullNamespace("db space.col"), ShouldNotBeNil)
+		})
+		Convey("[null].[null] is invalid", func() {
+			So(ValidateDBName("\x00"), ShouldNotBeNil)
+			So(ValidateCollectionName("\x00"), ShouldNotBeNil)
+			So(ValidateFullNamespace("\x00.\x00"), ShouldNotBeNil)
+		})
+		Convey("[empty] is invalid", func() {
+			So(ValidateFullNamespace(""), ShouldNotBeNil)
+		})
+		Convey("db.col is valid", func() {
+			So(ValidateDBName("db"), ShouldBeNil)
+			So(ValidateCollectionName("col"), ShouldBeNil)
+			So(ValidateFullNamespace("db.col"), ShouldBeNil)
+		})
+
+	})
+
+}
diff --git a/mongorestore/filepath.go b/mongorestore/filepath.go
index 2b03855a5fd..55e834ae4ac 100644
--- a/mongorestore/filepath.go
+++ b/mongorestore/filepath.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"github.com/mongodb/mongo-tools/common/intents"
 	"github.com/mongodb/mongo-tools/common/log"
+	"github.com/mongodb/mongo-tools/common/util"
 	"io/ioutil"
 	"os"
 	"path/filepath"
@@ -76,7 +77,10 @@ func (restore *MongoRestore) CreateAllIntents(fullpath string) error {
 // CreateIntentsForDB drills down into a db folder, creating intents
 // for all of the collection dump files it encounters
 func (restore *MongoRestore) CreateIntentsForDB(db, fullpath string) error {
-	//TODO check that it's really a directory
+	if err := util.ValidateDBName(db); err != nil {
+		return fmt.Errorf("invalid database name '%v': %v", db, err)
+	}
+
 	log.Logf(log.DebugHigh, "reading collections for database %v in %v", db, fullpath)
 	entries, err := ioutil.ReadDir(fullpath)
 	if err != nil {
@@ -101,6 +105,11 @@ func (restore *MongoRestore) CreateIntentsForDB(db, fullpath string) error {
 							"has .metadata.json files", db)
 					continue
 				}
+				if err = util.ValidateCollectionName(collection); err != nil {
+					if collection != "system.indexes" { // for < 2.6 compatability (TODO remove in 3.0)
+						return fmt.Errorf("invalid collection name '%v.%v': %v", db, collection, err)
+					}
+				}
 				intent := &intents.Intent{
 					DB:       db,
 					C:        collection,
@@ -137,12 +146,27 @@ func hasMetadataFiles(files []os.FileInfo) bool {
 	return false
 }
 
+// CreateIntentForCollection builds an intent for the given db and collection name
+// along with a path to a .bson collection file. It searches the .bson's directory
+// for a matching metadata file. This method is not called by CreateIntentsForDB,
+// it is only used in the case where --db and --collection flags are set.
 func (restore *MongoRestore) CreateIntentForCollection(
 	db, collection, fullpath string) error {
+
 	log.Logf(log.DebugLow, "reading collection %v for database %v from %v",
 		collection, db, fullpath)
 
-	// first make sure the bson file exists and is valid
+	//first validate the collection and db names
+	if err := util.ValidateDBName(db); err != nil {
+		return fmt.Errorf("invalid database name '%v': %v", db, err)
+	}
+	if err := util.ValidateCollectionName(collection); err != nil {
+		if collection != "system.indexes" { // for < 2.6 compatability (TODO remove in 3.0)
+			return fmt.Errorf("invalid collection name '%v.%v': %v", db, collection, err)
+		}
+	}
+
+	// then make sure the bson file exists and is valid
 	file, err := os.Lstat(fullpath)
 	if err != nil {
 		return err