Repository: pingcap/tidb
File: pkg/privilege/privileges/cache.go (Go)

Code before:

type columnStr = string
type columnStrs = []columnStr
type privOnColumns = map[mysql.PrivilegeType]columnStrs

func privOnColumnsToString(p privOnColumns) string {
	var buf bytes.Buffer
	idx := 0
	for _, priv := range mysql.AllColumnPrivs {
		v, ok := p[priv]
		if !ok || len(v) == 0 {
			continue
		}

		if idx > 0 {
			buf.WriteString(", ")
		}
		privStr := PrivToString(priv, mysql.AllColumnPrivs, mysql.Priv2Str)
		fmt.Fprintf(&buf, "%s(", privStr)
		for i, col := range v {
			if i > 0 {
				fmt.Fprintf(&buf, ", ")
			}
			buf.WriteString(col)
		}
		buf.WriteString(")")
		idx++
	}
	return buf.String()
}

func collectColumnGrant(record *columnsPrivRecord, user, host string, columnPrivTable map[string]privOnColumns, sqlMode mysql.SQLMode) bool {
	if record.baseRecord.match(user, host) {
		recordKey := stringutil.Escape(record.DB, sqlMode) + "." + stringutil.Escape(record.TableName, sqlMode)

		privColumns, ok := columnPrivTable[recordKey]
		if !ok {
			privColumns = make(map[mysql.PrivilegeType]columnStrs)
		}

		for _, priv := range mysql.AllColumnPrivs {
			if priv&record.ColumnPriv > 0 {
				old := privColumns[priv]
				privColumns[priv] = append(old, record.ColumnName)
				columnPrivTable[recordKey] = privColumns
			}
		}
		return true
	}
	return false
}

Review comment (security):
Well, it's ok here in 8.5 that we still have full privilege data in-memory.
But in 9.0 later, access the whole data might cost a lot, considing if there is 1M users.
How to cherry-pick would be an issue later.

Code after:

func (p *MySQLPrivilege) fetchColumnPrivileges(users, hosts []string) (rows [][]types.Datum) {
	rows = make([][]types.Datum, 0, len(users))
	// We use a 2-level map to record if a `username`@`hostname` is grantable on table/column level
	// grantableUser records: username -> hostname -> grantable
	grantableUser := make(map[string]map[string]struct{}, 4)
	isGrantable := func(username, hostname string) string {
		if hostNames, exists := grantableUser[username]; exists {
			if _, exists = hostNames[hostname]; exists {
				return "YES"
			}
		}
		return "NO"
	}

	for i := range users {
		userName, hostName := users[i], hosts[i]
		columnsPrivRecords, exists := p.ColumnsPrivMap[userName]
		if !exists {
			continue
		}

		// need to get "grantable" information from mysql.tables_priv
		if _, exists := grantableUser[userName]; !exists {
			grantableUser[userName] = make(map[string]struct{}, 1)
			if tablesPrivRecords, exists := p.TablesPrivMap[userName]; exists {
				for _, tablesPrivRecord := range tablesPrivRecords {
					if tablesPrivRecord.TablePriv&mysql.GrantPriv > 0 {
						grantableUser[tablesPrivRecord.User][tablesPrivRecord.Host] = struct{}{}
					}
				}
			}
		}

		for _, columnsPrivRecord := range columnsPrivRecords {
			if columnsPrivRecord.Host != hostName {
				continue
			}

			for _, columnPriv := range mysql.AllColumnPrivs {
				if columnsPrivRecord.ColumnPriv&columnPriv > 0 {
					rows = append(rows, types.MakeDatums(
						fmt.Sprintf("'%s'@'%s'", userName, hostName),
						"def",
						columnsPrivRecord.DB,
						columnsPrivRecord.TableName,
						columnsPrivRecord.ColumnName,
						strings.ToUpper(columnPriv.String()),
						isGrantable(userName, hostName),
					))
				}