go.mondoo.com/cnquery@v0.0.0-20231005093811-59568235f6ea/providers-sdk/v1/plugin/grpc.go (about)

     1  // Copyright (c) Mondoo, Inc.
     2  // SPDX-License-Identifier: BUSL-1.1
     3  
     4  package plugin
     5  
     6  import (
     7  	plugin "github.com/hashicorp/go-plugin"
     8  	"golang.org/x/net/context"
     9  	"google.golang.org/grpc"
    10  )
    11  
    12  func init() {
    13  	var x ProviderPlugin = &GRPCClient{}
    14  	_ = x
    15  }
    16  
    17  // GRPCClient is an implementation of KV that talks over RPC.
    18  type GRPCClient struct {
    19  	broker *plugin.GRPCBroker
    20  	client ProviderPluginClient
    21  }
    22  
    23  func (m *GRPCClient) ParseCLI(req *ParseCLIReq) (*ParseCLIRes, error) {
    24  	return m.client.ParseCLI(context.Background(), req)
    25  }
    26  
    27  func (m *GRPCClient) connect(req *ConnectReq, callback ProviderCallback) {
    28  	helper := &GRPCProviderCallbackServer{Impl: callback}
    29  
    30  	var s *grpc.Server
    31  	serverFunc := func(opts []grpc.ServerOption) *grpc.Server {
    32  		s = grpc.NewServer(opts...)
    33  		RegisterProviderCallbackServer(s, helper)
    34  		return s
    35  	}
    36  
    37  	brokerID := m.broker.NextId()
    38  	req.CallbackServer = brokerID
    39  	go m.broker.AcceptAndServe(brokerID, serverFunc)
    40  
    41  	// Note: the reverse connection is not closed explicitly. It stays open
    42  	// until the process is eventually stopped. Connect should only be called
    43  	// once per connected asset, thus the reverse connection is also only
    44  	// open for the duration of said connection.
    45  	// In the future, we may want to explicitly disconnect and re-use providers.
    46  }
    47  
    48  func (m *GRPCClient) Connect(req *ConnectReq, callback ProviderCallback) (*ConnectRes, error) {
    49  	m.connect(req, callback)
    50  	return m.client.Connect(context.Background(), req)
    51  }
    52  
    53  func (m *GRPCClient) MockConnect(req *ConnectReq, callback ProviderCallback) (*ConnectRes, error) {
    54  	m.connect(req, callback)
    55  	return m.client.MockConnect(context.Background(), req)
    56  }
    57  
    58  func (m *GRPCClient) Shutdown(req *ShutdownReq) (*ShutdownRes, error) {
    59  	return m.client.Shutdown(context.Background(), req)
    60  }
    61  
    62  func (m *GRPCClient) GetData(req *DataReq) (*DataRes, error) {
    63  	return m.client.GetData(context.Background(), req)
    64  }
    65  
    66  func (m *GRPCClient) StoreData(req *StoreReq) (*StoreRes, error) {
    67  	return m.client.StoreData(context.Background(), req)
    68  }
    69  
    70  // Here is the gRPC server that GRPCClient talks to.
    71  type GRPCServer struct {
    72  	// This is the real implementation
    73  	Impl   ProviderPlugin
    74  	broker *plugin.GRPCBroker
    75  	UnimplementedProviderPluginServer
    76  }
    77  
    78  func (m *GRPCServer) ParseCLI(ctx context.Context, req *ParseCLIReq) (*ParseCLIRes, error) {
    79  	return m.Impl.ParseCLI(req)
    80  }
    81  
    82  func (m *GRPCServer) Connect(ctx context.Context, req *ConnectReq) (*ConnectRes, error) {
    83  	conn, err := m.broker.Dial(req.CallbackServer)
    84  	if err != nil {
    85  		return nil, err
    86  	}
    87  
    88  	// Note: we do not close the connection from this side. It will get closed
    89  	// when the plugin caller decides to kill the process.
    90  
    91  	a := &GRPCProviderCallbackClient{NewProviderCallbackClient(conn)}
    92  	return m.Impl.Connect(req, a)
    93  }
    94  
    95  func (m *GRPCServer) MockConnect(ctx context.Context, req *ConnectReq) (*ConnectRes, error) {
    96  	conn, err := m.broker.Dial(req.CallbackServer)
    97  	if err != nil {
    98  		return nil, err
    99  	}
   100  
   101  	// Note: we do not close the connection from this side. It will get closed
   102  	// when the plugin caller decides to kill the process.
   103  
   104  	a := &GRPCProviderCallbackClient{NewProviderCallbackClient(conn)}
   105  	return m.Impl.MockConnect(req, a)
   106  }
   107  
   108  func (m *GRPCServer) Shutdown(ctx context.Context, req *ShutdownReq) (*ShutdownRes, error) {
   109  	return m.Impl.Shutdown(req)
   110  }
   111  
   112  func (m *GRPCServer) GetData(ctx context.Context, req *DataReq) (*DataRes, error) {
   113  	return m.Impl.GetData(req)
   114  }
   115  
   116  func (m *GRPCServer) StoreData(ctx context.Context, req *StoreReq) (*StoreRes, error) {
   117  	return m.Impl.StoreData(req)
   118  }
   119  
   120  // GRPCClient is an implementation of ProviderCallback that talks over RPC.
   121  type GRPCProviderCallbackClient struct{ client ProviderCallbackClient }
   122  
   123  func (m *GRPCProviderCallbackClient) Collect(req *DataRes) error {
   124  	_, err := m.client.Collect(context.Background(), req)
   125  	return err
   126  }
   127  
   128  func (m *GRPCProviderCallbackClient) GetRecording(req *DataReq) (*ResourceData, error) {
   129  	return m.client.GetRecording(context.Background(), req)
   130  }
   131  
   132  func (m *GRPCProviderCallbackClient) GetData(req *DataReq) (*DataRes, error) {
   133  	return m.client.GetData(context.Background(), req)
   134  }
   135  
   136  // Here is the gRPC server that GRPCClient talks to.
   137  type GRPCProviderCallbackServer struct {
   138  	// This is the real implementation
   139  	Impl ProviderCallback
   140  	UnsafeProviderCallbackServer
   141  }
   142  
   143  var empty CollectRes
   144  
   145  func (m *GRPCProviderCallbackServer) Collect(ctx context.Context, req *DataRes) (resp *CollectRes, err error) {
   146  	return &empty, m.Impl.Collect(req)
   147  }
   148  
   149  func (m *GRPCProviderCallbackServer) GetRecording(ctx context.Context, req *DataReq) (resp *ResourceData, err error) {
   150  	return m.Impl.GetRecording(req)
   151  }
   152  
   153  func (m *GRPCProviderCallbackServer) GetData(ctx context.Context, req *DataReq) (resp *DataRes, err error) {
   154  	return m.Impl.GetData(req)
   155  }