From fc81ebe6509acc68bff59e232e7bcc0b47fad5b1 Mon Sep 17 00:00:00 2001
From: Neur0toxine <pashok9825@gmail.com>
Date: Fri, 17 Nov 2023 20:58:55 +0300
Subject: [PATCH] limit plugin api connections

---
 internal/server/driver/plugin/driver.go | 17 +++++++++++++----
 1 file changed, 13 insertions(+), 4 deletions(-)

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():