From 6dbfc64ad5d665af0ac6d4af7ca924c4056306b9 Mon Sep 17 00:00:00 2001 From: Pulkit Kathuria Date: Sat, 22 Jun 2024 12:59:11 +0900 Subject: [PATCH] (fix) reuse ssh client (#17) * (fix) reuse ssh client donot create multiple clients if client already exists this is to not kill the remote server with a lot of ssh open agents * (feat) auto recover a client * (feat) minor UI improvement --- .../components/partials/viewer/Inputs.astro | 4 +- go.mod | 1 + go.sum | 2 + pkg/api_handler.go | 4 -- pkg/files.go | 22 ++++---- pkg/globals.go | 10 ++++ pkg/ssh.go | 53 +++++++++++++++++++ pkg/types.go | 2 + pkg/watcher.go | 20 ++++--- 9 files changed, 89 insertions(+), 29 deletions(-) create mode 100644 pkg/globals.go create mode 100644 pkg/ssh.go diff --git a/frontend/src/components/partials/viewer/Inputs.astro b/frontend/src/components/partials/viewer/Inputs.astro index 4732e9a..eed5d04 100644 --- a/frontend/src/components/partials/viewer/Inputs.astro +++ b/frontend/src/components/partials/viewer/Inputs.astro @@ -14,7 +14,7 @@ type="text" x-model="input.query" @keyup="submit" - class="block font-mono w-full p-3 ps-32 text-sm text-slate-300 hover:text-slate-200 border border-gray-600 hover:border-green-500 focus:border-green-500 rounded-lg bg-gray-900 focus:outline-none" + class="block font-mono w-full p-3 ps-32 text-sm text-slate-300 hover:text-slate-200 border-2 border-gray-600 hover:border-green-700 focus:border-green-500 rounded-lg bg-gray-900 focus:outline-none" placeholder=".*" /> diff --git a/go.mod b/go.mod index 6708960..f606ec8 100644 --- a/go.mod +++ b/go.mod @@ -9,6 +9,7 @@ require ( github.com/gookit/color v1.5.4 github.com/gravwell/gravwell/v3 v3.8.32 github.com/k0kubun/pp v3.0.1+incompatible + github.com/k0kubun/pp/v3 v3.2.0 github.com/kevincobain2000/go-human-uuid v0.0.0-20240321072653-0174cf7003de github.com/labstack/echo/v4 v4.12.0 github.com/mcuadros/go-defaults v1.2.0 diff --git a/go.sum b/go.sum index 0418520..4499d26 100644 --- a/go.sum +++ b/go.sum @@ -97,6 +97,8 @@ github.com/k0kubun/colorstring v0.0.0-20150214042306-9440f1994b88 h1:uC1QfSlInpQ github.com/k0kubun/colorstring v0.0.0-20150214042306-9440f1994b88/go.mod h1:3w7q1U84EfirKl04SVQ/s7nPm1ZPhiXd34z40TNz36k= github.com/k0kubun/pp v3.0.1+incompatible h1:3tqvf7QgUnZ5tXO6pNAZlrvHgl6DvifjDrd9g2S9Z40= github.com/k0kubun/pp v3.0.1+incompatible/go.mod h1:GWse8YhT0p8pT4ir3ZgBbfZild3tgzSScAn6HmfYukg= +github.com/k0kubun/pp/v3 v3.2.0 h1:h33hNTZ9nVFNP3u2Fsgz8JXiF5JINoZfFq4SvKJwNcs= +github.com/k0kubun/pp/v3 v3.2.0/go.mod h1:ODtJQbQcIRfAD3N+theGCV1m/CBxweERz2dapdz1EwA= github.com/kevincobain2000/go-human-uuid v0.0.0-20240321072653-0174cf7003de h1:KZUjqJeNpkuyfo8g5g/5DuaMtA0NRf0uDUg1cb+TyWQ= github.com/kevincobain2000/go-human-uuid v0.0.0-20240321072653-0174cf7003de/go.mod h1:pwoguytL8YNxXpKQRE7XrnAstOJlDf7WFO8EUEAYtLI= github.com/kisielk/errcheck v1.5.0/go.mod h1:pFxgyoBC7bSaBwPgfKdkLd5X25qrDl4LWUI2bnpBCr8= diff --git a/pkg/api_handler.go b/pkg/api_handler.go index 5f9f85f..e1e258e 100644 --- a/pkg/api_handler.go +++ b/pkg/api_handler.go @@ -20,10 +20,6 @@ type FileInfo struct { Host string `json:"host"` } -var GlobalFilePaths []FileInfo -var GlobalPipeTmpFilePath string -var GlobalPathSSHConfig []SSHPathConfig - func NewAPIHandler() *APIHandler { return &APIHandler{ API: NewAPI(), diff --git a/pkg/files.go b/pkg/files.go index 6d67d25..9179aec 100644 --- a/pkg/files.go +++ b/pkg/files.go @@ -166,6 +166,7 @@ func GetFileInfos(pattern string, limit int, isRemote bool, sshConfig *SSHConfig color.Warn.Printf("limiting to %d files\n", limit) filePaths = filePaths[:limit] } + for _, filePath := range filePaths { isText, err := IsReadableFile(filePath, isRemote, sshConfig, false) if err != nil { @@ -317,11 +318,7 @@ func sshConnect(config *SSHConfig) (*ssh.Client, error) { } func sshOpenFile(filename string, config *SSHConfig) (*os.File, error) { - client, err := sshConnect(config) - if err != nil { - return nil, err - } - session, err := client.NewSession() + session, err := NewSession(config) if err != nil { return nil, err } @@ -336,7 +333,9 @@ func sshOpenFile(filename string, config *SSHConfig) (*os.File, error) { var stdout bytes.Buffer session.Stdout = &stdout if err := session.Run("cat " + filename); err != nil { - return nil, err + if err.Error() != ErrorMsgSessionAlreadyStarted { + return nil, err + } } // Write the remote file content to the temporary file @@ -353,11 +352,8 @@ func sshOpenFile(filename string, config *SSHConfig) (*os.File, error) { } func sshFilesByPattern(pattern string, config *SSHConfig) ([]string, error) { - client, err := sshConnect(config) - if err != nil { - return nil, err - } - session, err := client.NewSession() + session, err := NewSession(config) + if err != nil { return nil, err } @@ -368,7 +364,9 @@ func sshFilesByPattern(pattern string, config *SSHConfig) ([]string, error) { // Execute the ls command to list files matching the pattern if err := session.Run("ls " + pattern); err != nil { - return nil, err + if err.Error() != ErrorMsgSessionAlreadyStarted { + return nil, err + } } filePaths := buf.String() diff --git a/pkg/globals.go b/pkg/globals.go new file mode 100644 index 0000000..28bb749 --- /dev/null +++ b/pkg/globals.go @@ -0,0 +1,10 @@ +package pkg + +import ( + "golang.org/x/crypto/ssh" +) + +var GlobalFilePaths []FileInfo +var GlobalPipeTmpFilePath string +var GlobalPathSSHConfig []SSHPathConfig +var GlobalSSHClients = make(map[string]*ssh.Client) diff --git a/pkg/ssh.go b/pkg/ssh.go new file mode 100644 index 0000000..ad4850f --- /dev/null +++ b/pkg/ssh.go @@ -0,0 +1,53 @@ +package pkg + +import ( + "sync" + + "golang.org/x/crypto/ssh" +) + +var ( + clientMutex = &sync.Mutex{} +) + +func NewSession(config *SSHConfig) (*ssh.Session, error) { + client, err := NewOrReusableClient(config) + if err != nil { + return nil, err + } + session, err := client.NewSession() + if err != nil { + return nil, err + } + return session, nil +} + +func NewOrReusableClient(config *SSHConfig) (*ssh.Client, error) { + key := config.Host + ":" + config.Port + + clientMutex.Lock() + client := GlobalSSHClients[key] + clientMutex.Unlock() + + if client == nil { + c, err := sshConnect(config) + if err != nil { + return nil, err + } + clientMutex.Lock() + GlobalSSHClients[key] = c + clientMutex.Unlock() + client = c + } + + // check if client is still connected + _, _, err := client.SendRequest("", true, nil) + if err != nil { + clientMutex.Lock() + delete(GlobalSSHClients, key) + clientMutex.Unlock() + return NewOrReusableClient(config) + } + + return client, nil +} diff --git a/pkg/types.go b/pkg/types.go index 7954179..40e6568 100644 --- a/pkg/types.go +++ b/pkg/types.go @@ -7,4 +7,6 @@ const ( TypeDocker = "docker" TmpStdinPath = "/tmp/GOL-STDIN-" TmpContainerPath = "/tmp/GOL-CONTAINER-" + + ErrorMsgSessionAlreadyStarted = "ssh: session already started" ) diff --git a/pkg/watcher.go b/pkg/watcher.go index 77330c2..c3a44ad 100644 --- a/pkg/watcher.go +++ b/pkg/watcher.go @@ -156,26 +156,24 @@ func (w *Watcher) initializeScanner() (*os.File, *bufio.Scanner, error) { } func (w *Watcher) initializeRemoteScanner() (*os.File, *bufio.Scanner, error) { - client, err := ssh.Dial("tcp", fmt.Sprintf("%s:%s", w.sshHost, w.sshPort), w.sshConfig) - if err != nil { - return nil, nil, tracerr.New(err.Error()) + sshConfig := SSHConfig{ + Host: w.sshHost, + Port: w.sshPort, } - - session, err := client.NewSession() + session, err := NewSession(&sshConfig) if err != nil { return nil, nil, tracerr.New(err.Error()) } + defer session.Close() var b bytes.Buffer session.Stdout = &b - err = session.Run(fmt.Sprintf("cat %s", w.filePath)) - if err != nil { - return nil, nil, tracerr.New(err.Error()) + if err := session.Run(fmt.Sprintf("cat %s", w.filePath)); err != nil { + if err.Error() != ErrorMsgSessionAlreadyStarted { + return nil, nil, tracerr.New(err.Error()) + } } - session.Close() - client.Close() - scanner := bufio.NewScanner(strings.NewReader(b.String())) return nil, scanner, nil