privilege, executor: support SET DEFAULT ROLE

executor/simple.go

@@ -87,11 +87,152 @@ func (e *SimpleExec) Next(ctx context.Context, req *chunk.RecordBatch) (err erro
  err = e.executeSetRole(x)
  case *ast.RevokeRoleStmt:
  err = e.executeRevokeRole(x)
+ case *ast.SetDefaultRoleStmt:
+ err = e.executeSetDefaultRole(x)
  }
  e.done = true
  return err
 }
 
+func (e *SimpleExec) setDefaultRoleNone(s *ast.SetDefaultRoleStmt) error {
+ sqlExecutor := e.ctx.(sqlexec.SQLExecutor)
+ if _, err := sqlExecutor.Execute(context.Background(), "begin"); err != nil {
+ return err
+ }
+ for _, u := range s.UserList {
+ if u.Hostname == "" {
+ u.Hostname = "%"
+ }
+ sql := fmt.Sprintf("DELETE IGNORE FROM mysql.default_roles WHERE USER='%s' AND HOST='%s';", u.Username, u.Hostname)
+ if _, err := sqlExecutor.Execute(context.Background(), sql); err != nil {
+ logutil.Logger(context.Background()).Error(fmt.Sprintf("Error occur when executing %s", sql))
+ if _, rollbackErr := sqlExecutor.Execute(context.Background(), "rollback"); rollbackErr != nil {
+ return rollbackErr
+ }
+ return err
+ }
+ }
+ if _, err := sqlExecutor.Execute(context.Background(), "commit"); err != nil {
+ return err
+ }
+ return nil
+}
+
+func (e *SimpleExec) setDefaultRoleRegular(s *ast.SetDefaultRoleStmt) error {
+ for _, user := range s.UserList {
+ exists, err := userExists(e.ctx, user.Username, user.Hostname)
+ if err != nil {
+ return err
+ }
+ if !exists {
+ return ErrCannotUser.GenWithStackByArgs("SET DEFAULT ROLE", user.String())
+ }
+ }
+ for _, role := range s.RoleList {
+ exists, err := userExists(e.ctx, role.Username, role.Hostname)
+ if err != nil {
+ return err
+ }
+ if !exists {
+ return ErrCannotUser.GenWithStackByArgs("SET DEFAULT ROLE", role.String())
+ }
+ }
+ sqlExecutor := e.ctx.(sqlexec.SQLExecutor)
+ if _, err := sqlExecutor.Execute(context.Background(), "begin"); err != nil {
+ return err
+ }
+ for _, user := range s.UserList {
+ if user.Hostname == "" {
+ user.Hostname = "%"
+ }
+ sql := fmt.Sprintf("DELETE IGNORE FROM mysql.default_roles WHERE USER='%s' AND HOST='%s';", user.Username, user.Hostname)
+ if _, err := sqlExecutor.Execute(context.Background(), sql); err != nil {
+ logutil.Logger(context.Background()).Error(fmt.Sprintf("Error occur when executing %s", sql))
+ if _, rollbackErr := sqlExecutor.Execute(context.Background(), "rollback"); rollbackErr != nil {
+ return rollbackErr
+ }
+ return err
+ }
+ for _, role := range s.RoleList {
+ sql := fmt.Sprintf("INSERT IGNORE INTO mysql.default_roles values('%s', '%s', '%s', '%s');", user.Hostname, user.Username, role.Hostname, role.Username)
+ checker := privilege.GetPrivilegeManager(e.ctx)
+ ok := checker.FindEdge(e.ctx, role, user)
+ if ok {
+ if _, err := sqlExecutor.Execute(context.Background(), sql); err != nil {
+ logutil.Logger(context.Background()).Error(fmt.Sprintf("Error occur when executing %s", sql))
+ if _, rollbackErr := sqlExecutor.Execute(context.Background(), "rollback"); rollbackErr != nil {
+ return rollbackErr
+ }
+ return err
+ }
+ } else {
+ if _, rollbackErr := sqlExecutor.Execute(context.Background(), "rollback"); rollbackErr != nil {
+ return rollbackErr
+ }
+ return ErrRoleNotGranted.GenWithStackByArgs(role.String(), user.String())
+ }
+ }
+ }
+ if _, err := sqlExecutor.Execute(context.Background(), "commit"); err != nil {
+ return err
+ }
+ return nil
+}
+
+func (e *SimpleExec) setDefaultRoleAll(s *ast.SetDefaultRoleStmt) error {
+ for _, user := range s.UserList {
+ exists, err := userExists(e.ctx, user.Username, user.Hostname)
+ if err != nil {
+ return err
+ }
+ if !exists {
+ return ErrCannotUser.GenWithStackByArgs("SET DEFAULT ROLE", user.String())
+ }
+ }
+ sqlExecutor := e.ctx.(sqlexec.SQLExecutor)
+ if _, err := sqlExecutor.Execute(context.Background(), "begin"); err != nil {
+ return err
+ }
+ for _, user := range s.UserList {
+ if user.Hostname == "" {
+ user.Hostname = "%"
+ }
+ sql := fmt.Sprintf("DELETE IGNORE FROM mysql.default_roles WHERE USER='%s' AND HOST='%s';", user.Username, user.Hostname)
+ if _, err := sqlExecutor.Execute(context.Background(), sql); err != nil {
+ logutil.Logger(context.Background()).Error(fmt.Sprintf("Error occur when executing %s", sql))
+ if _, rollbackErr := sqlExecutor.Execute(context.Background(), "rollback"); rollbackErr != nil {
+ return rollbackErr
+ }
+ return err
+ }
+ sql = fmt.Sprintf("INSERT IGNORE INTO mysql.default_roles(HOST,USER,DEFAULT_ROLE_HOST,DEFAULT_ROLE_USER) "+
+ "SELECT TO_HOST,TO_USER,FROM_HOST,FROM_USER FROM mysql.role_edges WHERE TO_HOST='%s' AND TO_USER='%s';", user.Hostname, user.Username)
+ if _, err := sqlExecutor.Execute(context.Background(), sql); err != nil {
+ if _, rollbackErr := sqlExecutor.Execute(context.Background(), "rollback"); rollbackErr != nil {
+ return rollbackErr
+ }
+ return err
+ }
+ }
+ if _, err := sqlExecutor.Execute(context.Background(), "commit"); err != nil {
+ return err
+ }
+ return nil
+}
+
+func (e *SimpleExec) executeSetDefaultRole(s *ast.SetDefaultRoleStmt) error {
+ switch s.SetRoleOpt {
+ case ast.SetRoleAll:
+ return e.setDefaultRoleAll(s)
+ case ast.SetRoleNone:
+ return e.setDefaultRoleNone(s)
+ case ast.SetRoleRegular:
+ return e.setDefaultRoleRegular(s)
+ }
+ err := domain.GetDomain(e.ctx).PrivilegeHandle().Update(e.ctx.(sessionctx.Context))
+ return err
+}
+
 func (e *SimpleExec) executeSetRole(s *ast.SetRoleStmt) error {
  checkDup := make(map[string]*auth.RoleIdentity, len(s.RoleList))
  // Check whether RoleNameList contain duplicate role name.

executor/simple_test.go

@@ -160,6 +160,56 @@ func (s *testSuite3) TestRole(c *C) {
  tk.MustExec(dropRoleSQL)
 }
 
+func (s *testSuite3) TestDefaultRole(c *C) {
+ tk := testkit.NewTestKit(c, s.store)
+
+ createRoleSQL := `CREATE ROLE r_1, r_2, r_3, u_1;`
+ tk.MustExec(createRoleSQL)
+
+ tk.MustExec("insert into mysql.role_edges (FROM_HOST,FROM_USER,TO_HOST,TO_USER) values ('%','r_1','%','u_1')")
+ tk.MustExec("insert into mysql.role_edges (FROM_HOST,FROM_USER,TO_HOST,TO_USER) values ('%','r_2','%','u_1')")
+
+ tk.MustExec("flush privileges;")
+
+ setRoleSQL := `SET DEFAULT ROLE r_3 TO u_1;`
+ _, err := tk.Exec(setRoleSQL)
+ c.Check(err, NotNil)
+
+ setRoleSQL = `SET DEFAULT ROLE r_1 TO u_1000;`
+ _, err = tk.Exec(setRoleSQL)
+ c.Check(err, NotNil)
+
+ setRoleSQL = `SET DEFAULT ROLE r_1, r_3 TO u_1;`
+ _, err = tk.Exec(setRoleSQL)
+ c.Check(err, NotNil)
+
+ setRoleSQL = `SET DEFAULT ROLE r_1 TO u_1;`
+ _, err = tk.Exec(setRoleSQL)
+ c.Check(err, IsNil)
+ result := tk.MustQuery(`SELECT DEFAULT_ROLE_USER FROM mysql.default_roles WHERE USER="u_1"`)
+ result.Check(testkit.Rows("r_1"))
+ setRoleSQL = `SET DEFAULT ROLE r_2 TO u_1;`
+ _, err = tk.Exec(setRoleSQL)
+ c.Check(err, IsNil)
+ result = tk.MustQuery(`SELECT DEFAULT_ROLE_USER FROM mysql.default_roles WHERE USER="u_1"`)
+ result.Check(testkit.Rows("r_2"))
+
+ setRoleSQL = `SET DEFAULT ROLE ALL TO u_1;`
+ _, err = tk.Exec(setRoleSQL)
+ c.Check(err, IsNil)
+ result = tk.MustQuery(`SELECT DEFAULT_ROLE_USER FROM mysql.default_roles WHERE USER="u_1"`)
+ result.Check(testkit.Rows("r_1", "r_2"))
+
+ setRoleSQL = `SET DEFAULT ROLE NONE TO u_1;`
+ _, err = tk.Exec(setRoleSQL)
+ c.Check(err, IsNil)
+ result = tk.MustQuery(`SELECT DEFAULT_ROLE_USER FROM mysql.default_roles WHERE USER="u_1"`)
+ result.Check(nil)
+
+ dropRoleSQL := `DROP USER r_1, r_2, r_3, u_1;`
+ tk.MustExec(dropRoleSQL)
+}
+
 func (s *testSuite3) TestUser(c *C) {
  tk := testkit.NewTestKit(c, s.store)
  // Make sure user test not in mysql.User.

planner/core/planbuilder.go

@@ -268,7 +268,7 @@ func (b *PlanBuilder) Build(node ast.Node) (Plan, error) {
  case *ast.BinlogStmt, *ast.FlushStmt, *ast.UseStmt,
  *ast.BeginStmt, *ast.CommitStmt, *ast.RollbackStmt, *ast.CreateUserStmt, *ast.SetPwdStmt,
  *ast.GrantStmt, *ast.DropUserStmt, *ast.AlterUserStmt, *ast.RevokeStmt, *ast.KillStmt, *ast.DropStatsStmt,
- *ast.GrantRoleStmt, *ast.RevokeRoleStmt, *ast.SetRoleStmt:
+ *ast.GrantRoleStmt, *ast.RevokeRoleStmt, *ast.SetRoleStmt, *ast.SetDefaultRoleStmt:
  return b.buildSimple(node.(ast.StmtNode))
  case ast.DDLNode:
  return b.buildDDL(x)
@@ -1095,7 +1095,7 @@ func (b *PlanBuilder) buildSimple(node ast.StmtNode) (Plan, error) {
  err := ErrSpecificAccessDenied.GenWithStackByArgs("CREATE USER")
  b.visitInfo = appendVisitInfo(b.visitInfo, mysql.CreateUserPriv, "", "", "", err)
  }
- case *ast.AlterUserStmt:
+ case *ast.AlterUserStmt, *ast.SetDefaultRoleStmt:
  err := ErrSpecificAccessDenied.GenWithStackByArgs("CREATE USER")
  b.visitInfo = appendVisitInfo(b.visitInfo, mysql.CreateUserPriv, "", "", "", err)
  case *ast.GrantStmt:

privilege/privilege.go

@@ -56,6 +56,12 @@ type Manager interface {
  // ActiveRoles active roles for current session.
  // The first illegal role will be returned.
  ActiveRoles(ctx sessionctx.Context, roleList []*auth.RoleIdentity) (bool, string)
+
+ // FindEdge find if there is an edge between role and user.
+ FindEdge(ctx sessionctx.Context, role *auth.RoleIdentity, user *auth.UserIdentity) bool
+
+ // GetDefaultRoles returns all default roles for certain user.
+ GetDefaultRoles(user, host string) []*auth.RoleIdentity
 }
 
 const key keyType = 0

privilege/privileges/cache.go

@@ -68,7 +68,7 @@ type dbRecord struct {
  User string
  Privileges mysql.PrivilegeType
 
- // patChars is compiled from Host and DB, cached for pattern match performance.
+ // hostPatChars is compiled from Host and DB, cached for pattern match performance.
  hostPatChars []byte
  hostPatTypes []byte
 
@@ -105,7 +105,19 @@ type columnsPrivRecord struct {
  patTypes []byte
 }
 
-// RoleGraphEdgesTable is used to cache relationship between and role.
+// defaultRoleRecord is used to cache mysql.default_roles
+type defaultRoleRecord struct {
+ Host string
+ User string
+ DefaultRoleUser string
+ DefaultRoleHost string
+
+ // patChars is compiled from Host, cached for pattern match performance.
+ patChars []byte
+ patTypes []byte
+}
+
+// roleGraphEdgesTable is used to cache relationship between and role.
 type roleGraphEdgesTable struct {
  roleList map[string]bool
 }
@@ -125,11 +137,12 @@ func (g roleGraphEdgesTable) Find(user, host string) bool {
 
 // MySQLPrivilege is the in-memory cache of mysql privilege tables.
 type MySQLPrivilege struct {
- User []UserRecord
- DB []dbRecord
- TablesPriv []tablesPrivRecord
- ColumnsPriv []columnsPrivRecord
- RoleGraph map[string]roleGraphEdgesTable
+ User []UserRecord
+ DB []dbRecord
+ TablesPriv []tablesPrivRecord
+ ColumnsPriv []columnsPrivRecord
+ DefaultRoles []defaultRoleRecord
+ RoleGraph map[string]roleGraphEdgesTable
 }
 
 // FindRole is used to detect whether there is edges between users and roles.
@@ -166,6 +179,14 @@ func (p *MySQLPrivilege) LoadAll(ctx sessionctx.Context) error {
  log.Warn("mysql.tables_priv missing")
  }
 
+ err = p.LoadDefaultRoles(ctx)
+ if err != nil {
+ if !noSuchTable(err) {
+ return errors.Trace(err)
+ }
+ log.Warn("mysql.default_roles missing")
+ }
+
  err = p.LoadColumnsPrivTable(ctx)
  if err != nil {
  if !noSuchTable(err) {
@@ -316,6 +337,11 @@ func (p *MySQLPrivilege) LoadColumnsPrivTable(ctx sessionctx.Context) error {
  return p.loadTable(ctx, "select HIGH_PRIORITY Host,DB,User,Table_name,Column_name,Timestamp,Column_priv from mysql.columns_priv", p.decodeColumnsPrivTableRow)
 }
 
+// LoadDefaultRoles loads the mysql.columns_priv table from database.
+func (p *MySQLPrivilege) LoadDefaultRoles(ctx sessionctx.Context) error {
+ return p.loadTable(ctx, "select HOST, USER, DEFAULT_ROLE_HOST, DEFAULT_ROLE_USER from mysql.default_roles", p.decodeDefaultRoleTableRow)
+}
+
 func (p *MySQLPrivilege) loadTable(sctx sessionctx.Context, sql string,
  decodeTableRow func(chunk.Row, []*ast.ResultField) error) error {
  ctx := context.Background()
@@ -455,6 +481,25 @@ func (p *MySQLPrivilege) decodeRoleEdgesTable(row chunk.Row, fs []*ast.ResultFie
  return nil
 }
 
+func (p *MySQLPrivilege) decodeDefaultRoleTableRow(row chunk.Row, fs []*ast.ResultField) error {
+ var value defaultRoleRecord
+ for i, f := range fs {
+ switch {
+ case f.ColumnAsName.L == "host":
+ value.Host = row.GetString(i)
+ value.patChars, value.patTypes = stringutil.CompilePattern(value.Host, '\\')
+ case f.ColumnAsName.L == "user":
+ value.User = row.GetString(i)
+ case f.ColumnAsName.L == "default_role_host":
+ value.DefaultRoleHost = row.GetString(i)
+ case f.ColumnAsName.L == "default_role_user":
+ value.DefaultRoleUser = row.GetString(i)
+ }
+ }
+ p.DefaultRoles = append(p.DefaultRoles, value)
+ return nil
+}
+
 func (p *MySQLPrivilege) decodeColumnsPrivTableRow(row chunk.Row, fs []*ast.ResultField) error {
  var value columnsPrivRecord
  for i, f := range fs {
@@ -522,6 +567,10 @@ func (record *columnsPrivRecord) match(user, host, db, table, col string) bool {
  patternMatch(host, record.patChars, record.patTypes)
 }
 
+func (record *defaultRoleRecord) match(user, host string) bool {
+ return record.User == user && patternMatch(host, record.patChars, record.patTypes)
+}
+
 // patternMatch matches "%" the same way as ".*" in regular expression, for example,
 // "10.0.%" would match "10.0.1" "10.0.1.118" ...
 func patternMatch(str string, patChars, patTypes []byte) bool {
@@ -766,6 +815,16 @@ func appendUserPrivilegesTableRow(rows [][]types.Datum, user UserRecord) [][]typ
  return rows
 }
 
+func (p *MySQLPrivilege) getDefaultRoles(user, host string) []*auth.RoleIdentity {
+ ret := make([]*auth.RoleIdentity, 0)
+ for _, r := range p.DefaultRoles {
+ if r.match(user, host) {
+ ret = append(ret, &auth.RoleIdentity{Username: r.DefaultRoleUser, Hostname: r.DefaultRoleHost})
+ }
+ }
+ return ret
+}
+
 // Handle wraps MySQLPrivilege providing thread safe access.
 type Handle struct {
  priv atomic.Value

privilege/privileges/cache_test.go

@@ -134,6 +134,25 @@ func (s *testCacheSuite) TestLoadColumnsPrivTable(c *C) {
  c.Assert(p.ColumnsPriv[1].ColumnPriv, Equals, mysql.SelectPriv)
 }
 
+func (s *testCacheSuite) TestLoadDefaultRoleTable(c *C) {
+ se, err := session.CreateSession4Test(s.store)
+ c.Assert(err, IsNil)
+ defer se.Close()
+ mustExec(c, se, "use mysql;")
+ mustExec(c, se, "truncate table default_roles")
+
+ mustExec(c, se, `INSERT INTO mysql.default_roles VALUES ("%", "test_default_roles", "localhost", "r_1")`)
+ mustExec(c, se, `INSERT INTO mysql.default_roles VALUES ("%", "test_default_roles", "localhost", "r_2")`)
+ var p privileges.MySQLPrivilege
+ err = p.LoadDefaultRoles(se)
+ c.Assert(err, IsNil)
+ c.Assert(p.DefaultRoles[0].Host, Equals, `%`)
+ c.Assert(p.DefaultRoles[0].User, Equals, "test_default_roles")
+ c.Assert(p.DefaultRoles[0].DefaultRoleHost, Equals, "localhost")
+ c.Assert(p.DefaultRoles[0].DefaultRoleUser, Equals, "r_1")
+ c.Assert(p.DefaultRoles[1].DefaultRoleHost, Equals, "localhost")
+}
+
 func (s *testCacheSuite) TestPatternMatch(c *C) {
  se, err := session.CreateSession4Test(s.store)
  c.Assert(err, IsNil)