extension

package
v1.1.0-beta.0...-68ac9ec Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Dec 10, 2024 License: Apache-2.0 Imports: 17 Imported by: 0

README

Extension Framework Usage

Introduction

The Extension Framework facilitates the seamless integration of standalone code with TiDB. This framework is useful when:

  • You aim to extend TiDB functionality but prefer to encapsulate new code within a standalone package, minimizing deep integration.
  • Considering confidentiality, you wish to extend TiDB without exposing the code publicly.

Usage Guidelines

A Simple Example

The subsequent example demonstrates how to log information when a new connection is established:

// pkg/extension/example/example.go

package example

import (
	"github.com/pingcap/tidb/pkg/extension"
	"github.com/pingcap/tidb/pkg/parser/terror"
	"github.com/pingcap/tidb/pkg/util/logutil"
	"go.uber.org/zap"
)

func createSessionHandler() *extension.SessionHandler {
	return &extension.SessionHandler{
		OnConnectionEvent: func(tp extension.ConnEventTp, info *extension.ConnEventInfo) {
			if tp == extension.ConnConnected {
				logutil.BgLogger().Info("new connection connected", zap.String("client IP", info.ClientIP))
			}
		},
	}
}

func init() {
	err := extension.Register(
		"example",
		extension.WithSessionHandlerFactory(createSessionHandler),
	)
	terror.MustNil(err)
}

// pkg/extension/_import/import_example.go

package extensionimport

import (
	_ "github.com/pingcap/tidb/pkg/extension/example"
)

Incorporate these files into your TiDB file structure, rebuild, and run the TiDB server. Observe the enhanced logging when a new connection attempts to connect to TiDB:

[2023/11/14 16:22:26.854 +08:00] [INFO] [example.go:28] ["new connection connected"] ["client IP"=127.0.0.1]

Let's delve into a step-by-step explanation of how this works:

In the primary file, example.go, a new function named createSessionHandler is introduced. This function acts as a factory and is intended to be invoked when a new connection is established. It returns an object of type extension.SessionHandler to guide the framework on handling the connection. The SessionHandler type provides various customizable fields. In the provided example, we utilize the OnConnectionEvent field to manage the extension.ConnConnected event for connections, logging the client's IP address.

The custom code is registered using extension.Register. In the init function of the example.go file, the first parameter passed to extension.Register is the extension name. Any string can be used as the name, with the only constraint being its uniqueness across all extensions. Following the extension name, an arbitrary number of options can be passed to extension.Register. In our example, we employ the extension.WithSessionHandlerFactory option to convey the createSessionHandler function to the previously defined framework.

It's essential to note that example.go resides in a newly created package. If this package is not imported by other packages, the init function will not be called. To address this, a new file, example_import.go, is created within an existing package, specifically pkg/extension/_import. This file imports pkg/extension/example, ensuring that the extension is registered when TiDB starts.

Further Explanation about Session Handler

The SessionHandler is designed to manage session events and is declared as follows:

// SessionHandler is used to listen session events
type SessionHandler struct {
	OnConnectionEvent func(ConnEventTp, *ConnEventInfo)
	OnStmtEvent       func(StmtEventTp, StmtEventInfo)
}

SessionHandler offers several customizable fields. These fields default to empty, implying that if left unset, the framework will not perform any actions when corresponding events occur.

OnConnectionEvent

The OnConnectionEvent function is used to observe all connection events. The types of connection events are defined by ConnEventTp, with key types including:

  • ConnConnected: the creation of a new connection that is not yet authenticated.
  • ConnHandshakeAccepted: successful authentication of a connection.
  • ConnHandshakeRejected: rejection due to authentication failure.
  • ConnClose: the close of a connection.

It is important to note that OnConnectionEvent does not return errors; it is designed as a listener, and any encountered errors in custom code must be managed before the function concludes. In essence, OnConnectionEvent does not disrupt the connection flow but merely observes.

OnStmtEvent

OnStmtEvent serves to monitor events related to statements. Similar to OnConnectionEvent, it functions as a listener without influencing the statement flow. Statement event types, defined by StmtEventTp, include:

  • StmtError: Denoting the completion of a statement with an error.
  • StmtSuccess: Indicating the successful completion of a statement.
Registering Extensions with Options

In addition to extension.WithSessionHandlerFactory, various options are available for extension customization. All options can be registered using the extension.Register function. For instance:

	err := extension.Register(
		"example",
		extension.WithSessionHandlerFactory(createSessionHandler),
		extension.WithBootstrapSQL("CREATE TABLE IF NOT EXISTS mysql.example(a int)"),
		// ...
	)
	terror.MustNil(err)

Alternatively, use extension.RegisterFactory for dynamic registration of extension options. This proves beneficial when global configurations influence the selection of options. For example:

	err := extension.RegisterFactory("example", func() ([]extension.Option, error) {
		cfg := config.GetGlobalConfig()
		var factory func() *extension.SessionHandler
		if cfg.SomeFlag {
			factory = createSessionHandler1
		} else {
			factory = createSessionHandler2
		}

		return []extension.Option{
			extension.WithSessionHandlerFactory(factory),
		}, nil
	})
	terror.MustNil(err)
Key Extension Options

Several important extension options are outlined below:

WithCustomSysVariables

The WithCustomSysVariables option registers custom system variables, accepting a slice of variable.SysVar for addition.

WithCustomDynPrivs

Use WithCustomDynPrivs to register custom dynamic privileges.

WithCustomFunctions

The WithCustomFunctions option registers custom functions.

AccessCheckFunc

The AccessCheckFunc option customizes the access check logic, enabling additional checks for table access.

WithSessionHandlerFactory

This option is instrumental in handling session events.

WithBootstrap

WithBootstrap customizes the bootstrap logic, receiving a function invoked during TiDB bootstrap. Note that WithBootstrapSQL and WithBootstrap are mutually exclusive, allowing the use of only one.

WithBootstrapSQL

WithBootstrapSQL customizes the bootstrap logic with a string containing SQL statements for bootstrapping. Note that WithBootstrapSQL and WithBootstrap are mutually exclusive, allowing the use of only one.

Authentication Plugin

It is possible to implement authentication plugins for TiDB using the extension framework. The framework allows TiDB users to create their own authentication and privilege verification schemes. See docs/design/2024-05-10-extension-authentication-plugin.md for design.

See below for a basic example:

import (
	"errors"

	"github.com/pingcap/tidb/pkg/extension"
	"github.com/pingcap/tidb/pkg/parser/mysql"
	"github.com/pingcap/tidb/pkg/parser/terror"
	"github.com/pingcap/tidb/pkg/sessionctx/variable"
)

func createAuthPlugin() *extension.AuthPlugin {
	return &extension.AuthPlugin{
		// Name of the plugin: `CREATE USER <username> IDENTIFIED WITH my_auth_plugin AS <pwd>`
		Name: "my_auth_plugin",
		// MySQL clients will use the `mysql_native_password` plugin
		RequiredClientSidePlugin: "mysql_native_password",
		AuthenticateUser: func(request extension.AuthenticateRequest) error {
			// Allow login as long as there is an input password
			if len(request.InputAuthString) > 0 {
				return nil
			}
			return errors.New("no password")
		},
		GenerateAuthString: func(pwd string) (string, bool) {
			// As long as password is not empty, allow it
			return pwd, pwd != ""
		},
		ValidateAuthString: func(pwdHash string) bool {
			// As long as password is not empty, allow it
			return pwdHash != ""
		},
	}
}

func init() {
	plugin := createAuthPlugin()
	// Can load a list of auth plugins
	authPlugins := []*extension.AuthPlugin{plugin}
	err := extension.Register(
		"extension_authentication_plugin",
		extension.WithCustomAuthPlugins(authPlugins),
		// Register the plugin name as a system variable
		extension.WithCustomSysVariables([]*variable.SysVar{
			{
				Scope: variable.ScopeGlobal,
				Name:  "extension_authentication_plugin",
				Value: mysql.AuthNativePassword,
				Type:  variable.TypeEnum,
				// Add `my_auth_plugin` and the names of other auth plugins here
				PossibleValues: []string{plugin.Name},
			},
		}),
	)
	terror.MustNil(err)
}

Once this auth plugin is loaded, users can run statements such as:

CREATE USER u1 IDENTIFIED WITH my_auth_plugin AS <pwd>;
CREATE USER u2 IDENTIFIED WITH my_auth_plugin BY <pwd>;

where the authentication and privilege checks of u1 and u2 will be done by the logic in my_auth_plugin.

For other attributes of extension.AuthPlugin, refer to extension/auth.go.

Documentation

Index

Constants

This section is empty.

Variables

View Source 
var RegisterDynamicPrivilege func(string) error

RegisterDynamicPrivilege is used to resolve dependency cycle

View Source 
var RegisterExtensionFunc func(*FunctionDef) error

RegisterExtensionFunc is to avoid dependency cycle

View Source 
var RemoveDynamicPrivilege func(string) bool

RemoveDynamicPrivilege is used to resolve dependency cycle

View Source 
var RemoveExtensionFunc func(string)

RemoveExtensionFunc is to avoid dependency cycle

Functions

func Register 

func Register(name string, options ...Option) error

Register registers a new extension with options

func RegisterFactory 

func RegisterFactory(name string, factory func() ([]Option, error)) error

RegisterFactory registers a new extension with a factory

func Reset 

func Reset()

Reset resets the registry. It is only used by test

func Setup 

func Setup() error

Setup setups extensions

Types

type AccessCheckFunc 

type AccessCheckFunc func(db, tbl, column string, priv mysql.PrivilegeType, sem bool) []string

AccessCheckFunc is a function that returns a dynamic privilege list for db/tbl/column access

type AuthPlugin 

type AuthPlugin struct {
	// Name is the name of the auth plugin. It will be registered as a system variable in TiDB which can be used inside the `CREATE USER ... IDENTIFIED WITH 'plugin_name'` statement.
	Name string

	// RequiredClientSidePlugin is the name of the client-side plugin required by the server-side plugin. It will be used to check if the client has the required plugin installed and require the client to use it if installed.
	// The user can require default MySQL plugins such as 'caching_sha2_password' or 'mysql_native_password'.
	// If this is empty then `AuthPlugin.Name` is used as the required client-side plugin.
	RequiredClientSidePlugin string

	// AuthenticateUser is called when a client connects to the server as a user and the server authenticates the user.
	// If an error is returned, the login attempt fails, otherwise it succeeds.
	// request: The request context for the authentication plugin to authenticate a user
	AuthenticateUser func(request AuthenticateRequest) error

	// GenerateAuthString is a function for user to implement customized ways to encode the password (e.g. hash/salt/clear-text). The returned string will be stored as the encoded password in the mysql.user table.
	// If the input password is considered as invalid, this should return an error.
	// pwd: User's input password in CREATE/ALTER USER statements in clear-text
	GenerateAuthString func(pwd string) (string, bool)

	// ValidateAuthString checks if the password hash stored in the mysql.user table or passed in from `IDENTIFIED AS` is valid.
	// This is called when retrieving an existing user to make sure the password stored is valid and not modified and make sure user is passing a valid password hash in `IDENTIFIED AS`.
	// pwdHash: hash of the password stored in the internal user table
	ValidateAuthString func(pwdHash string) bool

	// VerifyPrivilege is called for each user queries, and serves as an extra check for privileges for the user.
	// It will only be executed if the user has already been granted the privilege in SQL layer.
	// Returns true if user has the requested privilege.
	// request: The request context for the authorization plugin to authorize a user's static privilege
	VerifyPrivilege func(request VerifyStaticPrivRequest) bool

	// VerifyDynamicPrivilege is called for each user queries, and serves as an extra check for dynamic privileges for the user.
	// It will only be executed if the user has already been granted the dynamic privilege in SQL layer.
	// Returns true if user has the requested privilege.
	// request: The request context for the authorization plugin to authorize a user's dynamic privilege
	VerifyDynamicPrivilege func(request VerifyDynamicPrivRequest) bool
}

AuthPlugin contains attributes needed for an authentication plugin.

type AuthenticateRequest 

type AuthenticateRequest struct {
	// User The username in the connect attempt
	User string
	// StoredAuthString The user's auth string stored in mysql.user table
	StoredAuthString string
	// InputAuthString The user's auth string passed in from the connection attempt in bytes
	InputAuthString []byte
	// Salt Randomly generated salt for the current connection
	Salt []byte
	// ConnState The TLS connection state (contains the TLS certificate) if client is using TLS. It will be nil if the client is not using TLS
	ConnState *tls.ConnectionState
	// AuthConn Interface for the plugin to communicate with the client
	AuthConn conn.AuthConn
}

AuthenticateRequest contains the context for the authentication plugin to authenticate a user.

type BootstrapContext 

type BootstrapContext interface {
	context.Context
	// ExecuteSQL is used to execute a sql
	ExecuteSQL(ctx context.Context, sql string) ([]chunk.Row, error)
	// EtcdClient returns the etcd client
	EtcdClient() *clientv3.Client
	// SessionPool returns the session pool of domain
	SessionPool() SessionPool
}

BootstrapContext is the context used by extension in bootstrap

type ConnEventInfo 

type ConnEventInfo struct {
	*variable.ConnectionInfo
	SessionAlias string
	ActiveRoles  []*auth.RoleIdentity
	Error        error
}

ConnEventInfo is the connection info for the event

type ConnEventTp 

type ConnEventTp uint8

ConnEventTp is the type of the connection event 

const (
	// ConnConnected means connection connected, but not handshake yet
	ConnConnected ConnEventTp = iota
	// ConnHandshakeAccepted means connection is accepted after handshake
	ConnHandshakeAccepted
	// ConnHandshakeRejected means connections is rejected after handshake
	ConnHandshakeRejected
	// ConnReset means the connection is reset
	ConnReset
	// ConnDisconnected means the connection is disconnected
	ConnDisconnected
)

type Extensions 

type Extensions struct {
	// contains filtered or unexported fields
}

Extensions contains all extensions that have already setup

func GetExtensions 

func GetExtensions() (*Extensions, error)

GetExtensions returns all extensions after setup

func (*Extensions) Bootstrap 

func (es *Extensions) Bootstrap(ctx BootstrapContext) error

Bootstrap bootstraps all extensions

func (*Extensions) GetAccessCheckFuncs 

func (es *Extensions) GetAccessCheckFuncs() (funcs []AccessCheckFunc)

GetAccessCheckFuncs returns spec functions of the custom access check

func (*Extensions) GetAuthPlugins 

func (es *Extensions) GetAuthPlugins() map[string]*AuthPlugin

GetAuthPlugins returns the registered authentication plugins.

func (*Extensions) Manifests 

func (es *Extensions) Manifests() []*Manifest

Manifests returns a extension manifests

func (*Extensions) NewSessionExtensions 

func (es *Extensions) NewSessionExtensions() *SessionExtensions

NewSessionExtensions creates a new ConnExtensions object

type FunctionContext 

type FunctionContext interface {
	context.Context
	User() *auth.UserIdentity
	ActiveRoles() []*auth.RoleIdentity
	CurrentDB() string
	ConnectionInfo() *variable.ConnectionInfo
	EvalArgs(row chunk.Row) ([]types.Datum, error)
}

FunctionContext is an interface to provide context to the custom function

type FunctionDef 

type FunctionDef struct {
	// Name is the function's name
	Name string
	// EvalTp is the type of the return value
	EvalTp types.EvalType
	// ArgTps is the argument types
	ArgTps []types.EvalType
	// OptionalArgsLen is the length of the optional args
	OptionalArgsLen int
	// EvalStringFunc is the eval function when `EvalTp` is `types.ETString`
	EvalStringFunc func(ctx FunctionContext, row chunk.Row) (string, bool, error)
	// EvalIntFunc is the eval function when `EvalTp` is `types.ETInt`
	EvalIntFunc func(ctx FunctionContext, row chunk.Row) (int64, bool, error)
	// RequireDynamicPrivileges is a function to return a list of dynamic privileges to check.
	RequireDynamicPrivileges func(sem bool) []string
}

FunctionDef is the definition for the custom function

func (*FunctionDef) Validate 

func (def *FunctionDef) Validate() error

Validate validates the function definition

type Manifest 

type Manifest struct {
	// contains filtered or unexported fields
}

Manifest is an extension's manifest

func (*Manifest) Name 

func (m *Manifest) Name() string

Name returns the extension's name

type Option 

type Option func(m *Manifest)

Option represents an option to initialize an extension

func WithBootstrap 

func WithBootstrap(fn func(BootstrapContext) error) Option

WithBootstrap specifies the bootstrap func of an extension

func WithBootstrapSQL 

func WithBootstrapSQL(sqlList ...string) Option

WithBootstrapSQL the bootstrap SQL list

func WithClose 

func WithClose(fn func()) Option

WithClose specifies the close function of an extension. It will be invoked when `extension.Reset` is called

func WithCustomAccessCheck 

func WithCustomAccessCheck(fn AccessCheckFunc) Option

WithCustomAccessCheck specifies the custom db/tbl/column dynamic privilege check

func WithCustomAuthPlugins 

func WithCustomAuthPlugins(authPlugins []*AuthPlugin) Option

WithCustomAuthPlugins specifies the custom authentication plugins available for the system.

func WithCustomDynPrivs 

func WithCustomDynPrivs(privs []string) Option

WithCustomDynPrivs specifies dynamic privileges of an extension

func WithCustomFunctions 

func WithCustomFunctions(funcs []*FunctionDef) Option

WithCustomFunctions specifies custom functions

func WithCustomSysVariables 

func WithCustomSysVariables(vars []*variable.SysVar) Option

WithCustomSysVariables specifies custom variables of an extension

func WithSessionHandlerFactory 

func WithSessionHandlerFactory(factory func() *SessionHandler) Option

WithSessionHandlerFactory specifies a factory function to handle session

type SessionExtensions 

type SessionExtensions struct {
	// contains filtered or unexported fields
}

SessionExtensions is the extensions

func (*SessionExtensions) GetAuthPlugin 

func (es *SessionExtensions) GetAuthPlugin(name string) (*AuthPlugin, bool)

GetAuthPlugin returns the required registered extension auth plugin and whether it exists.

func (*SessionExtensions) HasStmtEventListeners 

func (es *SessionExtensions) HasStmtEventListeners() bool

HasStmtEventListeners returns a bool that indicates if any stmt event listener exists

func (*SessionExtensions) OnConnectionEvent 

func (es *SessionExtensions) OnConnectionEvent(tp ConnEventTp, event *ConnEventInfo)

OnConnectionEvent will be called when a connection event happens

func (*SessionExtensions) OnStmtEvent 

func (es *SessionExtensions) OnStmtEvent(tp StmtEventTp, event StmtEventInfo)

OnStmtEvent will be called when a stmt event happens

type SessionHandler 

type SessionHandler struct {
	OnConnectionEvent func(ConnEventTp, *ConnEventInfo)
	OnStmtEvent       func(StmtEventTp, StmtEventInfo)
}

SessionHandler is used to listen session events

type SessionPool 

type SessionPool interface {
	Get() (pools.Resource, error)
	Put(pools.Resource)
}

SessionPool is the pool for session

type StmtEventInfo 

type StmtEventInfo interface {
	// User returns the user of the session
	User() *auth.UserIdentity
	// ActiveRoles returns the active roles of the user
	ActiveRoles() []*auth.RoleIdentity
	// CurrentDB returns the current database
	CurrentDB() string
	// ConnectionInfo returns the connection info of the current session
	ConnectionInfo() *variable.ConnectionInfo
	// SessionAlias returns the session alias value set by user
	SessionAlias() string
	// StmtNode returns the parsed ast of the statement
	// When parse error, this method will return a nil value
	StmtNode() ast.StmtNode
	// ExecuteStmtNode will return the `ast.ExecuteStmt` node when the current statement is EXECUTE,
	// otherwise a nil value will be returned
	ExecuteStmtNode() *ast.ExecuteStmt
	// ExecutePreparedStmt will return the prepared stmt node for the EXECUTE statement.
	// If the current statement is not EXECUTE or prepared statement is not found, a nil value will be returned
	ExecutePreparedStmt() ast.StmtNode
	// PreparedParams will return the params for the EXECUTE statement
	PreparedParams() []types.Datum
	// OriginalText will return the text of the statement.
	// Notice that for the EXECUTE statement, the prepared statement text will be used as the return value
	OriginalText() string
	// SQLDigest will return the normalized and redact text of the `OriginalText()`
	SQLDigest() (normalized string, digest *parser.Digest)
	// AffectedRows will return the affected rows of the current statement
	AffectedRows() uint64
	// RelatedTables will return the related tables of the current statement
	// For statements succeeding to build logical plan, it uses the `visitinfo` to get the related tables
	// For statements failing to build logical plan, it traverses the ast node to get the related tables
	RelatedTables() []stmtctx.TableEntry
	// GetError will return the error when the current statement is failed
	GetError() error
}

StmtEventInfo is the information of stmt event

type StmtEventTp 

type StmtEventTp uint8

StmtEventTp is the type of the statement event 

const (
	// StmtError means the stmt is failed
	StmtError StmtEventTp = iota
	// StmtSuccess means the stmt is successfully executed
	StmtSuccess
)

type VerifyDynamicPrivRequest 

type VerifyDynamicPrivRequest struct {
	// User The username in the connect attempt
	User string
	// Host The host that the user is connecting from
	Host string
	// DynamicPriv the dynamic privilege required by the user's SQL statement
	DynamicPriv string
	// ConnState The TLS connection state (contains the TLS certificate) if client is using TLS. It will be nil if the client is not using TLS
	ConnState *tls.ConnectionState
	// ActiveRoles List of active MySQL roles for the current user
	ActiveRoles []*auth.RoleIdentity
	// WithGrant Whether the statement to be executed is granting the user privilege for executing GRANT statements
	WithGrant bool
}

VerifyDynamicPrivRequest contains the context for the plugin to authorize a user's dynamic privilege.

type VerifyStaticPrivRequest 

type VerifyStaticPrivRequest struct {
	// User The username in the connect attempt
	User string
	// Host The host that the user is connecting from
	Host string
	// DB The database to check for privilege
	DB string
	// Table The table to check for privilege
	Table string
	// Column The column to check for privilege (currently just a placeholder in TiDB as column-level privilege is not supported by TiDB yet)
	Column string
	// StaticPriv The privilege type of the SQL statement that will be executed
	StaticPriv mysql.PrivilegeType
	// ConnState The TLS connection state (contains the TLS certificate) if client is using TLS. It will be nil if the client is not using TLS
	ConnState *tls.ConnectionState
	// ActiveRoles List of active MySQL roles for the current user
	ActiveRoles []*auth.RoleIdentity
}

VerifyStaticPrivRequest contains the context for the plugin to authorize a user's static privilege.

Source Files

View all Source files

Directories

Path Synopsis

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.