sshpoke/internal/api/plugin/server.go

75 lines
2.0 KiB
Go

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(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)
if err := s.Serve(socket); err != nil {
log.Fatalf("cannot start plugin server on :%d: %s", port, err)
}
}