POC planbuilder authorization

Author: Hydrocharged

go.mod

@@ -43,4 +43,6 @@ require (
 	gopkg.in/check.v1 v1.0.0-20200902074654-038fdea0a05b // indirect
 )
 
+replace github.com/dolthub/vitess => ../vitess
+
 go 1.22.2

sql/planbuilder/auth.go

@@ -0,0 +1,103 @@
+// Copyright 2023 Dolthub, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package planbuilder
+
+import (
+	"fmt"
+	"github.com/dolthub/go-mysql-server/sql"
+	"github.com/dolthub/go-mysql-server/sql/mysql_db"
+	"github.com/dolthub/vitess/go/mysql"
+	ast "github.com/dolthub/vitess/go/vt/sqlparser"
+)
+
+// TODO: doc
+var Authorization = func(b *Builder, auth ast.AuthInformation) {
+	// TODO: expose necessary stuff from Builder to public
+	ctx := b.ctx
+	// TODO: database loading shouldn't be done for every call, this needs to be cached for each Parse call in some way
+	db, err := b.cat.Database(ctx, "mysql")
+	if err != nil {
+		b.handleErr(err)
+	}
+	mysqlDb, ok := db.(*mysql_db.MySQLDb)
+	if !ok {
+		b.handleErr(fmt.Errorf("FOR TESTING: could not load the `mysql` database")) // Check if this is likely
+	}
+	if !mysqlDb.Enabled() {
+		return
+	}
+	client := ctx.Session.Client()
+	user := func() *mysql_db.User {
+		rd := mysqlDb.Reader()
+		defer rd.Close()
+		return mysqlDb.GetUser(rd, client.User, client.Address, false)
+	}()
+	if user == nil {
+		// TODO: Builder should have a HandleAuthErr
+		b.handleErr(mysql.NewSQLError(mysql.ERAccessDeniedError, mysql.SSAccessDeniedError, "Access denied for user '%v'", ctx.Session.Client().User))
+	}
+
+	var privilegeTypes []sql.PrivilegeType
+	switch auth.AuthType {
+	case ast.AuthType_IGNORE:
+		// This means that authorization is being handled elsewhere, and should be ignored here
+		return
+	case ast.AuthType_DELETE:
+		privilegeTypes = []sql.PrivilegeType{sql.PrivilegeType_Delete}
+	case ast.AuthType_INSERT:
+		privilegeTypes = []sql.PrivilegeType{sql.PrivilegeType_Insert}
+	case ast.AuthType_REPLACE:
+		privilegeTypes = []sql.PrivilegeType{sql.PrivilegeType_Insert, sql.PrivilegeType_Delete}
+	case ast.AuthType_SELECT:
+		privilegeTypes = []sql.PrivilegeType{sql.PrivilegeType_Select}
+	case ast.AuthType_UPDATE:
+		privilegeTypes = []sql.PrivilegeType{sql.PrivilegeType_Update}
+	default:
+		b.handleErr(fmt.Errorf("FOR TESTING: default case hit for AuthType"))
+	}
+
+	hasPrivileges := true
+	switch auth.TargetType {
+	case ast.AuthTargetType_SingleTableIdentifier:
+		subject := sql.PrivilegeCheckSubject{
+			Database: authDatabaseName(ctx, auth.TargetNames[0]),
+			Table:    auth.TargetNames[1],
+		}
+		hasPrivileges = mysqlDb.UserHasPrivileges(ctx, sql.NewPrivilegedOperation(subject, privilegeTypes...))
+	case ast.AuthTargetType_MultipleTableIdentifiers:
+		for i := 0; i < len(auth.TargetNames); i += 2 {
+			subject := sql.PrivilegeCheckSubject{
+				Database: authDatabaseName(ctx, auth.TargetNames[i]),
+				Table:    auth.TargetNames[i+1],
+			}
+			hasPrivileges = hasPrivileges && mysqlDb.UserHasPrivileges(ctx, sql.NewPrivilegedOperation(subject, privilegeTypes...))
+		}
+	default:
+		b.handleErr(fmt.Errorf("FOR TESTING: default case hit for TargetType"))
+	}
+
+	if !hasPrivileges {
+		b.handleErr(sql.ErrPrivilegeCheckFailed.New(user.UserHostToString("'")))
+	}
+}
+
+// authDatabaseName uses the current database from the context if a database is not specified, otherwise it returns the
+// given database name.
+func authDatabaseName(ctx *sql.Context, dbName string) string {
+	if len(dbName) == 0 {
+		return ctx.GetCurrentDatabase()
+	}
+	return dbName
+}

sql/planbuilder/dml.go

@@ -152,6 +152,7 @@ func (b *Builder) buildInsert(inScope *scope, i *ast.Insert) (outScope *scope) {
 		checks := b.loadChecksFromTable(destScope, rt.Table)
 		outScope.node = ins.WithChecks(checks)
 	}
+	Authorization(b, i.Auth)
 
 	return
 }
@@ -475,6 +476,7 @@ func (b *Builder) buildDelete(inScope *scope, d *ast.Delete) (outScope *scope) {
 	del.RefsSingleRel = !outScope.refsSubquery
 	del.IsProcNested = b.ProcCtx().DbName != ""
 	outScope.node = del
+	Authorization(b, d.Auth)
 	return
 }
 

sql/planbuilder/from.go

@@ -383,6 +383,7 @@ func (b *Builder) buildDataSource(inScope *scope, te ast.TableExpr) (outScope *s
 		default:
 			b.handleErr(sql.ErrUnsupportedSyntax.New(ast.String(te)))
 		}
+		Authorization(b, t.Auth)
 
 	case *ast.TableFuncExpr:
 		return b.buildTableFunc(inScope, t)