github.com/Ilhicas/nomad@v1.0.4-0.20210304152020-e86851182bc3/client/allocrunner/consul_http_sock_hook_test.go (about)

     1  package allocrunner
     2  
     3  import (
     4  	"bytes"
     5  	"net"
     6  	"path/filepath"
     7  	"testing"
     8  
     9  	"github.com/hashicorp/nomad/client/allocdir"
    10  	"github.com/hashicorp/nomad/helper/testlog"
    11  	"github.com/hashicorp/nomad/nomad/mock"
    12  	"github.com/hashicorp/nomad/nomad/structs/config"
    13  	"github.com/stretchr/testify/require"
    14  )
    15  
    16  func TestConsulSocketHook_PrerunPostrun_Ok(t *testing.T) {
    17  	t.Parallel()
    18  
    19  	fakeConsul, err := net.Listen("tcp", "127.0.0.1:0")
    20  	require.NoError(t, err)
    21  	defer fakeConsul.Close()
    22  
    23  	consulConfig := &config.ConsulConfig{
    24  		Addr: fakeConsul.Addr().String(),
    25  	}
    26  
    27  	alloc := mock.ConnectNativeAlloc("bridge")
    28  
    29  	logger := testlog.HCLogger(t)
    30  
    31  	allocDir, cleanupDir := allocdir.TestAllocDir(t, logger, "ConnectNativeTask")
    32  	defer cleanupDir()
    33  
    34  	// start unix socket proxy
    35  	h := newConsulHTTPSocketHook(logger, alloc, allocDir, consulConfig)
    36  	require.NoError(t, h.Prerun())
    37  
    38  	httpSocket := filepath.Join(allocDir.AllocDir, allocdir.AllocHTTPSocket)
    39  	taskCon, err := net.Dial("unix", httpSocket)
    40  	require.NoError(t, err)
    41  
    42  	// write to consul from task to ensure data is proxied out of the netns
    43  	input := bytes.Repeat([]byte{'X'}, 5*1024)
    44  	errCh := make(chan error, 1)
    45  	go func() {
    46  		_, err := taskCon.Write(input)
    47  		errCh <- err
    48  	}()
    49  
    50  	// accept the connection from inside the netns
    51  	consulConn, err := fakeConsul.Accept()
    52  	require.NoError(t, err)
    53  	defer consulConn.Close()
    54  
    55  	output := make([]byte, len(input))
    56  	_, err = consulConn.Read(output)
    57  	require.NoError(t, err)
    58  	require.NoError(t, <-errCh)
    59  	require.Equal(t, input, output)
    60  
    61  	// read from consul to ensure http response bodies can come back
    62  	input = bytes.Repeat([]byte{'Y'}, 5*1024)
    63  	go func() {
    64  		_, err := consulConn.Write(input)
    65  		errCh <- err
    66  	}()
    67  
    68  	output = make([]byte, len(input))
    69  	_, err = taskCon.Read(output)
    70  	require.NoError(t, err)
    71  	require.NoError(t, <-errCh)
    72  	require.Equal(t, input, output)
    73  
    74  	// stop the unix socket proxy
    75  	require.NoError(t, h.Postrun())
    76  
    77  	// consul reads should now error
    78  	n, err := consulConn.Read(output)
    79  	require.Error(t, err)
    80  	require.Zero(t, n)
    81  
    82  	// task reads and writes should error
    83  	n, err = taskCon.Write(input)
    84  	require.Error(t, err)
    85  	require.Zero(t, n)
    86  	n, err = taskCon.Read(output)
    87  	require.Error(t, err)
    88  	require.Zero(t, n)
    89  }
    90  
    91  func TestConsulHTTPSocketHook_Prerun_Error(t *testing.T) {
    92  	t.Parallel()
    93  
    94  	logger := testlog.HCLogger(t)
    95  
    96  	allocDir, cleanupDir := allocdir.TestAllocDir(t, logger, "ConnectNativeTask")
    97  	defer cleanupDir()
    98  
    99  	consulConfig := new(config.ConsulConfig)
   100  
   101  	alloc := mock.Alloc()
   102  	connectNativeAlloc := mock.ConnectNativeAlloc("bridge")
   103  
   104  	{
   105  		// an alloc without a connect native task should not return an error
   106  		h := newConsulHTTPSocketHook(logger, alloc, allocDir, consulConfig)
   107  		require.NoError(t, h.Prerun())
   108  
   109  		// postrun should be a noop
   110  		require.NoError(t, h.Postrun())
   111  	}
   112  
   113  	{
   114  		// an alloc with a native task should return an error when consul is not
   115  		// configured
   116  		h := newConsulHTTPSocketHook(logger, connectNativeAlloc, allocDir, consulConfig)
   117  		require.EqualError(t, h.Prerun(), "consul address must be set on nomad client")
   118  
   119  		// Postrun should be a noop
   120  		require.NoError(t, h.Postrun())
   121  	}
   122  }