github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/engine/ghost/client/ghost_client.go (about)

     1  package client
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"fmt"
     7  	"io"
     8  
     9  	"google.golang.org/grpc"
    10  	"google.golang.org/grpc/credentials/insecure"
    11  
    12  	"github.com/onflow/flow-go/utils/unittest"
    13  
    14  	ghost "github.com/onflow/flow-go/engine/ghost/protobuf"
    15  	"github.com/onflow/flow-go/model/flow"
    16  	"github.com/onflow/flow-go/network"
    17  	"github.com/onflow/flow-go/network/channels"
    18  )
    19  
    20  // GhostClient is a client for the ghost node.
    21  //
    22  // The ghost node is a special node type, used for testing purposes. It can
    23  // "impersonate" any other node role, send messages to other nodes on the
    24  // network, and listen to broadcast messages.
    25  //
    26  // NOTE: currently the ghost node is limited to 1-K messages (ie. messages sent
    27  // to at least 2 other nodes). The ghost node WILL NOT receive a 1-1 message,
    28  // unless the message is explicitly sent to it.
    29  type GhostClient struct {
    30  	rpcClient ghost.GhostNodeAPIClient
    31  	close     func() error
    32  	codec     network.Codec
    33  }
    34  
    35  func NewGhostClient(addr string) (*GhostClient, error) {
    36  
    37  	conn, err := grpc.Dial(addr, grpc.WithTransportCredentials(insecure.NewCredentials()))
    38  	if err != nil {
    39  		return nil, err
    40  	}
    41  
    42  	grpcClient := ghost.NewGhostNodeAPIClient(conn)
    43  
    44  	return &GhostClient{
    45  		rpcClient: grpcClient,
    46  		close:     func() error { return conn.Close() },
    47  		codec:     unittest.NetworkCodec(),
    48  	}, nil
    49  }
    50  
    51  // Close closes the client connection.
    52  func (c *GhostClient) Close() error {
    53  	return c.close()
    54  }
    55  
    56  func (c *GhostClient) Send(ctx context.Context, channel channels.Channel, event interface{}, targetIDs ...flow.Identifier) error {
    57  
    58  	message, err := c.codec.Encode(event)
    59  	if err != nil {
    60  		return fmt.Errorf("could not encode event: %w", err)
    61  	}
    62  
    63  	var targets [][]byte
    64  	for _, t := range targetIDs {
    65  		id := make([]byte, len(t))
    66  		copy(id, t[:])
    67  		targets = append(targets, id)
    68  	}
    69  
    70  	req := ghost.SendEventRequest{
    71  		ChannelId: channel.String(),
    72  		TargetID:  targets,
    73  		Message:   message,
    74  	}
    75  
    76  	_, err = c.rpcClient.SendEvent(ctx, &req)
    77  	if err != nil {
    78  		return fmt.Errorf("failed to send event to the ghost node: %w", err)
    79  	}
    80  	return nil
    81  }
    82  
    83  func (c *GhostClient) Subscribe(ctx context.Context) (*FlowMessageStreamReader, error) {
    84  	req := ghost.SubscribeRequest{}
    85  	stream, err := c.rpcClient.Subscribe(ctx, &req)
    86  	if err != nil {
    87  		return nil, fmt.Errorf("failed to subscribe for events: %w", err)
    88  	}
    89  	return &FlowMessageStreamReader{stream: stream, codec: c.codec}, nil
    90  }
    91  
    92  type FlowMessageStreamReader struct {
    93  	stream ghost.GhostNodeAPI_SubscribeClient
    94  	codec  network.Codec
    95  }
    96  
    97  func (fmsr *FlowMessageStreamReader) Next() (flow.Identifier, interface{}, error) {
    98  	msg, err := fmsr.stream.Recv()
    99  	if errors.Is(err, io.EOF) {
   100  		// read done.
   101  		return flow.ZeroID, nil, fmt.Errorf("end of stream reached: %w", err)
   102  	}
   103  	if err != nil {
   104  		return flow.ZeroID, nil, fmt.Errorf("failed to read stream: %w", err)
   105  	}
   106  
   107  	event, err := fmsr.codec.Decode(msg.GetMessage())
   108  	if err != nil {
   109  		return flow.ZeroID, nil, fmt.Errorf("failed to decode event: %w", err)
   110  	}
   111  
   112  	originID := flow.HashToID(msg.GetSenderID())
   113  
   114  	return originID, event, nil
   115  }