LinkPort/internal/outbound/kubernetes_test.go

package outbound

import (
	"context"
	"encoding/json"
	"encoding/pem"
	"io"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/httpstream"
	streamspdy "k8s.io/apimachinery/pkg/util/httpstream/spdy"
	"k8s.io/apimachinery/pkg/util/intstr"
	appclient "k8s.io/client-go/kubernetes/typed/apps/v1"
	coreclient "k8s.io/client-go/kubernetes/typed/core/v1"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/portforward"
)

func TestKubernetesNativeForwardLifetime(t *testing.T) {
	closed := make(chan struct{})
	server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("Authorization") != "Bearer test-token" {
			http.Error(w, "unauthorized", 401)
			return
		}
		switch r.URL.Path {
		case "/api/v1/namespaces/testing/pods/db":
			w.Header().Set("Content-Type", "application/json")
			json.NewEncoder(w).Encode(&corev1.Pod{ObjectMeta: metav1.ObjectMeta{Name: "db"}, Status: corev1.PodStatus{Phase: corev1.PodRunning}})
		case "/api/v1/namespaces/testing/pods/db/portforward":
			if _, err := httpstream.Handshake(r, w, []string{portforward.PortForwardProtocolV1Name}); err != nil {
				return
			}
			conn := streamspdy.NewResponseUpgrader().UpgradeResponse(w, r, func(s httpstream.Stream, reply <-chan struct{}) error {
				go func() {
					<-reply
					defer s.Close()
					if s.Headers().Get(corev1.PortHeader) != "3306" {
						t.Errorf("wrong remote port: %v", s.Headers())
					}
					if s.Headers().Get(corev1.StreamType) == corev1.StreamTypeData {
						io.Copy(s, s)
					}
				}()
				return nil
			})
			if conn != nil {
				<-conn.CloseChan()
				close(closed)
			}
		default:
			http.NotFound(w, r)
		}
	}))
	defer server.Close()
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	ca := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: server.Certificate().Raw})
	conn, err := dialKubernetesConfig(ctx, &rest.Config{Host: server.URL, BearerToken: "test-token", TLSClientConfig: rest.TLSClientConfig{CAData: ca}}, "testing", "pod/db", 3306)
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()
	cancel()
	conn.SetDeadline(time.Now().Add(3 * time.Second))
	if _, err := conn.Write([]byte("native-forward")); err != nil {
		t.Fatal(err)
	}
	buf := make([]byte, len("native-forward"))
	if _, err := io.ReadFull(conn, buf); err != nil {
		t.Fatal(err)
	}
	if string(buf) != "native-forward" {
		t.Fatalf("unexpected response: %q", buf)
	}
	conn.Close()
	select {
	case <-closed:
	case <-time.After(3 * time.Second):
		t.Fatal("forward connection leaked after Close")
	}
}

func TestKubernetesCanceledSetupClosesRequest(t *testing.T) {
	entered, finished := make(chan struct{}), make(chan struct{})
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/api/v1/namespaces/testing/pods/db" {
			w.Header().Set("Content-Type", "application/json")
			json.NewEncoder(w).Encode(&corev1.Pod{ObjectMeta: metav1.ObjectMeta{Name: "db"}, Status: corev1.PodStatus{Phase: corev1.PodRunning}})
			return
		}
		close(entered)
		<-r.Context().Done()
		close(finished)
	}))
	defer server.Close()
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	result := make(chan error, 1)
	go func() {
		_, err := dialKubernetesConfig(ctx, &rest.Config{Host: server.URL}, "testing", "pod/db", 3306)
		result <- err
	}()
	select {
	case <-entered:
	case <-time.After(5 * time.Second):
		t.Fatal("upgrade not started")
	}
	cancel()
	select {
	case err := <-result:
		if err == nil {
			t.Fatal("expected cancellation")
		}
	case <-time.After(3 * time.Second):
		t.Fatal("canceled dial hung")
	}
	select {
	case <-finished:
	case <-time.After(3 * time.Second):
		t.Fatal("HTTP upgrade leaked")
	}
}

func TestKubernetesServiceNamedPort(t *testing.T) {
	pod := corev1.Pod{Spec: corev1.PodSpec{Containers: []corev1.Container{{Ports: []corev1.ContainerPort{{Name: "database", ContainerPort: 3306, Protocol: corev1.ProtocolTCP}}}}}}
	for _, tc := range []struct {
		target intstr.IntOrString
		want   int
		fail   bool
	}{
		{intstr.FromString("database"), 3306, false}, {intstr.FromInt32(5432), 5432, false}, {intstr.FromInt32(0), 3307, false}, {intstr.FromString("missing"), 0, true},
	} {
		n, err := serviceTargetPort(pod, corev1.ServicePort{Port: 3307, TargetPort: tc.target})
		if (err != nil) != tc.fail || n != tc.want {
			t.Fatalf("target %v: %d, %v", tc.target, n, err)
		}
	}
}

func TestKubernetesServiceResolution(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		switch r.URL.Path {
		case "/api/v1/namespaces/testing/services/db":
			json.NewEncoder(w).Encode(&corev1.Service{Spec: corev1.ServiceSpec{Selector: map[string]string{"app": "db"}, Ports: []corev1.ServicePort{{Port: 3307, TargetPort: intstr.FromString("sql")}}}})
		case "/api/v1/namespaces/testing/pods":
			if r.URL.Query().Get("labelSelector") != "app=db" {
				t.Errorf("wrong selector: %s", r.URL.RawQuery)
			}
			json.NewEncoder(w).Encode(&corev1.PodList{Items: []corev1.Pod{
				{ObjectMeta: metav1.ObjectMeta{Name: "pending"}, Status: corev1.PodStatus{Phase: corev1.PodPending}},
				{ObjectMeta: metav1.ObjectMeta{Name: "running"}, Status: corev1.PodStatus{Phase: corev1.PodRunning}, Spec: corev1.PodSpec{Containers: []corev1.Container{{Ports: []corev1.ContainerPort{{Name: "sql", ContainerPort: 3306}}}}}},
			}})
		default:
			http.NotFound(w, r)
		}
	}))
	defer server.Close()
	config := &rest.Config{Host: server.URL}
	core, err := coreclient.NewForConfig(config)
	if err != nil {
		t.Fatal(err)
	}
	apps, err := appclient.NewForConfig(config)
	if err != nil {
		t.Fatal(err)
	}
	pod, port, err := resolveKubernetesTarget(context.Background(), core, apps, "testing", "service/db", 3307)
	if err != nil || pod != "running" || port != 3306 {
		t.Fatalf("resolution: %s:%d %v", pod, port, err)
	}
	if _, _, err := resolveKubernetesTarget(context.Background(), core, apps, "testing", "service/db", 9999); err == nil {
		t.Fatal("accepted unknown service port")
	}
}