Neur0toxine/sshpoke
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

limit plugin api connections

Browse Source
This commit is contained in:
Pavel 2023-11-17 20:58:55 +03:00
parent 2a40999782
commit fc81ebe650
1 changed files with 13 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
internal/server/driver/plugin/driver.go
View File

@ -4,6 +4,7 @@ import (
"context" "context"
"errors" "errors"
"io" "io"
"sync/atomic"
"github.com/Neur0toxine/sshpoke/internal/config" "github.com/Neur0toxine/sshpoke/internal/config"
"github.com/Neur0toxine/sshpoke/internal/logger" "github.com/Neur0toxine/sshpoke/internal/logger"
@ -12,12 +13,15 @@ import (
"github.com/Neur0toxine/sshpoke/internal/server/driver/util" "github.com/Neur0toxine/sshpoke/internal/server/driver/util"
) )
var ErrAlreadyConnected = errors.New("already connected")
// Plugin driver uses RPC to communicate with external plugin. // Plugin driver uses RPC to communicate with external plugin.
type Plugin struct { type Plugin struct {
ctx context.Context ctx context.Context
name string name string
params Params params Params
send *Queue[model.Event] send *Queue[model.Event]
listening atomic.Bool
} }
type EventStream interface { type EventStream interface {
@ -57,6 +61,11 @@ func (d *Plugin) Token() string {
} }
func (d *Plugin) Listen(ctx context.Context, stream EventStream) error { 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 { for {
select { select {
case <-ctx.Done(): case <-ctx.Done():