1package plugin
2
3import (
4	"bytes"
5	"context"
6	"io"
7	"net"
8	"net/rpc"
9
10	hclog "github.com/hashicorp/go-hclog"
11	"github.com/hashicorp/go-plugin/internal/plugin"
12	"github.com/mitchellh/go-testing-interface"
13	"google.golang.org/grpc"
14)
15
16// TestOptions allows specifying options that can affect the behavior of the
17// test functions
18type TestOptions struct {
19	//ServerStdout causes the given value to be used in place of a blank buffer
20	//for RPCServer's Stdout
21	ServerStdout io.ReadCloser
22
23	//ServerStderr causes the given value to be used in place of a blank buffer
24	//for RPCServer's Stderr
25	ServerStderr io.ReadCloser
26}
27
28// The testing file contains test helpers that you can use outside of
29// this package for making it easier to test plugins themselves.
30
31// TestConn is a helper function for returning a client and server
32// net.Conn connected to each other.
33func TestConn(t testing.T) (net.Conn, net.Conn) {
34	// Listen to any local port. This listener will be closed
35	// after a single connection is established.
36	l, err := net.Listen("tcp", "127.0.0.1:0")
37	if err != nil {
38		t.Fatalf("err: %s", err)
39	}
40
41	// Start a goroutine to accept our client connection
42	var serverConn net.Conn
43	doneCh := make(chan struct{})
44	go func() {
45		defer close(doneCh)
46		defer l.Close()
47		var err error
48		serverConn, err = l.Accept()
49		if err != nil {
50			t.Fatalf("err: %s", err)
51		}
52	}()
53
54	// Connect to the server
55	clientConn, err := net.Dial("tcp", l.Addr().String())
56	if err != nil {
57		t.Fatalf("err: %s", err)
58	}
59
60	// Wait for the server side to acknowledge it has connected
61	<-doneCh
62
63	return clientConn, serverConn
64}
65
66// TestRPCConn returns a rpc client and server connected to each other.
67func TestRPCConn(t testing.T) (*rpc.Client, *rpc.Server) {
68	clientConn, serverConn := TestConn(t)
69
70	server := rpc.NewServer()
71	go server.ServeConn(serverConn)
72
73	client := rpc.NewClient(clientConn)
74	return client, server
75}
76
77// TestPluginRPCConn returns a plugin RPC client and server that are connected
78// together and configured.
79func TestPluginRPCConn(t testing.T, ps map[string]Plugin, opts *TestOptions) (*RPCClient, *RPCServer) {
80	// Create two net.Conns we can use to shuttle our control connection
81	clientConn, serverConn := TestConn(t)
82
83	// Start up the server
84	server := &RPCServer{Plugins: ps, Stdout: new(bytes.Buffer), Stderr: new(bytes.Buffer)}
85	if opts != nil {
86		if opts.ServerStdout != nil {
87			server.Stdout = opts.ServerStdout
88		}
89		if opts.ServerStderr != nil {
90			server.Stderr = opts.ServerStderr
91		}
92	}
93	go server.ServeConn(serverConn)
94
95	// Connect the client to the server
96	client, err := NewRPCClient(clientConn, ps)
97	if err != nil {
98		t.Fatalf("err: %s", err)
99	}
100
101	return client, server
102}
103
104// TestGRPCConn returns a gRPC client conn and grpc server that are connected
105// together and configured. The register function is used to register services
106// prior to the Serve call. This is used to test gRPC connections.
107func TestGRPCConn(t testing.T, register func(*grpc.Server)) (*grpc.ClientConn, *grpc.Server) {
108	// Create a listener
109	l, err := net.Listen("tcp", "127.0.0.1:0")
110	if err != nil {
111		t.Fatalf("err: %s", err)
112	}
113
114	server := grpc.NewServer()
115	register(server)
116	go server.Serve(l)
117
118	// Connect to the server
119	conn, err := grpc.Dial(
120		l.Addr().String(),
121		grpc.WithBlock(),
122		grpc.WithInsecure())
123	if err != nil {
124		t.Fatalf("err: %s", err)
125	}
126
127	// Connection successful, close the listener
128	l.Close()
129
130	return conn, server
131}
132
133// TestPluginGRPCConn returns a plugin gRPC client and server that are connected
134// together and configured. This is used to test gRPC connections.
135func TestPluginGRPCConn(t testing.T, ps map[string]Plugin) (*GRPCClient, *GRPCServer) {
136	// Create a listener
137	l, err := net.Listen("tcp", "127.0.0.1:0")
138	if err != nil {
139		t.Fatalf("err: %s", err)
140	}
141
142	// Start up the server
143	server := &GRPCServer{
144		Plugins: ps,
145		DoneCh:  make(chan struct{}),
146		Server:  DefaultGRPCServer,
147		Stdout:  new(bytes.Buffer),
148		Stderr:  new(bytes.Buffer),
149		logger:  hclog.Default(),
150	}
151	if err := server.Init(); err != nil {
152		t.Fatalf("err: %s", err)
153	}
154	go server.Serve(l)
155
156	// Connect to the server
157	conn, err := grpc.Dial(
158		l.Addr().String(),
159		grpc.WithBlock(),
160		grpc.WithInsecure())
161	if err != nil {
162		t.Fatalf("err: %s", err)
163	}
164
165	brokerGRPCClient := newGRPCBrokerClient(conn)
166	broker := newGRPCBroker(brokerGRPCClient, nil)
167	go broker.Run()
168	go brokerGRPCClient.StartStream()
169
170	// Create the client
171	client := &GRPCClient{
172		Conn:       conn,
173		Plugins:    ps,
174		broker:     broker,
175		doneCtx:    context.Background(),
176		controller: plugin.NewGRPCControllerClient(conn),
177	}
178
179	return client, server
180}
181