package plugin import ( "context" "errors" "fmt" "net" "github.com/Neur0toxine/sshpoke/internal/server" "github.com/Neur0toxine/sshpoke/internal/server/driver/plugin" "github.com/Neur0toxine/sshpoke/pkg/convert" "github.com/Neur0toxine/sshpoke/pkg/plugin/pb" "go.uber.org/zap" "google.golang.org/grpc" "google.golang.org/grpc/metadata" "google.golang.org/protobuf/types/known/emptypb" ) var ErrUnauthorized = errors.New("unauthorized") type pluginAPI struct { pb.UnimplementedPluginServiceServer log *zap.SugaredLogger } func (p *pluginAPI) Event(_ *emptypb.Empty, stream pb.PluginService_EventServer) error { pl := p.receiverForContext(stream.Context()) if pl == nil { return ErrUnauthorized } p.log.Debugw("attached plugin event stream", "serverName", pl.Name()) err := pl.Listen(stream.Context(), &Stream{stream: stream}) if err != nil { p.log.Debugw("detached plugin event stream", "serverName", pl.Name(), "error", err) return err } p.log.Debugw("detached plugin event stream", "serverName", pl.Name()) return nil } func (p *pluginAPI) EventStatus(ctx context.Context, msg *pb.EventStatusMessage) (*emptypb.Empty, error) { pl := p.receiverForContext(ctx) if pl == nil { return nil, ErrUnauthorized } pl.PushEventStatus(convert.MessageToAppEventStatus(msg)) return &emptypb.Empty{}, nil } func (p *pluginAPI) receiverForContext(ctx context.Context) plugin.Plugin { md, ok := metadata.FromIncomingContext(ctx) if !ok { return nil } tokens := md.Get("token") if len(tokens) != 1 { return nil } return server.DefaultManager.PluginByToken(tokens[0]) } func StartServer(ctx context.Context, port int, log *zap.SugaredLogger) { socket, err := net.Listen("tcp", fmt.Sprintf(":%d", port)) if err != nil { log.Errorf("cannot start plugin API server on port %d: %s", port, err) return } s := grpc.NewServer() pb.RegisterPluginServiceServer(s, &pluginAPI{log: log}) log.Debugf("starting plugin server on :%d", port) go func() { <-ctx.Done() s.GracefulStop() socket.Close() }() if err := s.Serve(socket); err != nil { log.Fatalf("cannot start plugin server on :%d: %s", port, err) } }