package server

import (
	"crypto/rsa"
	"crypto/tls"
	"fmt"
	"sync"
	"sync/atomic"

	"github.com/go-mysql-org/go-mysql/mysql"
)

// userConfigurableServerCapabilities lists handshake flags that may be toggled
// via SetCapability / UnsetCapability. Other flags (e.g. CLIENT_SSL) are derived
// from server construction and must not be changed independently.
const userConfigurableServerCapabilities = mysql.CLIENT_LOCAL_FILES |
	mysql.CLIENT_MULTI_RESULTS |
	mysql.CLIENT_PS_MULTI_RESULTS

// Defines a basic MySQL server with configs.
//
// We do not aim at implementing the whole MySQL connection suite to have the best compatibilities for the clients.
// The MySQL server can be configured to switch auth methods covering 'mysql_old_password', 'mysql_native_password',
// 'mysql_clear_password', 'authentication_windows_client', 'sha256_password', 'caching_sha2_password', etc.
//
// However, since some old auth methods are considered broken with security issues. MySQL major versions like 5.7 and 8.0 default to
// 'mysql_native_password' or 'caching_sha2_password', and most MySQL clients should have already supported at least one of the three auth
// methods 'mysql_native_password', 'caching_sha2_password', and 'sha256_password'. Thus here we will only support these three
// auth methods, and use 'mysql_native_password' as default for maximum compatibility with the clients and leave the other two as
// config options.
//
// The MySQL doc states that 'mysql_old_password' will be used if 'CLIENT_PROTOCOL_41' or 'CLIENT_SECURE_CONNECTION' flag is not set.
// We choose to drop the support for insecure 'mysql_old_password' auth method and require client capability 'CLIENT_PROTOCOL_41' and 'CLIENT_SECURE_CONNECTION'
// are set. Besides, if 'CLIENT_PLUGIN_AUTH' is not set, we fallback to 'mysql_native_password' auth method.
type Server struct {
	serverVersion     string // e.g. "8.0.12"
	protocolVersion   int    // minimal 10
	capability        uint32 // server capability flag
	collationID       uint8
	defaultAuthMethod string // default authentication method, 'mysql_native_password'
	rsaPrivateKey     *rsa.PrivateKey
	rsaPublicKeyBytes []byte
	tlsConfig         *tls.Config
	cacheShaPassword  *sync.Map // 'user@host' -> SHA256(SHA256(PASSWORD))
	authProvider      AuthenticationProvider
}

// NewDefaultServer: New mysql server with default settings.
//
// NOTES:
// TLS support will be enabled by default with auto-generated CA and server certificates (however, you can still use
// non-TLS connection). By default, it will verify the client certificate if present. You can enable TLS support on
// the client side without providing a client-side certificate. So only when you need the server to verify client
// identity for maximum security, you need to set a signed certificate for the client.
func NewDefaultServer() *Server {
	caPem, caKey := generateCA()
	certPem, keyPem := generateAndSignRSACerts(caPem, caKey)
	tlsConf := NewServerTLSConfig(caPem, certPem, keyPem, tls.VerifyClientCertIfGiven)
	rsaPrivateKey, rsaPublicKeyBytes := getRSAKeyPairFromPEM(keyPem)
	return &Server{
		serverVersion:   "8.0.11",
		protocolVersion: 10,
		capability: mysql.CLIENT_LONG_PASSWORD | mysql.CLIENT_LONG_FLAG | mysql.CLIENT_CONNECT_WITH_DB | mysql.CLIENT_PROTOCOL_41 |
			mysql.CLIENT_TRANSACTIONS | mysql.CLIENT_SECURE_CONNECTION | mysql.CLIENT_PLUGIN_AUTH | mysql.CLIENT_SSL |
			mysql.CLIENT_PLUGIN_AUTH_LENENC_CLIENT_DATA | mysql.CLIENT_CONNECT_ATTRS | mysql.CLIENT_SESSION_TRACK |
			mysql.CLIENT_MULTI_RESULTS | mysql.CLIENT_PS_MULTI_RESULTS,
		collationID:       mysql.DEFAULT_COLLATION_ID,
		defaultAuthMethod: mysql.AUTH_NATIVE_PASSWORD,
		rsaPrivateKey:     rsaPrivateKey,
		rsaPublicKeyBytes: rsaPublicKeyBytes,
		tlsConfig:         tlsConf,
		cacheShaPassword:  new(sync.Map),
		authProvider:      &DefaultAuthenticationProvider{},
	}
}

// NewServer: New mysql server with customized settings.
//
// NOTES:
// You can control the authentication methods and TLS settings here.
//
// For auth method, you can specify one of the supported methods 'mysql_native_password', 'caching_sha2_password', and 'sha256_password'.
// The specified auth method will be enforced by the server in the connection phase. That means, client will be asked to switch auth method
// if the supplied auth method is different from the server default.
//
// For TLS support, you can specify self-signed or CA-signed certificates and decide whether the client needs to provide
// a signed or unsigned certificate to provide different level of security.
//
// The rsaKey parameter is used for password encryption on non-TLS connections with 'caching_sha2_password' and 'sha256_password'.
// If it's is nil, it will attempt to extract an RSA key from tlsConfig.Certificates[0].
// If no RSA key is available, non-TLS connections will not be supported for these auth methods (TLS connections will still work).
func NewServer(serverVersion string, collationID uint8, defaultAuthMethod string, rsaKey *rsa.PrivateKey, tlsConfig *tls.Config) *Server {
	authProvider := &DefaultAuthenticationProvider{}
	return NewServerWithAuth(serverVersion, collationID, defaultAuthMethod, rsaKey, tlsConfig, authProvider)
}

func NewServerWithAuth(serverVersion string, collationID uint8, defaultAuthMethod string, rsaKey *rsa.PrivateKey, tlsConfig *tls.Config, authProvider AuthenticationProvider) *Server {
	if authProvider == nil || !authProvider.Validate(defaultAuthMethod) {
		panic(fmt.Sprintf("server authentication method '%s' is not supported", defaultAuthMethod))
	}

	if rsaKey == nil && (defaultAuthMethod == mysql.AUTH_CACHING_SHA2_PASSWORD || defaultAuthMethod == mysql.AUTH_SHA256_PASSWORD) {
		if tlsConfig != nil && len(tlsConfig.Certificates) > 0 {
			rsaKey, _ = tlsConfig.Certificates[0].PrivateKey.(*rsa.PrivateKey)
		} else {
			panic("either rsaKey or tlsConfig is required for caching_sha2_password and sha256_password authentication")
		}
	}

	capFlag := mysql.CLIENT_LONG_PASSWORD | mysql.CLIENT_LONG_FLAG | mysql.CLIENT_CONNECT_WITH_DB | mysql.CLIENT_PROTOCOL_41 |
		mysql.CLIENT_TRANSACTIONS | mysql.CLIENT_SECURE_CONNECTION | mysql.CLIENT_PLUGIN_AUTH | mysql.CLIENT_CONNECT_ATTRS |
		mysql.CLIENT_PLUGIN_AUTH_LENENC_CLIENT_DATA | mysql.CLIENT_SESSION_TRACK |
		mysql.CLIENT_MULTI_RESULTS | mysql.CLIENT_PS_MULTI_RESULTS
	if tlsConfig != nil {
		capFlag |= mysql.CLIENT_SSL
	}
	return &Server{
		serverVersion:     serverVersion,
		protocolVersion:   10,
		capability:        capFlag,
		collationID:       collationID,
		defaultAuthMethod: defaultAuthMethod,
		rsaPrivateKey:     rsaKey,
		rsaPublicKeyBytes: rsaPublicKeyBytes(rsaKey),
		tlsConfig:         tlsConfig,
		cacheShaPassword:  new(sync.Map),
		authProvider:      authProvider,
	}
}

func isAuthMethodSupported(authMethod string) bool {
	return authMethod == mysql.AUTH_NATIVE_PASSWORD || authMethod == mysql.AUTH_CACHING_SHA2_PASSWORD || authMethod == mysql.AUTH_SHA256_PASSWORD || authMethod == mysql.AUTH_CLEAR_PASSWORD
}

func (s *Server) InvalidateCache(username string, host string) {
	s.cacheShaPassword.Delete(fmt.Sprintf("%s@%s", username, host))
}

// Capability returns the capability flags advertised in the initial handshake.
func (s *Server) Capability() uint32 {
	return atomic.LoadUint32(&s.capability)
}

// SetCapability enables additional server capabilities advertised in the handshake.
// Only CLIENT_LOCAL_FILES, CLIENT_MULTI_RESULTS, and CLIENT_PS_MULTI_RESULTS may
// be set; other flags are managed by server construction (e.g. CLIENT_SSL requires TLS).
func (s *Server) SetCapability(capability uint32) error {
	if err := validateUserConfigurableCapability(capability); err != nil {
		return err
	}
	for {
		old := atomic.LoadUint32(&s.capability)
		if atomic.CompareAndSwapUint32(&s.capability, old, old|capability) {
			break
		}
	}
	return nil
}

// UnsetCapability disables server capabilities advertised in the handshake.
// Only CLIENT_LOCAL_FILES, CLIENT_MULTI_RESULTS, and CLIENT_PS_MULTI_RESULTS may
// be cleared via this API.
func (s *Server) UnsetCapability(capability uint32) error {
	if err := validateUserConfigurableCapability(capability); err != nil {
		return err
	}
	for {
		old := atomic.LoadUint32(&s.capability)
		if atomic.CompareAndSwapUint32(&s.capability, old, old&^capability) {
			break
		}
	}
	return nil
}

func validateUserConfigurableCapability(capability uint32) error {
	if capability == 0 {
		return fmt.Errorf("capability must not be zero")
	}
	if invalid := capability &^ userConfigurableServerCapabilities; invalid != 0 {
		return fmt.Errorf("server capability %#x contains non-user-configurable flags %#x; only CLIENT_LOCAL_FILES, CLIENT_MULTI_RESULTS, and CLIENT_PS_MULTI_RESULTS are supported", capability, invalid)
	}
	return nil
}
