internal/agent/syncclient/tcphandler_test.go

Ref: main v0.0.5   Size: 3.3 KiB   History

package syncclient

import (
	"context"
	"io"
	"net"
	"testing"

	"github.com/a73x/eitri/internal/agent/state"
	"github.com/a73x/eitri/internal/pb"
	"github.com/a73x/eitri/internal/transport"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// tcpStore returns a real state.Store in a temp dir with recs saved — the
// smallest fake for the TCPOpen guards, which only read Record.IP by vmID.
func tcpStore(t *testing.T, recs ...state.Record) *state.Store {
	t.Helper()
	st, err := state.Open(t.TempDir())
	require.NoError(t, err)
	for _, r := range recs {
		require.NoError(t, st.SaveVM(r))
	}
	return st
}

// readTCPOpened reads one framed AgentMessage from r and returns its TCPOpened.
func readTCPOpened(t *testing.T, r io.Reader) *pb.TCPOpened {
	t.Helper()
	var msg pb.AgentMessage
	require.NoError(t, transport.ReadMsg(r, &msg, transport.DefaultMaxFrame))
	to := msg.GetTcpOpened()
	require.NotNil(t, to, "reply must be a TCPOpened frame")
	return to
}

func TestTCPOpenRefusedUnknownVM(t *testing.T) {
	c := &Client{St: tcpStore(t)} // no record for vm1
	agentEnd, serverEnd := net.Pipe()
	defer serverEnd.Close()
	go c.handleTCPStream(context.Background(), agentEnd, "vm1", 22)

	to := readTCPOpened(t, serverEnd)
	assert.False(t, to.GetOk())
	assert.Contains(t, to.GetError(), "vm not on this host")
}

func TestTCPOpenRefusedEmptyIP(t *testing.T) {
	c := &Client{St: tcpStore(t, state.Record{Spec: state.VMSpec{VMID: "vm1"}, IP: ""})}
	agentEnd, serverEnd := net.Pipe()
	defer serverEnd.Close()
	go c.handleTCPStream(context.Background(), agentEnd, "vm1", 22)

	to := readTCPOpened(t, serverEnd)
	assert.False(t, to.GetOk())
	assert.Contains(t, to.GetError(), "vm has no address")
}

func TestTCPOpenRefusedPortNotAllowed(t *testing.T) {
	c := &Client{St: tcpStore(t, state.Record{Spec: state.VMSpec{VMID: "vm1"}, IP: "10.0.0.5"})}
	agentEnd, serverEnd := net.Pipe()
	defer serverEnd.Close()
	go c.handleTCPStream(context.Background(), agentEnd, "vm1", 80)

	to := readTCPOpened(t, serverEnd)
	assert.False(t, to.GetOk())
	assert.Contains(t, to.GetError(), "port not allowed")
}

func TestTCPOpenHappyPathRoundTrip(t *testing.T) {
	// An in-process echo listener stands in for the VM's sshd on :22.
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	require.NoError(t, err)
	defer ln.Close()
	go func() {
		conn, err := ln.Accept()
		if err != nil {
			return
		}
		_, _ = io.Copy(conn, conn) // echo until closed
		_ = conn.Close()
	}()

	c := &Client{
		St: tcpStore(t, state.Record{Spec: state.VMSpec{VMID: "vm1"}, IP: "10.0.0.5"}),
		// Production pins :22; the test dialer models that by reaching the
		// in-process echo listener regardless of the guest IP it is handed.
		dialGuest: func(ip string) (net.Conn, error) {
			return net.Dial("tcp", ln.Addr().String())
		},
	}

	agentEnd, serverEnd := net.Pipe()
	defer serverEnd.Close()
	go c.handleTCPStream(context.Background(), agentEnd, "vm1", 22)

	to := readTCPOpened(t, serverEnd)
	require.True(t, to.GetOk(), "happy path should reply ok=true")

	// Bytes written to the stream come back from the echo listener.
	want := []byte("ping")
	go func() { _, _ = serverEnd.Write(want) }()
	got := make([]byte, len(want))
	_, err = io.ReadFull(serverEnd, got)
	require.NoError(t, err)
	assert.Equal(t, want, got, "tunnel round-trips bytes to the VM and back")
}