go.mondoo.com/cnquery@v0.0.0-20231005093811-59568235f6ea/shared/grpc.go (about)

     1  // Copyright (c) Mondoo, Inc.
     2  // SPDX-License-Identifier: BUSL-1.1
     3  
     4  package shared
     5  
     6  import (
     7  	"io"
     8  
     9  	hclog "github.com/hashicorp/go-hclog"
    10  	plugin "github.com/hashicorp/go-plugin"
    11  	"go.mondoo.com/cnquery/providers"
    12  	"go.mondoo.com/cnquery/shared/proto"
    13  	"golang.org/x/net/context"
    14  	"google.golang.org/grpc"
    15  )
    16  
    17  func init() {
    18  	var x CNQuery = &GRPCClient{}
    19  	_ = x
    20  }
    21  
    22  // GRPCClient is an implementation of KV that talks over RPC.
    23  type GRPCClient struct {
    24  	broker *plugin.GRPCBroker
    25  	client proto.CNQueryClient
    26  }
    27  
    28  type IOWriter struct {
    29  	io.Writer
    30  }
    31  
    32  func (i *IOWriter) Write(x []byte) (int, error) {
    33  	return i.Writer.Write(x)
    34  }
    35  
    36  func (i *IOWriter) WriteString(x string) error {
    37  	_, err := i.Writer.Write([]byte(x))
    38  	return err
    39  }
    40  
    41  func (m *GRPCClient) RunQuery(conf *proto.RunQueryConfig, runtime *providers.Runtime, out OutputHelper) error {
    42  	helper := &GRPCOutputHelperServer{Impl: &IOWriter{out}}
    43  
    44  	var s *grpc.Server
    45  	serverFunc := func(opts []grpc.ServerOption) *grpc.Server {
    46  		s = grpc.NewServer(opts...)
    47  		proto.RegisterOutputHelperServer(s, helper)
    48  
    49  		return s
    50  	}
    51  
    52  	brokerID := m.broker.NextId()
    53  	conf.CallbackServer = brokerID
    54  	go m.broker.AcceptAndServe(brokerID, serverFunc)
    55  
    56  	_, err := m.client.RunQuery(context.Background(), conf)
    57  
    58  	s.Stop()
    59  	return err
    60  }
    61  
    62  // Here is the gRPC server that GRPCClient talks to.
    63  type GRPCServer struct {
    64  	// This is the real implementation
    65  	Impl   CNQuery
    66  	broker *plugin.GRPCBroker
    67  	proto.UnimplementedCNQueryServer
    68  }
    69  
    70  func (m *GRPCServer) RunQuery(ctx context.Context, req *proto.RunQueryConfig) (*proto.Empty, error) {
    71  	conn, err := m.broker.Dial(req.CallbackServer)
    72  	if err != nil {
    73  		return nil, err
    74  	}
    75  	defer conn.Close()
    76  
    77  	a := &GRPCOutputHelperClient{proto.NewOutputHelperClient(conn)}
    78  	panic("UNCLEAR HOW TO GET RUNTIME YET")
    79  	return &proto.Empty{}, m.Impl.RunQuery(req, nil, a)
    80  }
    81  
    82  // GRPCClient is an implementation of KV that talks over RPC.
    83  type GRPCOutputHelperClient struct{ client proto.OutputHelperClient }
    84  
    85  // NOTE: we do NOT return the number of written bytes, because we felt
    86  // it was not important to have this value at this stage of the tool.
    87  // Once the use-case arises it will be added to the plugin, so please
    88  // submit an issue if you feel it should be there.
    89  // The type is still built this way to support the most common interfaces
    90  // for io.Writer.
    91  func (m *GRPCOutputHelperClient) Write(b []byte) (int, error) {
    92  	_, err := m.client.Write(context.Background(), &proto.String{
    93  		Data: string(b),
    94  	})
    95  	if err != nil {
    96  		hclog.Default().Info("out.Write", "client", "start", "err", err)
    97  		return 0, err
    98  	}
    99  	return 0, nil
   100  }
   101  
   102  func (m *GRPCOutputHelperClient) WriteString(s string) error {
   103  	_, err := m.client.Write(context.Background(), &proto.String{
   104  		Data: s,
   105  	})
   106  	if err != nil {
   107  		hclog.Default().Info("out.Write", "client", "start", "err", err)
   108  		return err
   109  	}
   110  	return nil
   111  }
   112  
   113  // Here is the gRPC server that GRPCClient talks to.
   114  type GRPCOutputHelperServer struct {
   115  	// This is the real implementation
   116  	Impl OutputHelper
   117  	proto.UnsafeOutputHelperServer
   118  }
   119  
   120  var empty proto.Empty
   121  
   122  func (m *GRPCOutputHelperServer) Write(ctx context.Context, req *proto.String) (resp *proto.Empty, err error) {
   123  	return &empty, m.Impl.WriteString(req.Data)
   124  }