simpler plugin proto

This commit is contained in:
Pavel 2023-11-17 20:53:52 +03:00
parent 9ec7bd14b8
commit 2a40999782
7 changed files with 43 additions and 27 deletions

2
go.mod
View File

@ -11,6 +11,7 @@ require (
github.com/spf13/cobra v1.8.0 github.com/spf13/cobra v1.8.0
github.com/spf13/viper v1.17.0 github.com/spf13/viper v1.17.0
go.uber.org/zap v1.26.0 go.uber.org/zap v1.26.0
golang.design/x/lockfree v0.0.1
golang.org/x/crypto v0.13.0 golang.org/x/crypto v0.13.0
google.golang.org/grpc v1.58.2 google.golang.org/grpc v1.58.2
google.golang.org/protobuf v1.31.0 google.golang.org/protobuf v1.31.0
@ -18,6 +19,7 @@ require (
require ( require (
github.com/Microsoft/go-winio v0.6.1 // indirect github.com/Microsoft/go-winio v0.6.1 // indirect
github.com/changkun/lockfree v0.0.1 // indirect
github.com/distribution/reference v0.5.0 // indirect github.com/distribution/reference v0.5.0 // indirect
github.com/docker/distribution v2.8.3+incompatible // indirect github.com/docker/distribution v2.8.3+incompatible // indirect
github.com/docker/go-units v0.5.0 // indirect github.com/docker/go-units v0.5.0 // indirect

4
go.sum
View File

@ -43,6 +43,8 @@ github.com/BurntSushi/xgb v0.0.0-20160522181843-27f122750802/go.mod h1:IVnqGOEym
github.com/Microsoft/go-winio v0.6.1 h1:9/kr64B9VUZrLm5YYwbGtUJnMgqWVOdUAXu6Migciow= github.com/Microsoft/go-winio v0.6.1 h1:9/kr64B9VUZrLm5YYwbGtUJnMgqWVOdUAXu6Migciow=
github.com/Microsoft/go-winio v0.6.1/go.mod h1:LRdKpFKfdobln8UmuiYcKPot9D2v6svN5+sAH+4kjUM= github.com/Microsoft/go-winio v0.6.1/go.mod h1:LRdKpFKfdobln8UmuiYcKPot9D2v6svN5+sAH+4kjUM=
github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU= github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU=
github.com/changkun/lockfree v0.0.1 h1:5WefVJLglY4IHRqOQmh6Ao6wkJYaJkarshKU8VUtId4=
github.com/changkun/lockfree v0.0.1/go.mod h1:3bKiaXn/iNzIPlSvSOMSVbRQUQtAp8qUAyBUtzU11s4=
github.com/chzyer/logex v1.1.10/go.mod h1:+Ywpsq7O8HXn0nuIou7OrIPyXbp3wmkHB+jjWRnGsAI= github.com/chzyer/logex v1.1.10/go.mod h1:+Ywpsq7O8HXn0nuIou7OrIPyXbp3wmkHB+jjWRnGsAI=
github.com/chzyer/readline v0.0.0-20180603132655-2972be24d48e/go.mod h1:nSuG5e5PlCu98SY8svDHJxuZscDgtXS6KTTbou5AhLI= github.com/chzyer/readline v0.0.0-20180603132655-2972be24d48e/go.mod h1:nSuG5e5PlCu98SY8svDHJxuZscDgtXS6KTTbou5AhLI=
github.com/chzyer/test v0.0.0-20180213035817-a1ea475d72b1/go.mod h1:Q3SI9o4m/ZMnBNeIyt5eFwwo7qiLfzFZmjNmxjkiQlU= github.com/chzyer/test v0.0.0-20180213035817-a1ea475d72b1/go.mod h1:Q3SI9o4m/ZMnBNeIyt5eFwwo7qiLfzFZmjNmxjkiQlU=
@ -242,6 +244,8 @@ go.uber.org/multierr v1.11.0 h1:blXXJkSxSSfBVBlC76pxqeO+LN3aDfLQo+309xJstO0=
go.uber.org/multierr v1.11.0/go.mod h1:20+QtiLqy0Nd6FdQB9TLXag12DsQkrbs3htMFfDN80Y= go.uber.org/multierr v1.11.0/go.mod h1:20+QtiLqy0Nd6FdQB9TLXag12DsQkrbs3htMFfDN80Y=
go.uber.org/zap v1.26.0 h1:sI7k6L95XOKS281NhVKOFCUNIvv9e0w4BF8N3u+tCRo= go.uber.org/zap v1.26.0 h1:sI7k6L95XOKS281NhVKOFCUNIvv9e0w4BF8N3u+tCRo=
go.uber.org/zap v1.26.0/go.mod h1:dtElttAiwGvoJ/vj4IwHBS/gXsEu/pZ50mUIRWuG0so= go.uber.org/zap v1.26.0/go.mod h1:dtElttAiwGvoJ/vj4IwHBS/gXsEu/pZ50mUIRWuG0so=
golang.design/x/lockfree v0.0.1 h1:IHFNwZgM5bnZYWkEbzn5lWHMYr8WsRBdCJ/RBVY0xMM=
golang.design/x/lockfree v0.0.1/go.mod h1:iaZUx6UgZaOdePjzI6wFd+seYMl1i0rsG8+xKvA8c4I=
golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
golang.org/x/crypto v0.0.0-20190510104115-cbcb75029529/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI= golang.org/x/crypto v0.0.0-20190510104115-cbcb75029529/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
golang.org/x/crypto v0.0.0-20190605123033-f99c8df09eb5/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI= golang.org/x/crypto v0.0.0-20190605123033-f99c8df09eb5/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=

View File

@ -23,7 +23,7 @@ type pluginAPI struct {
pb.UnimplementedPluginServiceServer pb.UnimplementedPluginServiceServer
} }
func (p *pluginAPI) Event(stream pb.PluginService_EventServer) error { func (p *pluginAPI) Event(_ *emptypb.Empty, stream pb.PluginService_EventServer) error {
pl := p.receiverForContext(stream.Context()) pl := p.receiverForContext(stream.Context())
if pl == nil { if pl == nil {
return ErrUnauthorized return ErrUnauthorized

View File

@ -11,11 +11,6 @@ type Stream struct {
stream pb.PluginService_EventServer stream pb.PluginService_EventServer
} }
func (s *Stream) Recv() error {
_, err := s.stream.Recv()
return err
}
func (s *Stream) Send(event model.Event) error { func (s *Stream) Send(event model.Event) error {
return s.stream.Send(s.eventToMessage(event)) return s.stream.Send(s.eventToMessage(event))
} }

View File

@ -17,19 +17,18 @@ type Plugin struct {
ctx context.Context ctx context.Context
name string name string
params Params params Params
send chan model.Event send *Queue[model.Event]
} }
type EventStream interface { type EventStream interface {
Send(event model.Event) error Send(event model.Event) error
Recv() error
} }
func New(ctx context.Context, name string, params config.DriverParams) (iface.Driver, error) { func New(ctx context.Context, name string, params config.DriverParams) (iface.Driver, error) {
drv := &Plugin{ drv := &Plugin{
name: name, name: name,
ctx: ctx, ctx: ctx,
send: make(chan model.Event), send: NewQueue[model.Event](),
} }
if err := util.UnmarshalParams(params, &drv.params); err != nil { if err := util.UnmarshalParams(params, &drv.params); err != nil {
return nil, err return nil, err
@ -41,7 +40,7 @@ func (d *Plugin) Handle(event model.Event) error {
if d.isDone() { if d.isDone() {
return nil return nil
} }
d.send <- event d.send.Enqueue(event)
return nil return nil
} }
@ -63,21 +62,10 @@ func (d *Plugin) Listen(ctx context.Context, stream EventStream) error {
case <-ctx.Done(): case <-ctx.Done():
return nil return nil
default: default:
} event, exists := d.send.Dequeue()
if !exists {
err := stream.Recv() continue
if errors.Is(err, io.EOF) { }
return nil
}
if err != nil {
logger.Sugar.Errorw("error reading poll event from plugin",
"server", d.name, "error", err)
return err
}
select {
case <-ctx.Done():
return nil
case event := <-d.send:
err := stream.Send(event) err := stream.Send(event)
if errors.Is(err, io.EOF) { if errors.Is(err, io.EOF) {
return nil return nil
@ -98,7 +86,6 @@ func (d *Plugin) HandleStatus(event model.EventRequest) {
func (d *Plugin) isDone() bool { func (d *Plugin) isDone() bool {
select { select {
case <-d.ctx.Done(): case <-d.ctx.Done():
close(d.send)
return true return true
default: default:
return false return false

View File

@ -0,0 +1,28 @@
package plugin
import "golang.design/x/lockfree"
type Queue[T any] struct {
queue *lockfree.Queue
}
func NewQueue[T any]() *Queue[T] {
return &Queue[T]{queue: lockfree.NewQueue()}
}
func (q *Queue[T]) Enqueue(val T) {
q.queue.Enqueue(val)
}
func (q *Queue[T]) Dequeue() (result T, ok bool) {
item := q.queue.Dequeue()
if item == nil {
ok = false
return
}
return item.(T), true
}
func (q *Queue[T]) Len() uint64 {
return q.queue.Length()
}

View File

@ -5,7 +5,7 @@ option go_package = "github.com/Neur0toxine/sshpoke/pkg/plugin";
option java_multiple_files = true; option java_multiple_files = true;
service PluginService { service PluginService {
rpc Event (stream google.protobuf.Empty) returns (stream EventMessage); rpc Event (google.protobuf.Empty) returns (stream EventMessage);
rpc EventStatus (EventStatusMessage) returns (google.protobuf.Empty); rpc EventStatus (EventStatusMessage) returns (google.protobuf.Empty);
rpc Shutdown (stream google.protobuf.Empty) returns (google.protobuf.Empty); rpc Shutdown (stream google.protobuf.Empty) returns (google.protobuf.Empty);
} }