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