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