Files
codit/backend/internal/handlers/ssh_broker_runtime.go

142 lines
2.4 KiB
Go

package handlers
import "errors"
import "io"
import "sync"
import "golang.org/x/crypto/ssh"
import "golang.org/x/net/websocket"
type sshActiveSession struct {
mu sync.Mutex
ws *websocket.Conn
client *ssh.Client
session *ssh.Session
stdin io.WriteCloser
closeReason string
closed bool
}
type SSHSessionRegistry struct {
mu sync.Mutex
items map[string]*sshActiveSession
}
func NewSSHSessionRegistry() *SSHSessionRegistry {
var registry *SSHSessionRegistry
registry = &SSHSessionRegistry{
items: map[string]*sshActiveSession{},
}
return registry
}
func (r *SSHSessionRegistry) Register(id string, item *sshActiveSession) error {
if r == nil || item == nil {
return nil
}
r.mu.Lock()
defer r.mu.Unlock()
if r.items[id] != nil {
return errors.New("ssh session already active")
}
r.items[id] = item
return nil
}
func (r *SSHSessionRegistry) Unregister(id string, item *sshActiveSession) {
if r == nil || item == nil {
return
}
r.mu.Lock()
defer r.mu.Unlock()
if r.items[id] == item {
delete(r.items, id)
}
}
func (r *SSHSessionRegistry) RequestClose(id string, reason string) bool {
var item *sshActiveSession
if r == nil {
return false
}
r.mu.Lock()
item = r.items[id]
r.mu.Unlock()
if item == nil {
return false
}
item.RequestClose(reason)
return true
}
func (s *sshActiveSession) SetResources(ws *websocket.Conn, client *ssh.Client, session *ssh.Session, stdin io.WriteCloser) {
if s == nil {
return
}
s.mu.Lock()
if ws != nil {
s.ws = ws
}
if client != nil {
s.client = client
}
if session != nil {
s.session = session
}
if stdin != nil {
s.stdin = stdin
}
s.mu.Unlock()
}
func (s *sshActiveSession) CloseReason() string {
var reason string
if s == nil {
return ""
}
s.mu.Lock()
reason = s.closeReason
s.mu.Unlock()
return reason
}
func (s *sshActiveSession) RequestClose(reason string) {
var ws *websocket.Conn
var client *ssh.Client
var session *ssh.Session
var stdin io.WriteCloser
if s == nil {
return
}
s.mu.Lock()
if s.closeReason == "" {
s.closeReason = reason
}
if s.closed {
s.mu.Unlock()
return
}
s.closed = true
ws = s.ws
client = s.client
session = s.session
stdin = s.stdin
s.mu.Unlock()
if stdin != nil {
_ = stdin.Close()
}
if session != nil {
_ = session.Close()
}
if client != nil {
_ = client.Close()
}
if ws != nil {
_ = ws.Close()
}
}