Skip to content

Commit

Permalink
privilege: using system session to execute internal sql of RBAC (#13820
Browse files Browse the repository at this point in the history
…) (#13890)
  • Loading branch information
Lingyu Song authored and sre-bot committed Dec 5, 2019
1 parent ab73fde commit 251574b
Showing 1 changed file with 40 additions and 16 deletions.
56 changes: 40 additions & 16 deletions executor/simple.go
Original file line number Diff line number Diff line change
Expand Up @@ -132,7 +132,12 @@ func (e *SimpleExec) Next(ctx context.Context, req *chunk.Chunk) (err error) {
}

func (e *SimpleExec) setDefaultRoleNone(s *ast.SetDefaultRoleStmt) error {
sqlExecutor := e.ctx.(sqlexec.SQLExecutor)
restrictedCtx, err := e.getSysSession()
if err != nil {
return err
}
defer e.releaseSysSession(restrictedCtx)
sqlExecutor := restrictedCtx.(sqlexec.SQLExecutor)
if _, err := sqlExecutor.Execute(context.Background(), "begin"); err != nil {
return err
}
Expand Down Expand Up @@ -174,7 +179,13 @@ func (e *SimpleExec) setDefaultRoleRegular(s *ast.SetDefaultRoleStmt) error {
return ErrCannotUser.GenWithStackByArgs("SET DEFAULT ROLE", role.String())
}
}
sqlExecutor := e.ctx.(sqlexec.SQLExecutor)

restrictedCtx, err := e.getSysSession()
if err != nil {
return err
}
defer e.releaseSysSession(restrictedCtx)
sqlExecutor := restrictedCtx.(sqlexec.SQLExecutor)
if _, err := sqlExecutor.Execute(context.Background(), "begin"); err != nil {
return err
}
Expand Down Expand Up @@ -288,7 +299,6 @@ func (e *SimpleExec) setDefaultRoleForCurrentUser(s *ast.SetDefaultRoleStmt) (er
return err
}
defer e.releaseSysSession(restrictedCtx)

sqlExecutor := restrictedCtx.(sqlexec.SQLExecutor)

if _, err := sqlExecutor.Execute(context.Background(), "begin"); err != nil {
Expand Down Expand Up @@ -551,8 +561,15 @@ func (e *SimpleExec) executeRevokeRole(s *ast.RevokeRoleStmt) error {
}
}

restrictedCtx, err := e.getSysSession()
if err != nil {
return err
}
defer e.releaseSysSession(restrictedCtx)
sqlExecutor := restrictedCtx.(sqlexec.SQLExecutor)

// begin a transaction to insert role graph edges.
if _, err := e.ctx.(sqlexec.SQLExecutor).Execute(context.Background(), "begin"); err != nil {
if _, err := sqlExecutor.Execute(context.Background(), "begin"); err != nil {
return errors.Trace(err)
}
for _, user := range s.Users {
Expand All @@ -561,7 +578,7 @@ func (e *SimpleExec) executeRevokeRole(s *ast.RevokeRoleStmt) error {
return errors.Trace(err)
}
if !exists {
if _, err := e.ctx.(sqlexec.SQLExecutor).Execute(context.Background(), "rollback"); err != nil {
if _, err := sqlExecutor.Execute(context.Background(), "rollback"); err != nil {
return errors.Trace(err)
}
return ErrCannotUser.GenWithStackByArgs("REVOKE ROLE", user.String())
Expand All @@ -571,22 +588,22 @@ func (e *SimpleExec) executeRevokeRole(s *ast.RevokeRoleStmt) error {
role.Hostname = "%"
}
sql := fmt.Sprintf(`DELETE IGNORE FROM %s.%s WHERE FROM_HOST='%s' and FROM_USER='%s' and TO_HOST='%s' and TO_USER='%s'`, mysql.SystemDB, mysql.RoleEdgeTable, role.Hostname, role.Username, user.Hostname, user.Username)
if _, err := e.ctx.(sqlexec.SQLExecutor).Execute(context.Background(), sql); err != nil {
if _, err := e.ctx.(sqlexec.SQLExecutor).Execute(context.Background(), "rollback"); err != nil {
if _, err := sqlExecutor.Execute(context.Background(), sql); err != nil {
if _, err := sqlExecutor.Execute(context.Background(), "rollback"); err != nil {
return errors.Trace(err)
}
return ErrCannotUser.GenWithStackByArgs("REVOKE ROLE", role.String())
}
sql = fmt.Sprintf(`DELETE IGNORE FROM %s.%s WHERE DEFAULT_ROLE_HOST='%s' and DEFAULT_ROLE_USER='%s' and HOST='%s' and USER='%s'`, mysql.SystemDB, mysql.DefaultRoleTable, role.Hostname, role.Username, user.Hostname, user.Username)
if _, err := e.ctx.(sqlexec.SQLExecutor).Execute(context.Background(), sql); err != nil {
if _, err := e.ctx.(sqlexec.SQLExecutor).Execute(context.Background(), "rollback"); err != nil {
if _, err := sqlExecutor.Execute(context.Background(), sql); err != nil {
if _, err := sqlExecutor.Execute(context.Background(), "rollback"); err != nil {
return errors.Trace(err)
}
return ErrCannotUser.GenWithStackByArgs("REVOKE ROLE", role.String())
}
}
}
if _, err := e.ctx.(sqlexec.SQLExecutor).Execute(context.Background(), "commit"); err != nil {
if _, err := sqlExecutor.Execute(context.Background(), "commit"); err != nil {
return err
}
domain.GetDomain(e.ctx).NotifyUpdatePrivilege(e.ctx)
Expand Down Expand Up @@ -760,29 +777,36 @@ func (e *SimpleExec) executeGrantRole(s *ast.GrantRoleStmt) error {
}
}

restrictedCtx, err := e.getSysSession()
if err != nil {
return err
}
defer e.releaseSysSession(restrictedCtx)
sqlExecutor := restrictedCtx.(sqlexec.SQLExecutor)

// begin a transaction to insert role graph edges.
if _, err := e.ctx.(sqlexec.SQLExecutor).Execute(context.Background(), "begin"); err != nil {
if _, err := sqlExecutor.Execute(context.Background(), "begin"); err != nil {
return err
}

for _, user := range s.Users {
for _, role := range s.Roles {
sql := fmt.Sprintf(`INSERT IGNORE INTO %s.%s (FROM_HOST, FROM_USER, TO_HOST, TO_USER) VALUES ('%s','%s','%s','%s')`, mysql.SystemDB, mysql.RoleEdgeTable, role.Hostname, role.Username, user.Hostname, user.Username)
if _, err := e.ctx.(sqlexec.SQLExecutor).Execute(context.Background(), sql); err != nil {
if _, err := sqlExecutor.Execute(context.Background(), sql); err != nil {
failedUsers = append(failedUsers, user.String())
logutil.Logger(context.Background()).Error(fmt.Sprintf("Error occur when executing %s", sql))
if _, err := e.ctx.(sqlexec.SQLExecutor).Execute(context.Background(), "rollback"); err != nil {
if _, err := sqlExecutor.Execute(context.Background(), "rollback"); err != nil {
return err
}
return ErrCannotUser.GenWithStackByArgs("GRANT ROLE", user.String())
}
}
}
if _, err := e.ctx.(sqlexec.SQLExecutor).Execute(context.Background(), "commit"); err != nil {
if _, err := sqlExecutor.Execute(context.Background(), "commit"); err != nil {
return err
}
err := domain.GetDomain(e.ctx).PrivilegeHandle().Update(e.ctx.(sessionctx.Context))
return err
domain.GetDomain(e.ctx).NotifyUpdatePrivilege(e.ctx)
return nil
}

func (e *SimpleExec) executeDropUser(s *ast.DropUserStmt) error {
Expand Down

0 comments on commit 251574b

Please sign in to comment.