sshpoke/internal/plugin/server.go

package plugin

import (
	"context"
	"errors"
	"fmt"
	"net"

	"github.com/Neur0toxine/sshpoke/internal/config"
	"github.com/Neur0toxine/sshpoke/internal/logger"
	"github.com/Neur0toxine/sshpoke/internal/model"
	"github.com/Neur0toxine/sshpoke/internal/server"
	"github.com/Neur0toxine/sshpoke/internal/server/driver/plugin"
	pb "github.com/Neur0toxine/sshpoke/pkg/plugin"
	"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
}

func (p *pluginAPI) Event(_ *emptypb.Empty, stream pb.PluginService_EventServer) error {
	pl := p.receiverForContext(stream.Context())
	if pl == nil {
		return ErrUnauthorized
	}
	logger.Sugar.Debugw("attached plugin event stream", "serverName", pl.Name())
	err := pl.Listen(stream.Context(), &Stream{stream: stream})
	if err != nil {
		logger.Sugar.Debugw("detached plugin event stream", "serverName", pl.Name(), "error", err)
		return err
	}
	logger.Sugar.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.HandleStatus(model.EventRequest{
		ID:    msg.Id,
		Error: msg.Error,
	})
	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 StartAPIServer() {
	port := config.Default.PluginAPIPort
	if port == 0 {
		port = 3000
	}
	socket, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
	if err != nil {
		logger.Sugar.Errorf("cannot start plugin API server on port %d: %s", port, err)
		return
	}
	s := grpc.NewServer()
	pb.RegisterPluginServiceServer(s, &pluginAPI{})
	logger.Sugar.Debugf("starting plugin server on :%d", port)
	if err := s.Serve(socket); err != nil {
		logger.Sugar.Fatalf("cannot start plugin server on :%d: %s", port, err)
	}
}