From 1949335ab6596dd6e2654f89f741dff6ced38eaa Mon Sep 17 00:00:00 2001
From: Lingyu Song <songlingyu@pingcap.com>
Date: Wed, 5 Jun 2019 10:44:28 +0800
Subject: [PATCH] privileges: add SkipWithGrant check for RBAC methods (#10681)

---
 privilege/privileges/privileges.go | 16 ++++++++++++++++
 session/session_test.go            |  5 ++++-
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/privilege/privileges/privileges.go b/privilege/privileges/privileges.go
index b9a74b09df24f..3e1e950454e17 100644
--- a/privilege/privileges/privileges.go
+++ b/privilege/privileges/privileges.go
@@ -185,6 +185,9 @@ func (p *UserPrivileges) UserPrivilegesTable() [][]types.Datum {
 
 // ShowGrants implements privilege.Manager ShowGrants interface.
 func (p *UserPrivileges) ShowGrants(ctx sessionctx.Context, user *auth.UserIdentity, roles []*auth.RoleIdentity) (grants []string, err error) {
+	if SkipWithGrant {
+		return nil, errNonexistingGrant.GenWithStackByArgs("root", "%")
+	}
 	mysqlPrivilege := p.Handle.Get()
 	u := user.Username
 	h := user.Hostname
@@ -202,6 +205,9 @@ func (p *UserPrivileges) ShowGrants(ctx sessionctx.Context, user *auth.UserIdent
 
 // ActiveRoles implements privilege.Manager ActiveRoles interface.
 func (p *UserPrivileges) ActiveRoles(ctx sessionctx.Context, roleList []*auth.RoleIdentity) (bool, string) {
+	if SkipWithGrant {
+		return true, ""
+	}
 	mysqlPrivilege := p.Handle.Get()
 	u := p.user
 	h := p.host
@@ -218,6 +224,9 @@ func (p *UserPrivileges) ActiveRoles(ctx sessionctx.Context, roleList []*auth.Ro
 
 // FindEdge implements privilege.Manager FindRelationship interface.
 func (p *UserPrivileges) FindEdge(ctx sessionctx.Context, role *auth.RoleIdentity, user *auth.UserIdentity) bool {
+	if SkipWithGrant {
+		return false
+	}
 	mysqlPrivilege := p.Handle.Get()
 	ok := mysqlPrivilege.FindRole(user.Username, user.Hostname, role)
 	if !ok {
@@ -229,6 +238,9 @@ func (p *UserPrivileges) FindEdge(ctx sessionctx.Context, role *auth.RoleIdentit
 
 // GetDefaultRoles returns all default roles for certain user.
 func (p *UserPrivileges) GetDefaultRoles(user, host string) []*auth.RoleIdentity {
+	if SkipWithGrant {
+		return make([]*auth.RoleIdentity, 0, 10)
+	}
 	mysqlPrivilege := p.Handle.Get()
 	ret := mysqlPrivilege.getDefaultRoles(user, host)
 	return ret
@@ -236,6 +248,10 @@ func (p *UserPrivileges) GetDefaultRoles(user, host string) []*auth.RoleIdentity
 
 // GetAllRoles return all roles of user.
 func (p *UserPrivileges) GetAllRoles(user, host string) []*auth.RoleIdentity {
+	if SkipWithGrant {
+		return make([]*auth.RoleIdentity, 0, 10)
+	}
+
 	mysqlPrivilege := p.Handle.Get()
 	return mysqlPrivilege.getAllRoles(user, host)
 }
diff --git a/session/session_test.go b/session/session_test.go
index 7cc34eab65959..600cb4c9fb217 100644
--- a/session/session_test.go
+++ b/session/session_test.go
@@ -702,7 +702,10 @@ func (s *testSessionSuite) TestSkipWithGrant(c *C) {
 	c.Assert(tk.Se.Auth(&auth.UserIdentity{Username: "xxx", Hostname: `%`}, []byte("yyy"), []byte("zzz")), IsTrue)
 	c.Assert(tk.Se.Auth(&auth.UserIdentity{Username: "root", Hostname: `%`}, []byte(""), []byte("")), IsTrue)
 	tk.MustExec("create table t (id int)")
-
+	tk.MustExec("create role r_1")
+	tk.MustExec("grant r_1 to root")
+	tk.MustExec("set role all")
+	tk.MustExec("show grants for root")
 	privileges.SkipWithGrant = save2
 }