evan/next-terminal
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

优化代码

Browse Source
This commit is contained in:
dushixiang
2022-05-06 21:10:29 +08:00
parent c58776fa24
commit 7357cebc34
16 changed files with 159 additions and 208 deletions
Download Patch File Download Diff File Expand all files Collapse all files

50
server/global/gateway/gateway.go
View File

@ -8,8 +8,6 @@ import (
"os"
"sync"
"next-terminal/server/log"
"next-terminal/server/utils"
"golang.org/x/crypto/ssh"
@ -22,11 +20,7 @@ type Gateway struct {
SshClient *ssh.Client
Message string // 失败原因
tunnels *sync.Map
Add chan *Tunnel
Del chan string
exit chan bool
tunnels sync.Map
}
func NewGateway(id string, connected bool, message string, client *ssh.Client) *Gateway {
@ -35,42 +29,14 @@ func NewGateway(id string, connected bool, message string, client *ssh.Client) *
Connected: connected,
Message: message,
SshClient: client,
Add: make(chan *Tunnel),
Del: make(chan string),
tunnels: new(sync.Map),
exit: make(chan bool, 1),
}
}
func (g *Gateway) Run() {
for {
select {
case t := <-g.Add:
g.tunnels.Store(t.ID, t)
log.Info("add tunnel: %s", t.ID)
go t.Open()
case k := <-g.Del:
if val, ok := g.tunnels.Load(k); ok {
if vval, vok := val.(*Tunnel); vok {
vval.Close()
g.tunnels.Delete(k)
}
}
case <-g.exit:
return
}
}
}
func (g *Gateway) Close() {
g.tunnels.Range(func(key, value interface{}) bool {
if val, ok := value.(*Tunnel); ok {
val.Close()
}
g.CloseSshTunnel(key.(string))
return true
})
g.exit <- true
}
func (g *Gateway) OpenSshTunnel(id, ip string, port int) (exposedIP string, exposedPort int, err error) {
@ -110,11 +76,17 @@ func (g *Gateway) OpenSshTunnel(id, ip string, port int) (exposedIP string, expo
cancel: cancel,
listener: listener,
}
g.Add <- tunnel
go tunnel.Open()
g.tunnels.Store(tunnel.ID, tunnel)
return tunnel.LocalHost, tunnel.LocalPort, nil
}
func (g Gateway) CloseSshTunnel(id string) {
g.Del <- id
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)
}
}
}

43
server/global/gateway/manager.go
View File

@ -7,49 +7,32 @@ import (
)
type Manager struct {
gateways *sync.Map
Add chan *Gateway
Del chan string
gateways sync.Map
}
func NewManager() *Manager {
return &Manager{
Add: make(chan *Gateway),
Del: make(chan string),
gateways: new(sync.Map),
}
return &Manager{}
}
func (m *Manager) Start() {
for {
select {
case g := <-m.Add:
m.gateways.Store(g.ID, g)
log.Info("add gateway: %s", g.ID)
go g.Run()
case k := <-m.Del:
if val, ok := m.gateways.Load(k); ok {
if vv, vok := val.(*Gateway); vok {
vv.Close()
m.gateways.Delete(k)
}
}
}
}
}
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) {
m.gateways.Store(g.ID, g)
log.Infof("add gateway: %s", g.ID)
}
func (m *Manager) Del(id string) {
m.gateways.Delete(id)
log.Infof("del gateway: %s", id)
}
var GlobalGatewayManager *Manager
func init() {
GlobalGatewayManager = NewManager()
go GlobalGatewayManager.Start()
}

50
server/global/gateway/tunnel.go
View File

@ -10,23 +10,26 @@ import (
)
type Tunnel struct {
ID string // 唯一标识
LocalHost string // 本地监听地址
LocalPort int // 本地端口
RemoteHost string // 远程连接地址
RemotePort int // 远程端口
Gateway *Gateway
ctx context.Context
cancel context.CancelFunc
listener net.Listener
err error
ID string // 唯一标识
LocalHost string // 本地监听地址
LocalPort int // 本地端口
RemoteHost string // 远程连接地址
RemotePort int // 远程端口
Gateway *Gateway
ctx context.Context
cancel context.CancelFunc
listener net.Listener
localConnections []net.Conn
remoteConnections []net.Conn
}
func (r *Tunnel) Open() {
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)
@ -35,6 +38,7 @@ func (r *Tunnel) Open() {
log.Debugf("接受连接失败 %v, 退出循环", err.Error())
return
}
r.localConnections = append(r.localConnections, localConn)
log.Debugf("客户端 %v 连接至 %v", localConn.RemoteAddr().String(), localAddr)
remoteAddr := fmt.Sprintf("%s:%d", r.RemoteHost, r.RemotePort)
@ -44,27 +48,27 @@ func (r *Tunnel) Open() {
log.Debugf("连接远程主机 %v 失败", remoteAddr)
return
}
r.remoteConnections = append(r.remoteConnections, remoteConn)
log.Debugf("连接远程主机 %v 成功", remoteAddr)
go copyConn(r.ctx, localConn, remoteConn)
go copyConn(r.ctx, remoteConn, localConn)
go copyConn(localConn, remoteConn)
go copyConn(remoteConn, localConn)
log.Debugf("转发数据 [%v]->[%v]", localAddr, remoteAddr)
}
}
func (r Tunnel) Close() {
func (r *Tunnel) Close() {
for i := range r.localConnections {
_ = r.localConnections[i].Close()
}
r.localConnections = nil
for i := range r.remoteConnections {
_ = r.remoteConnections[i].Close()
}
r.remoteConnections = nil
r.cancel()
}
func copyConn(ctx context.Context, writer, reader net.Conn) {
func copyConn(writer, reader net.Conn) {
_, _ = io.Copy(writer, reader)
for {
select {
case <-ctx.Done():
_ = writer.Close()
_ = reader.Close()
return
}
}
}

62
server/global/security/security.go
View File

@ -1,6 +1,11 @@
package security
import "sort"
import (
"sort"
"sync"
"next-terminal/server/log"
)
type Security struct {
ID string
@ -10,45 +15,29 @@ type Security struct {
}
type Manager struct {
securities map[string]*Security
securities sync.Map
values []*Security
Add chan *Security
Del chan string
}
func NewManager() *Manager {
return &Manager{
Add: make(chan *Security),
Del: make(chan string),
securities: map[string]*Security{},
}
}
func (m *Manager) Start() {
for {
select {
case s := <-m.Add:
m.securities[s.ID] = s
m.LoadData()
case s := <-m.Del:
if _, ok := m.securities[s]; ok {
delete(m.securities, s)
m.LoadData()
}
}
}
return &Manager{}
}
func (m *Manager) Clear() {
m.securities = map[string]*Security{}
m.securities.Range(func(k, _ interface{}) bool {
m.securities.Delete(k)
return true
})
}
func (m *Manager) LoadData() {
var values []*Security
for _, security := range m.securities {
values = append(values, security)
}
m.securities.Range(func(key, value interface{}) bool {
if security, ok := value.(*Security); ok {
values = append(values, security)
}
return true
})
sort.Slice(values, func(i, j int) bool {
// 优先级数字越小代表优先级越高，因此此处用小于号
@ -58,13 +47,24 @@ func (m *Manager) LoadData() {
m.values = values
}
func (m Manager) Values() []*Security {
func (m *Manager) Values() []*Security {
return m.values
}
func (m *Manager) Add(s *Security) {
m.securities.Store(s.ID, s)
m.LoadData()
log.Infof("add security: %s", s.ID)
}
func (m *Manager) Del(id string) {
m.securities.Delete(id)
m.LoadData()
log.Infof("del security: %s", id)
}
var GlobalSecurityManager *Manager
func init() {
GlobalSecurityManager = NewManager()
go GlobalSecurityManager.Start()
}

103
server/global/session/session.go
View File

@ -1,11 +1,11 @@
package session
import (
"fmt"
"sync"
"next-terminal/server/dto"
"next-terminal/server/guacd"
"next-terminal/server/log"
"next-terminal/server/term"
"github.com/gorilla/websocket"
@ -42,80 +42,79 @@ func (s *Session) WriteString(str string) error {
return s.WebSocket.WriteMessage(websocket.TextMessage, message)
}
func (s *Session) Close() {
if s.GuacdTunnel != nil {
_ = s.GuacdTunnel.Close()
}
if s.NextTerminal != nil {
s.NextTerminal.Close()
}
if s.WebSocket != nil {
_ = s.WebSocket.Close()
}
}
type Manager struct {
id string
sessions map[string]*Session
Add chan *Session
Del chan string
exit chan bool
sessions sync.Map
}
func NewManager() *Manager {
return &Manager{
Add: make(chan *Session),
Del: make(chan string),
sessions: map[string]*Session{},
exit: make(chan bool, 1),
}
return &Manager{}
}
func NewObserver(id string) *Manager {
return &Manager{
id: id,
Add: make(chan *Session),
Del: make(chan string),
sessions: map[string]*Session{},
exit: make(chan bool, 1),
id: id,
}
}
func (m *Manager) Start() {
defer fmt.Printf("Session Manager %v End\n", m.id)
fmt.Printf("Session Manager %v Open\n", m.id)
for {
select {
case s := <-m.Add:
m.sessions[s.ID] = s
case k := <-m.Del:
if _, ok := m.sessions[k]; ok {
ss := m.sessions[k]
if ss.GuacdTunnel != nil {
_ = ss.GuacdTunnel.Close()
}
if ss.NextTerminal != nil {
ss.NextTerminal.Close()
}
func (m *Manager) GetById(id string) *Session {
value, ok := m.sessions.Load(id)
if ok {
return value.(*Session)
}
return nil
}
if ss.WebSocket != nil {
_ = ss.WebSocket.Close()
}
if ss.Observer != nil {
ss.Observer.Close()
}
delete(m.sessions, k)
}
case <-m.exit:
return
func (m *Manager) Add(s *Session) {
m.sessions.Store(s.ID, s)
log.Infof("add session: %s", s.ID)
}
func (m *Manager) Del(id string) {
session := m.GetById(id)
if session != nil {
session.Close()
if session.Observer != nil {
session.Observer.Clear()
}
}
m.sessions.Delete(id)
log.Infof("del session: %s", id)
}
func (m *Manager) Close() {
m.exit <- true
func (m *Manager) Clear() {
m.sessions.Range(func(key, value interface{}) bool {
if session, ok := value.(*Session); ok {
session.Close()
}
m.sessions.Delete(key)
return true
})
}
func (m Manager) GetById(id string) *Session {
return m.sessions[id]
}
func (m Manager) All() map[string]*Session {
return m.sessions
func (m *Manager) Range(f func(key string, value *Session)) {
m.sessions.Range(func(key, value interface{}) bool {
if session, ok := value.(*Session); ok {
f(key.(string), session)
}
return true
})
}
var GlobalSessionManager *Manager
func init() {
GlobalSessionManager = NewManager()
go GlobalSessionManager.Start()
}