diff --git a/internal/server/driver/plugin/driver.go b/internal/server/driver/plugin/driver.go index 8d61248..8ce7923 100644 --- a/internal/server/driver/plugin/driver.go +++ b/internal/server/driver/plugin/driver.go @@ -4,6 +4,7 @@ import ( "context" "errors" "io" + "sync/atomic" "github.com/Neur0toxine/sshpoke/internal/config" "github.com/Neur0toxine/sshpoke/internal/logger" @@ -12,12 +13,15 @@ import ( "github.com/Neur0toxine/sshpoke/internal/server/driver/util" ) +var ErrAlreadyConnected = errors.New("already connected") + // Plugin driver uses RPC to communicate with external plugin. type Plugin struct { - ctx context.Context - name string - params Params - send *Queue[model.Event] + ctx context.Context + name string + params Params + send *Queue[model.Event] + listening atomic.Bool } type EventStream interface { @@ -57,6 +61,11 @@ func (d *Plugin) Token() string { } func (d *Plugin) Listen(ctx context.Context, stream EventStream) error { + if d.listening.Load() { + return ErrAlreadyConnected + } + d.listening.Store(true) + defer d.listening.Store(false) for { select { case <-ctx.Done():