diff --git a/server/plugin/plg_backend_sftp/index.go b/server/plugin/plg_backend_sftp/index.go index f59c9839..114d6375 100644 --- a/server/plugin/plg_backend_sftp/index.go +++ b/server/plugin/plg_backend_sftp/index.go @@ -1,14 +1,16 @@ package plg_backend_sftp import ( - . "github.com/mickael-kerjean/filestash/server/common" - "github.com/pkg/sftp" - "golang.org/x/crypto/ssh" "io" "net" "os" "regexp" "strings" + "sync" + + . "github.com/mickael-kerjean/filestash/server/common" + "github.com/pkg/sftp" + "golang.org/x/crypto/ssh" ) var SftpCache AppCache @@ -16,6 +18,7 @@ var SftpCache AppCache type Sftp struct { SSHClient *ssh.Client SFTPClient *sftp.Client + wg *sync.WaitGroup } func init() { @@ -24,6 +27,11 @@ func init() { SftpCache = NewAppCache() SftpCache.OnEvict(func(key string, value interface{}) { c := value.(*Sftp) + if c.wg != nil { + c.wg.Wait() + } else { + Log.Warning("plg_backend_sftp::wg is nil on close") + } if c != nil { c.Close() } @@ -51,6 +59,15 @@ func (s Sftp) Init(params map[string]string, app *App) (IBackend, error) { c := SftpCache.Get(params) if c != nil { d := c.(*Sftp) + if d.wg != nil { + d.wg.Add(1) + go func() { + <-app.Context.Done() + d.wg.Done() + }() + } else { + Log.Warning("plg_backend_sftp::wg is nil on get") + } return d, nil } @@ -146,6 +163,12 @@ func (s Sftp) Init(params map[string]string, app *App) (IBackend, error) { return &s, err } s.SFTPClient = session + s.wg = new(sync.WaitGroup) + s.wg.Add(1) + go func() { + <-app.Context.Done() + s.wg.Done() + }() SftpCache.Set(params, &s) return &s, nil }