From 8c4c41cf0b81aa20e6a0439e4de444946f8f302f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E5=AE=97=E6=B4=8B?= <liuzongyang@cqliving.com>
Date: Wed, 20 Dec 2023 17:41:24 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=E8=BE=BE=E6=A2=A6=E6=94=AF=E6=8C=81ssh?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../src/views/ops/db/InstanceEdit.vue         |  2 +-
 mayfly_go_web/src/views/ops/db/SqlExec.vue    |  8 +-
 server/internal/db/dbm/dialect_dm.go          | 82 +++++++++++++++++++
 3 files changed, 87 insertions(+), 5 deletions(-)

diff --git a/mayfly_go_web/src/views/ops/db/InstanceEdit.vue b/mayfly_go_web/src/views/ops/db/InstanceEdit.vue
index 8764622d..fc55bbba 100644
--- a/mayfly_go_web/src/views/ops/db/InstanceEdit.vue
+++ b/mayfly_go_web/src/views/ops/db/InstanceEdit.vue
@@ -149,7 +149,7 @@ const dbTypes = [
     },
     {
         type: 'dm',
-        label: '达梦(暂不支持ssh)',
+        label: '达梦',
     },
 ];
 
diff --git a/mayfly_go_web/src/views/ops/db/SqlExec.vue b/mayfly_go_web/src/views/ops/db/SqlExec.vue
index 04bdf85d..9df780aa 100644
--- a/mayfly_go_web/src/views/ops/db/SqlExec.vue
+++ b/mayfly_go_web/src/views/ops/db/SqlExec.vue
@@ -161,11 +161,11 @@
 </template>
 
 <script lang="ts" setup>
-import { defineAsyncComponent, onMounted, reactive, ref, toRefs, onBeforeUnmount } from 'vue';
+import { defineAsyncComponent, onBeforeUnmount, onMounted, reactive, ref, toRefs } from 'vue';
 import { ElMessage, ElMessageBox } from 'element-plus';
 import { formatByteSize } from '@/common/utils/format';
-import { DbInst, TabInfo, TabType, registerDbCompletionItemProvider } from './db';
-import { TagTreeNode, NodeType } from '../component/tag';
+import { DbInst, registerDbCompletionItemProvider, TabInfo, TabType } from './db';
+import { NodeType, TagTreeNode } from '../component/tag';
 import TagTree from '../component/TagTree.vue';
 import { dbApi } from './api';
 import { dispposeCompletionItemProvider } from '@/components/monaco/completionItemProvider';
@@ -174,7 +174,7 @@ import { ContextmenuItem } from '@/components/contextmenu';
 import { getDbDialect } from './dialect/index';
 import { sleep } from '@/common/utils/loading';
 import { TagResourceTypeEnum } from '@/common/commonEnum';
-import { Splitpanes, Pane } from 'splitpanes';
+import { Pane, Splitpanes } from 'splitpanes';
 import { useEventListener } from '@vueuse/core';
 
 const DbSqlEditor = defineAsyncComponent(() => import('./component/sqleditor/DbSqlEditor.vue'));
diff --git a/server/internal/db/dbm/dialect_dm.go b/server/internal/db/dbm/dialect_dm.go
index 81745346..dfe3f9a7 100644
--- a/server/internal/db/dbm/dialect_dm.go
+++ b/server/internal/db/dbm/dialect_dm.go
@@ -4,13 +4,86 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
+	"io"
+	machineapp "mayfly-go/internal/machine/application"
 	"mayfly-go/pkg/errorx"
+	"mayfly-go/pkg/logx"
 	"mayfly-go/pkg/utils/anyx"
+	"net"
 	"strings"
 
 	_ "gitee.com/chunanyong/dm"
 )
 
+type ConnectionInfo struct {
+	Port       int
+	Listener   net.Listener
+	remoteConn net.Conn
+}
+
+var connectionMap = make(map[string]ConnectionInfo)
+
+func getLocalListener() (net.Listener, int, error) {
+	// Setup localListener (type net.Listener)
+	localListener, err := net.Listen("tcp", "localhost:0")
+	if err != nil {
+		return nil, 0, err
+	}
+
+	// 获取本地端口
+	localPort := localListener.Addr().(*net.TCPAddr).Port
+	return localListener, localPort, nil
+}
+
+func acceptConn(listener net.Listener, sshConn net.Conn) {
+	for {
+		localConn, err := listener.Accept()
+		if err != nil {
+			logx.Warn("端口转发出错", err)
+			return
+		}
+		go forward(localConn, sshConn)
+	}
+}
+
+func forward(localConn net.Conn, remoteConn net.Conn) {
+	copyConn := func(writer, reader net.Conn) {
+		_, err := io.Copy(writer, reader)
+		if err != nil {
+			logx.Warnf("io.Copy error: %s", err)
+		}
+	}
+	go copyConn(localConn, remoteConn)
+	go copyConn(remoteConn, localConn)
+}
+
+func openSsh(d *DbInfo) error {
+
+	sshTunnelMachine, err := machineapp.GetMachineApp().GetSshTunnelMachine(d.SshTunnelMachineId)
+	if err != nil {
+		return err
+	}
+	remoteConn, err := sshTunnelMachine.GetDialConn("tcp", fmt.Sprintf("%s:%d", d.Host, d.Port))
+	if err != nil {
+		return err
+	}
+	// 获取sshConn的本地端口
+	localLister, localPort, err := getLocalListener()
+	// defer localLister.Close()
+	go acceptConn(localLister, remoteConn)
+	connectionMap[d.Network] = ConnectionInfo{
+		Port:       localPort,
+		Listener:   localLister,
+		remoteConn: remoteConn,
+	}
+	d.Host = "127.0.0.1"
+	d.Port = localPort
+
+	return nil
+}
+
+// 创建一个成员变量存放ssh隧道转发对应的本地连接
+
 func getDmDB(d *DbInfo) (*sql.DB, error) {
 	driverName := "dm"
 	// SSH Conect 暂时不支持隧道连接
@@ -25,6 +98,15 @@ func getDmDB(d *DbInfo) (*sql.DB, error) {
 			dbParam = db
 		}
 	}
+
+	// 开启ssh隧道
+	if d.SshTunnelMachineId > 0 {
+		err := openSsh(d)
+		if err != nil {
+			return nil, err
+		}
+	}
+
 	dsn := fmt.Sprintf("dm://%s:%s@%s:%d/%s", d.Username, d.Password, d.Host, d.Port, dbParam)
 	return sql.Open(driverName, dsn)
 }