优化接入网关，解决协程泄漏的问题
This commit is contained in:
dushixiang
@ -1,13 +1,13 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net"
|
||||
"os"
|
||||
"sync"
|
||||
|
||||
"next-terminal/server/term"
|
||||
"next-terminal/server/utils"
|
||||
|
||||
"golang.org/x/crypto/ssh"
|
||||
@ -15,33 +15,35 @@ import (
|
||||
|
||||
// Gateway 接入网关
|
||||
type Gateway struct {
|
||||
ID string // 接入网关ID
|
||||
Connected bool // 是否已连接
|
||||
SshClient *ssh.Client
|
||||
Message string // 失败原因
|
||||
ID string // 接入网关ID
|
||||
IP string
|
||||
Port int
|
||||
Username string
|
||||
Password string
|
||||
PrivateKey string
|
||||
Passphrase string
|
||||
Connected bool // 是否已连接
|
||||
Message string // 失败原因
|
||||
SshClient *ssh.Client
|
||||
|
||||
tunnels sync.Map
|
||||
}
|
||||
|
||||
func NewGateway(id string, connected bool, message string, client *ssh.Client) *Gateway {
|
||||
return &Gateway{
|
||||
ID: id,
|
||||
Connected: connected,
|
||||
Message: message,
|
||||
SshClient: client,
|
||||
}
|
||||
}
|
||||
|
||||
func (g *Gateway) Close() {
|
||||
g.tunnels.Range(func(key, value interface{}) bool {
|
||||
g.CloseSshTunnel(key.(string))
|
||||
return true
|
||||
})
|
||||
mutex sync.Mutex
|
||||
tunnels map[string]*Tunnel
|
||||
}
|
||||
|
||||
func (g *Gateway) OpenSshTunnel(id, ip string, port int) (exposedIP string, exposedPort int, err error) {
|
||||
g.mutex.Lock()
|
||||
defer g.mutex.Unlock()
|
||||
if !g.Connected {
|
||||
return "", 0, errors.New(g.Message)
|
||||
sshClient, err := term.NewSshClient(g.IP, g.Port, g.Username, g.Password, g.PrivateKey, g.Passphrase)
|
||||
if err != nil {
|
||||
g.Connected = false
|
||||
g.Message = "接入网关不可用：" + err.Error()
|
||||
return "", 0, errors.New(g.Message)
|
||||
} else {
|
||||
g.Connected = true
|
||||
g.SshClient = sshClient
|
||||
g.Message = "使用中"
|
||||
}
|
||||
}
|
||||
|
||||
localPort, err := utils.GetAvailablePort()
|
||||
@ -63,30 +65,39 @@ func (g *Gateway) OpenSshTunnel(id, ip string, port int) (exposedIP string, expo
|
||||
return "", 0, err
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
tunnel := &Tunnel{
|
||||
ID: id,
|
||||
LocalHost: hostname,
|
||||
//LocalHost: "docker.for.mac.host.internal",
|
||||
LocalPort: localPort,
|
||||
Gateway: g,
|
||||
RemoteHost: ip,
|
||||
RemotePort: port,
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
id: id,
|
||||
localHost: hostname,
|
||||
//localHost: "docker.for.mac.host.internal",
|
||||
localPort: localPort,
|
||||
remoteHost: ip,
|
||||
remotePort: port,
|
||||
listener: listener,
|
||||
}
|
||||
go tunnel.Open()
|
||||
g.tunnels.Store(tunnel.ID, tunnel)
|
||||
go tunnel.Open(g.SshClient)
|
||||
g.tunnels[tunnel.id] = tunnel
|
||||
|
||||
return tunnel.LocalHost, tunnel.LocalPort, nil
|
||||
return tunnel.localHost, tunnel.localPort, nil
|
||||
}
|
||||
|
||||
func (g *Gateway) CloseSshTunnel(id string) {
|
||||
if value, ok := g.tunnels.Load(id); ok {
|
||||
if tunnel, vok := value.(*Tunnel); vok {
|
||||
tunnel.Close()
|
||||
g.tunnels.Delete(id)
|
||||
}
|
||||
g.mutex.Lock()
|
||||
defer g.mutex.Unlock()
|
||||
t := g.tunnels[id]
|
||||
if t != nil {
|
||||
t.Close()
|
||||
delete(g.tunnels, id)
|
||||
}
|
||||
|
||||
if len(g.tunnels) == 0 {
|
||||
_ = g.SshClient.Close()
|
||||
g.Connected = false
|
||||
g.Message = "暂未使用"
|
||||
}
|
||||
}
|
||||
|
||||
func (g *Gateway) Close() {
|
||||
for id := range g.tunnels {
|
||||
g.CloseSshTunnel(id)
|
||||
}
|
||||
}
|
||||
|
||||
@ -4,35 +4,50 @@ import (
|
||||
"sync"
|
||||
|
||||
"next-terminal/server/log"
|
||||
"next-terminal/server/model"
|
||||
)
|
||||
|
||||
type Manager struct {
|
||||
type manager struct {
|
||||
gateways sync.Map
|
||||
}
|
||||
|
||||
func NewManager() *Manager {
|
||||
return &Manager{}
|
||||
}
|
||||
|
||||
func (m *Manager) GetById(id string) *Gateway {
|
||||
func (m *manager) GetById(id string) *Gateway {
|
||||
if val, ok := m.gateways.Load(id); ok {
|
||||
return val.(*Gateway)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *Manager) Add(g *Gateway) {
|
||||
func (m *manager) Add(model *model.AccessGateway) *Gateway {
|
||||
g := &Gateway{
|
||||
ID: model.ID,
|
||||
IP: model.IP,
|
||||
Port: model.Port,
|
||||
Username: model.Username,
|
||||
Password: model.Password,
|
||||
PrivateKey: model.PrivateKey,
|
||||
Passphrase: model.Passphrase,
|
||||
Connected: false,
|
||||
SshClient: nil,
|
||||
Message: "暂未使用",
|
||||
tunnels: make(map[string]*Tunnel),
|
||||
}
|
||||
m.gateways.Store(g.ID, g)
|
||||
log.Infof("add gateway: %s", g.ID)
|
||||
log.Infof("add Gateway: %s", g.ID)
|
||||
return g
|
||||
}
|
||||
|
||||
func (m *Manager) Del(id string) {
|
||||
func (m *manager) Del(id string) {
|
||||
g := m.GetById(id)
|
||||
if g != nil {
|
||||
g.Close()
|
||||
}
|
||||
m.gateways.Delete(id)
|
||||
log.Infof("del gateway: %s", id)
|
||||
log.Infof("del Gateway: %s", id)
|
||||
}
|
||||
|
||||
var GlobalGatewayManager *Manager
|
||||
var GlobalGatewayManager *manager
|
||||
|
||||
func init() {
|
||||
GlobalGatewayManager = NewManager()
|
||||
GlobalGatewayManager = &manager{}
|
||||
}
|
||||
|
||||
@ -1,59 +1,55 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
|
||||
"next-terminal/server/log"
|
||||
|
||||
"golang.org/x/crypto/ssh"
|
||||
)
|
||||
|
||||
type Tunnel struct {
|
||||
ID string // 唯一标识
|
||||
LocalHost string // 本地监听地址
|
||||
LocalPort int // 本地端口
|
||||
RemoteHost string // 远程连接地址
|
||||
RemotePort int // 远程端口
|
||||
Gateway *Gateway
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
id string // 唯一标识
|
||||
localHost string // 本地监听地址
|
||||
localPort int // 本地端口
|
||||
remoteHost string // 远程连接地址
|
||||
remotePort int // 远程端口
|
||||
listener net.Listener
|
||||
localConnections []net.Conn
|
||||
remoteConnections []net.Conn
|
||||
}
|
||||
|
||||
func (r *Tunnel) Open() {
|
||||
localAddr := fmt.Sprintf("%s:%d", r.LocalHost, r.LocalPort)
|
||||
func (r *Tunnel) Open(sshClient *ssh.Client) {
|
||||
localAddr := fmt.Sprintf("%s:%d", r.localHost, r.localPort)
|
||||
|
||||
go func() {
|
||||
<-r.ctx.Done()
|
||||
_ = r.listener.Close()
|
||||
log.Debugf("SSH 隧道 %v 关闭", localAddr)
|
||||
}()
|
||||
for {
|
||||
log.Debugf("等待客户端访问 %v", localAddr)
|
||||
log.Debugf("隧道 %v 等待客户端访问 %v", r.id, localAddr)
|
||||
localConn, err := r.listener.Accept()
|
||||
if err != nil {
|
||||
log.Debugf("接受连接失败 %v, 退出循环", err.Error())
|
||||
log.Debugf("隧道 %v 接受连接失败 %v, 退出循环", r.id, err.Error())
|
||||
log.Debug("-------------------------------------------------")
|
||||
return
|
||||
}
|
||||
r.localConnections = append(r.localConnections, localConn)
|
||||
|
||||
log.Debugf("客户端 %v 连接至 %v", localConn.RemoteAddr().String(), localAddr)
|
||||
remoteAddr := fmt.Sprintf("%s:%d", r.RemoteHost, r.RemotePort)
|
||||
log.Debugf("连接远程主机 %v ...", remoteAddr)
|
||||
remoteConn, err := r.Gateway.SshClient.Dial("tcp", remoteAddr)
|
||||
log.Debugf("隧道 %v 新增本地连接 %v", r.id, localConn.RemoteAddr().String())
|
||||
remoteAddr := fmt.Sprintf("%s:%d", r.remoteHost, r.remotePort)
|
||||
log.Debugf("隧道 %v 连接远程地址 %v ...", r.id, remoteAddr)
|
||||
remoteConn, err := sshClient.Dial("tcp", remoteAddr)
|
||||
if err != nil {
|
||||
log.Debugf("连接远程主机 %v 失败", remoteAddr)
|
||||
log.Debugf("隧道 %v 连接远程地址 %v, 退出循环", r.id, err.Error())
|
||||
log.Debug("-------------------------------------------------")
|
||||
return
|
||||
}
|
||||
r.remoteConnections = append(r.remoteConnections, remoteConn)
|
||||
|
||||
log.Debugf("连接远程主机 %v 成功", remoteAddr)
|
||||
log.Debugf("隧道 %v 连接远程主机成功", r.id)
|
||||
go copyConn(localConn, remoteConn)
|
||||
go copyConn(remoteConn, localConn)
|
||||
log.Debugf("转发数据 [%v]->[%v]", localAddr, remoteAddr)
|
||||
log.Debugf("隧道 %v 开始转发数据 [%v]->[%v]", r.id, localAddr, remoteAddr)
|
||||
log.Debug("~~~~~~~~~~~~~~~~~~~~分割线~~~~~~~~~~~~~~~~~~~~~~~~")
|
||||
}
|
||||
}
|
||||
|
||||
@ -66,7 +62,8 @@ func (r *Tunnel) Close() {
|
||||
_ = r.remoteConnections[i].Close()
|
||||
}
|
||||
r.remoteConnections = nil
|
||||
r.cancel()
|
||||
_ = r.listener.Close()
|
||||
log.Debugf("隧道 %v 监听器关闭", r.id)
|
||||
}
|
||||
|
||||
func copyConn(writer, reader net.Conn) {
|
||||
|
||||
Reference in New Issue
Block a user