diff --git a/weed/filer/abstract_sql/abstract_sql_store.go b/weed/filer/abstract_sql/abstract_sql_store.go index 52f959d80a..14a552315d 100644 --- a/weed/filer/abstract_sql/abstract_sql_store.go +++ b/weed/filer/abstract_sql/abstract_sql_store.go @@ -50,6 +50,9 @@ func (store *AbstractSqlStore) CanDropWholeBucket() bool { return store.SupportBucketTable && !store.SkipDDL } func (store *AbstractSqlStore) OnBucketCreation(bucket string) { + if !isValidBucket(bucket) { + return + } store.dbsLock.Lock() defer store.dbsLock.Unlock() @@ -61,6 +64,9 @@ func (store *AbstractSqlStore) OnBucketCreation(bucket string) { store.dbs[bucket] = true } func (store *AbstractSqlStore) OnBucketDeletion(bucket string) { + if !isValidBucket(bucket) { + return + } store.dbsLock.Lock() defer store.dbsLock.Unlock() diff --git a/weed/filer/abstract_sql/abstract_sql_store_test.go b/weed/filer/abstract_sql/abstract_sql_store_test.go index 818360327c..027c5b308e 100644 --- a/weed/filer/abstract_sql/abstract_sql_store_test.go +++ b/weed/filer/abstract_sql/abstract_sql_store_test.go @@ -2,8 +2,10 @@ package abstract_sql import ( "context" + "reflect" "testing" + "github.com/DATA-DOG/go-sqlmock" "github.com/seaweedfs/seaweedfs/weed/util" ) @@ -89,3 +91,52 @@ func TestSplitPoolForKv(t *testing.T) { } } } + +type recordingDDLGen struct { + SqlGenerator + calls []string +} + +func (g *recordingDDLGen) GetSqlCreateTable(tableName string) string { + g.calls = append(g.calls, "create "+tableName) + return "CREATE TABLE " + tableName +} + +func (g *recordingDDLGen) GetSqlDropTable(tableName string) string { + g.calls = append(g.calls, "drop "+tableName) + return "DROP" +} + +func TestBucketHooksSkipInvalidBucketNames(t *testing.T) { + db, mock, err := sqlmock.New() + if err != nil { + t.Fatal(err) + } + defer db.Close() + mock.ExpectExec("CREATE TABLE mybucket").WillReturnResult(sqlmock.NewResult(0, 0)) + mock.ExpectExec("DROP").WillReturnResult(sqlmock.NewResult(0, 0)) + + gen := &recordingDDLGen{} + store := &AbstractSqlStore{SqlGenerator: gen, DB: db, SupportBucketTable: true} + + for _, name := range []string{ + `.x01" (a int);CREATE TABLE pwned(flag text);--`, + `.z01","marker`, + `b01" (a int); DROP TABLE filemeta; --`, + DEFAULT_TABLE, + ".system", + "", + } { + store.OnBucketCreation(name) + store.OnBucketDeletion(name) + } + store.OnBucketCreation("mybucket") + store.OnBucketDeletion("mybucket") + + if want := []string{"create mybucket", "drop mybucket"}; !reflect.DeepEqual(gen.calls, want) { + t.Errorf("DDL calls = %q, want %q", gen.calls, want) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Error(err) + } +}