package main

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"database/sql"
	"fmt"

	"github.com/aws/aws-lambda-go/lambda"
	"github.com/go-sql-driver/mysql"
)

type Event struct {
	DataSourceName string     `json:"dsn"`
	TLS            *TLSConfig `json:"tls,omitempty"`
}

type TLSConfig struct {
	CA   []byte `json:"ca"`
	Cert []byte `json:"cert"`
	Key  []byte `json:"key"`
}

const (
	tlsConfigName = "custom"
)

func Handler(ctx context.Context, e Event) error {
	cfg, err := mysql.ParseDSN(e.DataSourceName)
	if err != nil {
		return err
	}
	if e.TLS != nil {
		rootCertPool := x509.NewCertPool()
		if ok := rootCertPool.AppendCertsFromPEM(e.TLS.CA); !ok {
			return fmt.Errorf(("can't append ca"))
		}
		clientCert := make([]tls.Certificate, 0, 1)
		certs, err := tls.X509KeyPair(e.TLS.Cert, e.TLS.Key)
		if err != nil {
			return fmt.Errorf("invalid cert and key: %w", err)
		}
		clientCert = append(clientCert, certs)
		if err := mysql.RegisterTLSConfig(tlsConfigName, &tls.Config{
			RootCAs:      rootCertPool,
			Certificates: clientCert,
		}); err != nil {
			return fmt.Errorf("can't register tls config: %w", err)
		}
		defer mysql.DeregisterTLSConfig(tlsConfigName)

		cfg.TLSConfig = tlsConfigName
	}
	connector, err := mysql.NewConnector(cfg)
	if err != nil {
		return err
	}
	db := sql.OpenDB(connector)
	return db.PingContext(ctx)
}

func main() {
	lambda.Start(Handler)
}