connect.go

package config

import (
	"fmt"
	"net"
	"os"
	"strconv"
	"time"

	"github.com/go-sql-driver/mysql"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"
	gormMysql "gorm.io/driver/mysql"
	"gorm.io/gorm"
)

type ViaSSHDialer struct {
	client *ssh.Client
}

func (s *ViaSSHDialer) Dial(addr string) (net.Conn, error) {
	return s.client.Dial("tcp", addr)
}

func NewMysql(configuration Config) *gorm.DB {

	var agentClient agent.Agent

	if conn, err := net.Dial("unix", os.Getenv("SSH_AUTH_SOCK")); err == nil {
		defer conn.Close()

		agentClient = agent.NewClient(conn)
	} else {
		panic(err)
	}

	signers, err := agentClient.Signers()
	if err != nil {
		panic(err)
	}

	sshConfig := &ssh.ClientConfig{
		User: configuration.Get("SSH_USER"),
		Auth: []ssh.AuthMethod{
			ssh.PublicKeys(signers...),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Timeout:         time.Minute * 5,
	}

	if agentClient != nil {
		sshConfig.Auth = append(sshConfig.Auth, ssh.PublicKeysCallback(agentClient.Signers))
	}

	sshConfig.Auth = append(sshConfig.Auth, ssh.PasswordCallback(func() (string, error) {
		return configuration.Get("SSH_PASS"), nil
	}))

	portSsh, err := strconv.Atoi(configuration.Get("SSH_PORT"))
	if err != nil {
		panic(err)
	}

	// Connect to the SSH Server
	if sshcon, err := ssh.Dial("tcp", fmt.Sprintf("%s:%d", configuration.Get("SSH_HOST"), portSsh), sshConfig); err == nil {
		defer sshcon.Close()

		mysql.RegisterDial("mysql+tcp", (&ViaSSHDialer{sshcon}).Dial)
	} else {
		panic(err)
	}

	port, err := strconv.Atoi(configuration.Get("DB_PORT"))
	if err != nil {
		panic(err)
	}

	dsn := fmt.Sprintf(
		"%s:%s@mysql+tcp(%s:%d)/%s?parseTime=true", configuration.Get("DB_USER"), configuration.Get("DB_PASSWORD"), configuration.Get("DB_HOST"), port, configuration.Get("DB_NAME"),
	)

	db, err := gorm.Open(gormMysql.Open(dsn), &gorm.Config{})

	if err != nil {
		panic(err)
	}

	return db
}