consul/agent/grpc-middleware/handshake_test.go

package middleware

import (
	"net"
	"testing"

	"github.com/stretchr/testify/require"
	"google.golang.org/grpc/credentials"
)

type fakeTransportCredentials struct {
	credentials.TransportCredentials
	callback func(conn net.Conn) (net.Conn, credentials.AuthInfo, error)
}

func (f fakeTransportCredentials) ServerHandshake(conn net.Conn) (net.Conn, credentials.AuthInfo, error) {
	return f.callback(conn)
}

func TestGRPCMiddleware_optionalTransportCredentials_ServerHandshake(t *testing.T) {

	plainConn := LabelledConn{protocol: ProtocolPlaintext}
	tlsConn := LabelledConn{protocol: ProtocolTLS}
	tests := []struct {
		name           string
		conn           net.Conn
		callback       func(conn net.Conn) (net.Conn, credentials.AuthInfo, error)
		expectConn     net.Conn
		expectAuthInfo credentials.AuthInfo
		expectErr      string
	}{
		{
			name:           "plaintext_noop",
			conn:           plainConn,
			expectConn:     plainConn,
			expectAuthInfo: nil,
		},
		{
			name: "tls_with_missing_auth",
			conn: tlsConn,
			callback: func(conn net.Conn) (net.Conn, credentials.AuthInfo, error) {
				return conn, nil, nil
			},
			expectConn:     nil,
			expectAuthInfo: nil,
			expectErr:      "missing auth info after handshake",
		},
		{
			name: "tls_success",
			conn: tlsConn,
			callback: func(conn net.Conn) (net.Conn, credentials.AuthInfo, error) {
				return conn, credentials.TLSInfo{}, nil
			},
			expectConn:     tlsConn,
			expectAuthInfo: credentials.TLSInfo{},
			expectErr:      "",
		},
		{
			name:           "invalid_protocol",
			conn:           LabelledConn{protocol: -1},
			expectConn:     nil,
			expectAuthInfo: nil,
			expectErr:      "invalid protocol for grpc connection",
		},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			creds := optionalTransportCredentials{
				TransportCredentials: fakeTransportCredentials{
					callback: tc.callback,
				},
			}
			conn, authInfo, err := creds.ServerHandshake(tc.conn)
			require.Equal(t, tc.expectConn, conn)
			require.Equal(t, tc.expectAuthInfo, authInfo)
			if tc.expectErr == "" {
				require.NoError(t, err)
			} else {
				require.Error(t, err)
				require.Contains(t, err.Error(), tc.expectErr)
			}
		})
	}
}