open-consul/agent/grpc/server_test.go

package grpc

import (
	"context"
	"fmt"
	"io"
	"net"
	"testing"
	"time"

	"github.com/hashicorp/consul/agent/grpc/internal/testservice"
	"github.com/hashicorp/consul/agent/metadata"
	"github.com/hashicorp/consul/agent/pool"
	"github.com/stretchr/testify/require"
	"golang.org/x/sync/errgroup"
)

type testServer struct {
	addr     net.Addr
	name     string
	dc       string
	shutdown func()
}

func (s testServer) Metadata() *metadata.Server {
	return &metadata.Server{ID: s.name, Datacenter: s.dc, Addr: s.addr}
}

func newTestServer(t *testing.T, name string, dc string) testServer {
	addr := &net.IPAddr{IP: net.ParseIP("127.0.0.1")}
	handler := NewHandler(addr)

	testservice.RegisterSimpleServer(handler.srv, &simple{name: name, dc: dc})

	lis, err := net.Listen("tcp", "127.0.0.1:0")
	require.NoError(t, err)

	rpc := &fakeRPCListener{t: t, handler: handler}

	g := errgroup.Group{}
	g.Go(func() error {
		return rpc.listen(lis)
	})
	g.Go(func() error {
		return handler.Run()
	})
	return testServer{
		addr: lis.Addr(),
		name: name,
		dc:   dc,
		shutdown: func() {
			if err := lis.Close(); err != nil {
				t.Logf("listener closed with error: %v", err)
			}
			if err := handler.Shutdown(); err != nil {
				t.Logf("grpc server shutdown: %v", err)
			}
			if err := g.Wait(); err != nil {
				t.Logf("grpc server error: %v", err)
			}
		},
	}
}

type simple struct {
	name string
	dc   string
}

func (s *simple) Flow(_ *testservice.Req, flow testservice.Simple_FlowServer) error {
	for flow.Context().Err() == nil {
		resp := &testservice.Resp{ServerName: "one", Datacenter: s.dc}
		if err := flow.Send(resp); err != nil {
			return err
		}
		time.Sleep(time.Millisecond)
	}
	return nil
}

func (s *simple) Something(_ context.Context, _ *testservice.Req) (*testservice.Resp, error) {
	return &testservice.Resp{ServerName: s.name, Datacenter: s.dc}, nil
}

// fakeRPCListener mimics agent/consul.Server.listen to handle the RPCType byte.
// In the future we should be able to refactor Server and extract this RPC
// handling logic so that we don't need to use a fake.
// For now, since this logic is in agent/consul, we can't easily use Server.listen
// so we fake it.
type fakeRPCListener struct {
	t       *testing.T
	handler *Handler
}

func (f *fakeRPCListener) listen(listener net.Listener) error {
	for {
		conn, err := listener.Accept()
		if err != nil {
			return err
		}

		go f.handleConn(conn)
	}
}

func (f *fakeRPCListener) handleConn(conn net.Conn) {
	buf := make([]byte, 1)

	if _, err := conn.Read(buf); err != nil {
		if err != io.EOF {
			fmt.Println("ERROR", err.Error())
		}
		conn.Close()
		return
	}
	typ := pool.RPCType(buf[0])

	if typ == pool.RPCGRPC {
		f.handler.Handle(conn)
		return
	}

	fmt.Println("ERROR: unexpected byte", typ)
	conn.Close()
}