package plugin import ( "context" "errors" "io" "sync/atomic" "github.com/Neur0toxine/sshpoke/internal/config" "github.com/Neur0toxine/sshpoke/internal/logger" "github.com/Neur0toxine/sshpoke/internal/server/driver/iface" "github.com/Neur0toxine/sshpoke/internal/server/driver/util" "github.com/Neur0toxine/sshpoke/pkg/dto" ) var ErrAlreadyConnected = errors.New("already connected") // Driver plugin uses RPC to communicate with external plugin. type Driver struct { ctx context.Context name string params Params send *Queue[dto.Event] listening atomic.Bool } type EventStream interface { Send(event dto.Event) error } type Plugin interface { iface.Driver Token() string Listen(ctx context.Context, stream EventStream) error HandleStatus(event dto.EventStatus) } func New(ctx context.Context, name string, params config.DriverParams) (iface.Driver, error) { drv := &Driver{ name: name, ctx: ctx, send: NewQueue[dto.Event](), } if err := util.UnmarshalParams(params, &drv.params); err != nil { return nil, err } return drv, nil } func (d *Driver) Handle(event dto.Event) error { if d.isDone() { d.send.Enqueue(dto.Event{Type: dto.EventShutdown}) return nil } d.send.Enqueue(event) return nil } func (d *Driver) Name() string { return d.name } func (d *Driver) Driver() config.DriverType { return config.DriverPlugin } func (d *Driver) Token() string { return d.params.Token } func (d *Driver) 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(): return nil default: event, exists := d.send.Dequeue() if !exists { continue } err := stream.Send(event) if errors.Is(err, io.EOF) { return nil } if err != nil { logger.Sugar.Errorw("error writing event to plugin", "server", d.name, "error", err) return err } } } } func (d *Driver) HandleStatus(event dto.EventStatus) { logger.Sugar.Errorw("plugin error", "serverName", d.name, "id", event.ID, "error", event.Error) } func (d *Driver) isDone() bool { select { case <-d.ctx.Done(): return true default: return false } } func (d *Driver) WaitForShutdown() { <-d.ctx.Done() return }