LinkPort/internal/outbound/kubernetes.go

package outbound

import (
	"agentdb/internal/model"
	"context"
	"crypto/tls"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	streamspdy "k8s.io/apimachinery/pkg/util/httpstream/spdy"
	"k8s.io/apimachinery/pkg/util/intstr"
	utilnet "k8s.io/apimachinery/pkg/util/net"
	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/clientcmd"
	"k8s.io/client-go/tools/portforward"
	"k8s.io/client-go/transport/spdy"
)

type kubernetesConn struct {
	net.Conn
	stop func()
	done <-chan struct{}
	once sync.Once
}

func (c *kubernetesConn) Close() error {
	err := c.Conn.Close()
	c.once.Do(func() { c.stop(); <-c.done })
	return err
}

// The tunnel owns its lifetime. A caller's dial deadline only bounds setup.
func dialKubernetes(ctx context.Context, network, address string, o model.Outbound) (net.Conn, error) {
	if network != "tcp" {
		return nil, errors.New("Kubernetes port-forward requires TCP")
	}
	if err := o.ValidateKubernetes(); err != nil {
		return nil, err
	}
	_, rawPort, err := net.SplitHostPort(address)
	port, parseErr := strconv.Atoi(rawPort)
	if err != nil || parseErr != nil || port < 1 || port > 65535 {
		return nil, errors.New("Kubernetes port-forward invalid target port")
	}
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	rules := clientcmd.NewDefaultClientConfigLoadingRules()
	rules.ExplicitPath = o.KubeConfig
	overrides := &clientcmd.ConfigOverrides{CurrentContext: o.KubeContext}
	loader := clientcmd.NewNonInteractiveDeferredLoadingClientConfig(rules, overrides)
	config, err := loader.ClientConfig()
	if err != nil {
		return nil, errors.New("Kubernetes config could not be loaded")
	}
	namespace := o.Namespace
	if namespace == "" {
		namespace, _, err = loader.Namespace()
		if err != nil {
			return nil, errors.New("Kubernetes namespace could not be loaded")
		}
	}
	return dialKubernetesConfig(ctx, config, namespace, o.Resource, port)
}

func dialKubernetesConfig(ctx context.Context, config *rest.Config, namespace, resource string, port int) (net.Conn, error) {
	apiURL, err := url.Parse(config.Host)
	if err != nil {
		return nil, errors.New("Kubernetes API address is invalid")
	}
	proxy := config.Proxy
	if proxy == nil {
		proxy = http.ProxyFromEnvironment
	}
	proxyURL, err := proxy(&http.Request{URL: apiURL})
	if err != nil || proxyURL != nil {
		return nil, errors.New("Kubernetes port-forward requires direct API access; kubeconfig proxy-url and HTTP proxy are not supported")
	}
	core, err := coreclient.NewForConfig(config)
	if err != nil {
		return nil, errors.New("Kubernetes client configuration failed")
	}
	apps, err := appclient.NewForConfig(config)
	if err != nil {
		return nil, errors.New("Kubernetes client configuration failed")
	}
	pod, remotePort, err := resolveKubernetesTarget(ctx, core, apps, namespace, resource, port)
	if err != nil {
		return nil, err
	}
	// Do not apply a request-wide timeout to the upgraded streaming connection.
	streamConfig := rest.CopyConfig(config)
	streamConfig.Timeout = 0
	lifetime, cancel := context.WithCancel(context.Background())
	baseDial := streamConfig.Dial
	if baseDial == nil {
		baseDial = (&net.Dialer{}).DialContext
	}
	streamConfig.Dial = func(ctx context.Context, network, address string) (net.Conn, error) {
		conn, err := baseDial(ctx, network, address)
		if err != nil {
			return nil, err
		}
		return &kubernetesSocket{Conn: conn, stop: context.AfterFunc(lifetime, func() { conn.Close() })}, nil
	}
	upgradeTransport, err := rest.TransportFor(streamConfig)
	if err != nil {
		cancel()
		return nil, errors.New("Kubernetes port-forward transport configuration failed")
	}
	tlsConfig, err := utilnet.TLSClientConfig(upgradeTransport)
	if err != nil {
		cancel()
		return nil, errors.New("Kubernetes port-forward TLS configuration failed")
	}
	if tlsConfig == nil {
		tlsConfig = &tls.Config{MinVersion: tls.VersionTLS12}
	}
	// DialURL treats a nil TLS configuration as insecure when using custom dialers.
	// Supply an explicit verified default while retaining client certificate callbacks.
	upgradeTransport = &http.Transport{TLSClientConfig: tlsConfig, DialContext: streamConfig.Dial}
	upgrader, err := streamspdy.NewRoundTripperWithConfig(streamspdy.RoundTripperConfig{UpgradeTransport: upgradeTransport, PingPeriod: 5 * time.Second})
	if err != nil {
		cancel()
		return nil, errors.New("Kubernetes port-forward transport configuration failed")
	}
	transport, err := rest.HTTPWrappersForConfig(streamConfig, upgrader)
	if err != nil {
		cancel()
		return nil, errors.New("Kubernetes port-forward authentication configuration failed")
	}
	stopCh, ready := make(chan struct{}), make(chan struct{})
	var once sync.Once
	stop := func() { once.Do(func() { cancel(); close(stopCh) }) }
	url := core.RESTClient().Post().Resource("pods").Namespace(namespace).Name(pod).SubResource("portforward").URL()
	dialer := spdy.NewDialer(upgrader, &http.Client{Transport: contextTransport{transport, lifetime}}, http.MethodPost, url)
	forwarder, err := portforward.NewOnAddresses(dialer, []string{"127.0.0.1"}, []string{"0:" + strconv.Itoa(remotePort)}, stopCh, ready, io.Discard, io.Discard)
	if err != nil {
		stop()
		return nil, errors.New("Kubernetes port-forward initialization failed")
	}
	done := make(chan struct{})
	go func() { defer close(done); _ = forwarder.ForwardPorts() }()
	owned := true
	defer func() {
		if owned {
			stop()
			<-done
		}
	}()
	select {
	case <-ctx.Done():
		return nil, fmt.Errorf("Kubernetes port-forward setup: %w", ctx.Err())
	case <-done:
		return nil, errors.New("Kubernetes port-forward startup failed; check API access and pods/portforward permission")
	case <-ready:
	}
	ports, err := forwarder.GetPorts()
	if err != nil || len(ports) != 1 {
		return nil, errors.New("Kubernetes port-forward listener unavailable")
	}
	conn, err := (&net.Dialer{}).DialContext(ctx, "tcp", net.JoinHostPort("127.0.0.1", strconv.Itoa(int(ports[0].Local))))
	if err != nil {
		return nil, fmt.Errorf("Kubernetes port-forward local connection: %w", err)
	}
	if err := ctx.Err(); err != nil {
		conn.Close()
		return nil, err
	}
	select {
	case <-done:
		conn.Close()
		return nil, errors.New("Kubernetes port-forward stopped")
	default:
	}
	owned = false
	return &kubernetesConn{Conn: conn, stop: stop, done: done}, nil
}

type contextTransport struct {
	base http.RoundTripper
	ctx  context.Context
}

// SPDY reads upgrade responses directly, so request cancellation must also
// close the underlying socket when an API server never sends its response.
type kubernetesSocket struct {
	net.Conn
	stop func() bool
}

func (c *kubernetesSocket) Close() error { c.stop(); return c.Conn.Close() }

func (t contextTransport) RoundTrip(r *http.Request) (*http.Response, error) {
	return t.base.RoundTrip(r.Clone(t.ctx))
}

func resolveKubernetesTarget(ctx context.Context, core coreclient.CoreV1Interface, apps appclient.AppsV1Interface, namespace, resource string, port int) (string, int, error) {
	kind, name, _ := strings.Cut(resource, "/")
	var selector string
	var servicePort *corev1.ServicePort
	switch kind {
	case "pod":
		pod, err := core.Pods(namespace).Get(ctx, name, metav1.GetOptions{})
		if err != nil {
			return "", 0, errors.New("Kubernetes Pod lookup failed; check name and read permission")
		}
		if pod.Status.Phase != corev1.PodRunning || pod.DeletionTimestamp != nil {
			return "", 0, errors.New("Kubernetes Pod is not running")
		}
		return pod.Name, port, nil
	case "service":
		svc, err := core.Services(namespace).Get(ctx, name, metav1.GetOptions{})
		if err != nil {
			return "", 0, errors.New("Kubernetes Service lookup failed; check name and read permission")
		}
		if len(svc.Spec.Selector) == 0 {
			return "", 0, errors.New("Kubernetes Service requires a Pod selector")
		}
		selector = labels.SelectorFromSet(svc.Spec.Selector).String()
		for i := range svc.Spec.Ports {
			p := &svc.Spec.Ports[i]
			if int(p.Port) == port && (p.Protocol == "" || p.Protocol == corev1.ProtocolTCP) {
				servicePort = p
				break
			}
		}
		if servicePort == nil {
			return "", 0, errors.New("Kubernetes Service has no matching TCP port")
		}
	case "deployment":
		dep, err := apps.Deployments(namespace).Get(ctx, name, metav1.GetOptions{})
		if err != nil {
			return "", 0, errors.New("Kubernetes Deployment lookup failed; check name and read permission")
		}
		if dep.Spec.Selector == nil {
			return "", 0, errors.New("Kubernetes Deployment has no Pod selector")
		}
		s, err := metav1.LabelSelectorAsSelector(dep.Spec.Selector)
		if err != nil || s.Empty() {
			return "", 0, errors.New("Kubernetes Deployment has invalid Pod selector")
		}
		selector = s.String()
	default:
		return "", 0, errors.New("unsupported Kubernetes resource")
	}
	pods, err := core.Pods(namespace).List(ctx, metav1.ListOptions{LabelSelector: selector})
	if err != nil {
		return "", 0, errors.New("Kubernetes Pod list failed; check read permission")
	}
	sort.SliceStable(pods.Items, func(i, j int) bool { return podReady(pods.Items[i]) && !podReady(pods.Items[j]) })
	for _, pod := range pods.Items {
		if pod.Status.Phase != corev1.PodRunning || pod.DeletionTimestamp != nil {
			continue
		}
		if servicePort == nil {
			return pod.Name, port, nil
		}
		n, err := serviceTargetPort(pod, *servicePort)
		if err == nil {
			return pod.Name, n, nil
		}
	}
	return "", 0, errors.New("Kubernetes has no running Pod with the requested TCP target port")
}

func podReady(p corev1.Pod) bool {
	for _, c := range p.Status.Conditions {
		if c.Type == corev1.PodReady && c.Status == corev1.ConditionTrue {
			return true
		}
	}
	return false
}

func serviceTargetPort(pod corev1.Pod, port corev1.ServicePort) (int, error) {
	if port.TargetPort.Type == intstr.Int {
		p := int(port.TargetPort.IntVal)
		if p == 0 {
			p = int(port.Port)
		}
		if p > 0 && p <= 65535 {
			return p, nil
		}
	} else {
		for _, container := range pod.Spec.Containers {
			for _, p := range container.Ports {
				if p.Name == port.TargetPort.StrVal && (p.Protocol == "" || p.Protocol == corev1.ProtocolTCP) && p.ContainerPort > 0 && p.ContainerPort <= 65535 {
					return int(p.ContainerPort), nil
				}
			}
		}
	}
	return 0, errors.New("Kubernetes Service targetPort was not found on Pod")
}