cli/e2e/relay/gateway_test.go

package relay_test

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"sync"
	"testing"
	"time"

	"github.com/compose-spec/compose-go/v2/types"
	"github.com/google/uuid"
	"github.com/hanzokms/cli/e2e-tests/packages/client"
	helpers "github.com/hanzokms/cli/e2e-tests/util"
	openapitypes "github.com/oapi-codegen/runtime/types"
	kv "github.com/hanzokv/go/v9"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	testcontainers "github.com/testcontainers/testcontainers-go"
	"github.com/testcontainers/testcontainers-go/wait"
)

func TestGateway_RegistersAGateway(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	t.Cleanup(cancel)

	kms := helpers.NewKMSService().
		WithBackendEnvironment(types.NewMappingWithEquals([]string{
			// This is needed for the private ip (current host) to be accepted for the relay server
			"ALLOW_INTERNAL_IP_CONNECTIONS=true",
		})).
		Up(t, ctx)

	c := kms.ApiClient()
	identity := kms.CreateMachineIdentity(t, ctx, helpers.WithTokenAuth())
	require.NotNil(t, identity)

	relayName := helpers.RandomSlug(2)
	relayCmd := helpers.Command{
		Test: t,
		Args: []string{"relay", "start", "--domain", kms.ApiUrl(t)},
		Env: map[string]string{
			"KMS_API_URL":    kms.ApiUrl(t),
			"KMS_RELAY_NAME": relayName,
			"KMS_RELAY_HOST": "host.docker.internal",
			"KMS_TOKEN":      *identity.TokenAuthToken,
		},
	}
	relayCmd.Start(ctx)
	t.Cleanup(relayCmd.Stop)
	result := helpers.WaitForStderr(t, helpers.WaitForStderrOptions{
		EnsureCmdRunning: &relayCmd,
		ExpectedString:   "Relay is reachable by KMS",
	})
	require.Equal(t, helpers.WaitSuccess, result)

	tmpLogDir := t.TempDir()
	sessionRecordingPath := filepath.Join(tmpLogDir, "session-recording")
	require.NoError(t, os.MkdirAll(sessionRecordingPath, 0755))
	gatewayName := helpers.RandomSlug(2)
	gatewayCmd := helpers.Command{
		Test: t,
		Args: []string{"gateway", "start",
			fmt.Sprintf("--name=%s", gatewayName),
			fmt.Sprintf("--pam-session-recording-path=%s", sessionRecordingPath),
		},
		Env: map[string]string{
			"KMS_API_URL": kms.ApiUrl(t),
			"KMS_TOKEN":   *identity.TokenAuthToken,
		},
	}
	gatewayCmd.Start(ctx)
	t.Cleanup(gatewayCmd.Stop)

	result = helpers.WaitForStderr(t, helpers.WaitForStderrOptions{
		EnsureCmdRunning: &gatewayCmd,
		ExpectedString:   "Successfully registered gateway and received certificates",
	})
	require.Equal(t, helpers.WaitSuccess, result)

	result = helpers.WaitFor(t, helpers.WaitForOptions{
		EnsureCmdRunning: &gatewayCmd,
		Condition: func() helpers.ConditionResult {
			resp, err := c.ListGatewaysWithResponse(ctx)
			if err != nil {
				return helpers.ConditionWait
			}
			if resp.StatusCode() != http.StatusOK {
				return helpers.ConditionWait
			}
			for _, gateway := range *resp.JSON200 {
				slog.Info(
					"Gateway info",
					"id", gateway.Id,
					"name", gateway.Name,
					"identityId", gateway.IdentityId,
					"heartbeat", gateway.Heartbeat,
				)
				if gateway.Name == gatewayName && gateway.Heartbeat != nil {
					slog.Info("Confirmed gateway heartbeat")
					return helpers.ConditionSuccess
				}
			}
			return helpers.ConditionWait
		},
	})
	require.Equal(t, helpers.WaitSuccess, result)

	result = helpers.WaitForStderr(t, helpers.WaitForStderrOptions{
		EnsureCmdRunning: &gatewayCmd,
		ExpectedString:   "Gateway is reachable by KMS",
	})
	assert.Equal(t, helpers.WaitSuccess, result)
}

func TestGateway_RelayGatewayConnectivity(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	t.Cleanup(cancel)

	kms := helpers.NewKMSService().
		WithBackendEnvironment(types.NewMappingWithEquals([]string{
			// This is needed for the private ip (current host) to be accepted for the relay server
			"ALLOW_INTERNAL_IP_CONNECTIONS=true",
		})).
		Up(t, ctx)

	identity := kms.CreateMachineIdentity(t, ctx, helpers.WithTokenAuth())
	require.NotNil(t, identity)

	relayName := helpers.RandomSlug(2)
	relayCmd := helpers.Command{
		Test: t,
		Args: []string{"relay", "start", "--domain", kms.ApiUrl(t)},
		Env: map[string]string{
			"KMS_API_URL":    kms.ApiUrl(t),
			"KMS_RELAY_NAME": relayName,
			"KMS_RELAY_HOST": "host.docker.internal",
			"KMS_TOKEN":      *identity.TokenAuthToken,
		},
	}
	relayCmd.Start(ctx)
	t.Cleanup(relayCmd.Stop)
	result := helpers.WaitForStderr(t, helpers.WaitForStderrOptions{
		EnsureCmdRunning: &relayCmd,
		ExpectedString:   "Relay is reachable by KMS",
	})
	require.Equal(t, helpers.WaitSuccess, result)

	tmpLogDir := t.TempDir()
	sessionRecordingPath := filepath.Join(tmpLogDir, "session-recording")
	require.NoError(t, os.MkdirAll(sessionRecordingPath, 0755))
	gatewayName := helpers.RandomSlug(2)
	gatewayCmd := helpers.Command{
		Test: t,
		Args: []string{"gateway", "start",
			fmt.Sprintf("--name=%s", gatewayName),
			fmt.Sprintf("--pam-session-recording-path=%s", sessionRecordingPath),
		},
		Env: map[string]string{
			"KMS_API_URL": kms.ApiUrl(t),
			"KMS_TOKEN":   *identity.TokenAuthToken,
		},
	}
	gatewayCmd.Start(ctx)
	t.Cleanup(gatewayCmd.Stop)
	result = helpers.WaitForStderr(t, helpers.WaitForStderrOptions{
		EnsureCmdRunning: &gatewayCmd,
		ExpectedString:   "Gateway is reachable by KMS",
	})
	assert.Equal(t, helpers.WaitSuccess, result)

	c := kms.ApiClient()
	var gatewayId openapitypes.UUID
	resp, err := c.ListGatewaysWithResponse(ctx)
	require.NoError(t, err)
	require.Equal(t, http.StatusOK, resp.StatusCode())
	for _, gateway := range *resp.JSON200 {
		slog.Info(
			"Gateway info",
			"id", gateway.Id,
			"name", gateway.Name,
			"identityId", gateway.IdentityId,
			"heartbeat", gateway.Heartbeat,
		)
		if gateway.Name == gatewayName && gateway.Heartbeat != nil {
			gatewayId = gateway.Id
			slog.Info("Found gateway ID", "gatewayId", gatewayId)
			break
		}
	}
	require.NotZero(t, gatewayId, "Gateway ID should be set")

	projDesc := "e2e tests for PAM connectivity"
	template := "default"
	projectType := client.Pam
	projectResp, err := c.CreateProjectWithResponse(ctx, client.CreateProjectJSONRequestBody{
		ProjectName:        "pam-tests",
		ProjectDescription: &projDesc,
		Template:           &template,
		Type:               &projectType,
	})
	require.NoError(t, err)
	require.Equal(t, projectResp.StatusCode(), http.StatusOK)
	projectId := projectResp.JSON200.Project.Id

	t.Run("kubernetes", func(t *testing.T) {
		t.Parallel()
		ctx := t.Context()
		// Create a mock HTTP server running on a random port in a goroutine
		// The HTTP server implements a mock /version endpoint that returns dummy data
		// and marks a variable as true when the endpoint is hit
		var versionEndpointHit bool
		var versionEndpointHitMu sync.Mutex

		// Create a listener on a random port (port 0 means OS assigns an available port)
		listener, err := net.Listen("tcp", ":0")
		require.NoError(t, err)

		server := &http.Server{
			Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				if r.URL.Path == "/version" {
					versionEndpointHitMu.Lock()
					versionEndpointHit = true
					versionEndpointHitMu.Unlock()

					w.Header().Set("Content-Type", "application/json")
					w.WriteHeader(http.StatusOK)
					// Return dummy version data
					versionData := map[string]interface{}{
						"version": "1.0.0",
						"build":   "test-build",
					}
					json.NewEncoder(w).Encode(versionData)
				} else {
					w.WriteHeader(http.StatusNotFound)
				}
			}),
		}

		// Start the server in a goroutine
		go func() {
			if err := server.Serve(listener); err != nil && err != http.ErrServerClosed {
				t.Errorf("Mock HTTP server error: %v", err)
			}
		}()

		// Clean up the server when the test completes
		t.Cleanup(func() {
			shutdownCtx, shutdownCancel := context.WithTimeout(ctx, 5*time.Second)
			defer shutdownCancel()
			server.Shutdown(shutdownCtx)
		})

		// Get the server URL
		serverURL := fmt.Sprintf("http://%s", listener.Addr().String())
		slog.Info("Mock HTTP server started", "url", serverURL)

		k8sPamResResp, err := c.CreateKubernetesPamResourceWithResponse(
			ctx,
			client.CreateKubernetesPamResourceJSONRequestBody{
				ProjectId: uuid.MustParse(projectId),
				GatewayId: gatewayId,
				Name:      "k8s-resource",
				ConnectionDetails: struct {
					SslCertificate        *string `json:"sslCertificate,omitempty"`
					SslRejectUnauthorized bool    `json:"sslRejectUnauthorized"`
					Url                   string  `json:"url"`
				}{
					Url:                   serverURL,
					SslRejectUnauthorized: false,
				},
			})
		require.NoError(t, err)
		require.Equal(t, k8sPamResResp.StatusCode(), http.StatusOK)
		require.True(t, versionEndpointHit)
	})

	t.Run("kv", func(t *testing.T) {
		t.Parallel()
		ctx := t.Context()
		// The server under test is our own, so the container is our own image.
		kvContainer, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
			ContainerRequest: testcontainers.ContainerRequest{
				Image:        "ghcr.io/hanzoai/kv",
				ExposedPorts: []string{"6379/tcp"},
				WaitingFor:   wait.ForLog("Ready to accept connections"),
			},
			Started: true,
		})
		require.NoError(t, err)
		t.Cleanup(func() {
			if err := kvContainer.Terminate(ctx); err != nil {
				t.Logf("Failed to terminate KV container: %v", err)
			}
		})

		kvHost, err := kvContainer.Host(ctx)
		require.NoError(t, err)
		kvPort, err := kvContainer.MappedPort(ctx, "6379")
		require.NoError(t, err)
		slog.Info("KV container", "host", kvHost, "port", kvPort.Port())

		// Addr directly rather than a URL: the container reports a mapped port,
		// and host:port is what the client wants anyway.
		rdb := kv.NewClient(&kv.Options{Addr: net.JoinHostPort(kvHost, kvPort.Port())})
		t.Cleanup(func() { rdb.Close() })

		pong, err := rdb.Ping(ctx).Result()
		require.NoError(t, err)
		require.Equal(t, "PONG", pong)
		slog.Info("Verified KV is accessible", "addr", rdb.Options().Addr)

		// Create the PAM resource
		kvPortFloat := float32(kvPort.Int())
		kvPamResResp, err := c.CreateRedisPamResourceWithResponse(
			ctx,
			client.CreateRedisPamResourceJSONRequestBody{
				ProjectId: uuid.MustParse(projectId),
				GatewayId: gatewayId,
				Name:      "kv-resource",
				ConnectionDetails: struct {
					Host                  string  `json:"host"`
					Port                  float32 `json:"port"`
					SslCertificate        *string `json:"sslCertificate,omitempty"`
					SslEnabled            bool    `json:"sslEnabled"`
					SslRejectUnauthorized bool    `json:"sslRejectUnauthorized"`
				}{
					Host:                  kvHost,
					Port:                  kvPortFloat,
					SslEnabled:            false,
					SslRejectUnauthorized: false,
				},
			})
		require.NoError(t, err)
		require.Equal(t, kvPamResResp.StatusCode(), http.StatusOK)
		slog.Info("PAM resource created successfully")
	})
}