alibaba/kt-connect

package connect

import (
    "fmt"
    "github.com/alibaba/kt-connect/pkg/common"
    opt "github.com/alibaba/kt-connect/pkg/kt/command/options"
    "github.com/alibaba/kt-connect/pkg/kt/service/cluster"
    "github.com/alibaba/kt-connect/pkg/kt/service/dns"
    "github.com/alibaba/kt-connect/pkg/kt/transmission"
    "github.com/alibaba/kt-connect/pkg/kt/util"
    "github.com/rs/zerolog/log"
    coreV1 "k8s.io/api/core/v1"
    "strings"
    "time"
)

func setupDns(shadowPodName, shadowPodIp string) error {
    if strings.HasPrefix(opt.Get().Connect.DnsMode, util.DnsModeHosts) {
        log.Info().Msgf("Setting up dns in hosts mode")
        dump2HostsNamespaces := ""
        pos := len(util.DnsModeHosts)
        if len(opt.Get().Connect.DnsMode) > pos + 1 && opt.Get().Connect.DnsMode[pos:pos+1] == ":" {
            dump2HostsNamespaces = opt.Get().Connect.DnsMode[pos+1:]
        }
        if err := dumpToHost(dump2HostsNamespaces); err != nil {
            return err
        }
    } else if opt.Get().Connect.DnsMode == util.DnsModePodDns {
        log.Info().Msgf("Setting up dns in pod mode")
        return dns.SetNameServer(shadowPodIp)
    } else if strings.HasPrefix(opt.Get().Connect.DnsMode, util.DnsModeLocalDns) {
        log.Info().Msgf("Setting up dns in local mode")
        svcToIp, headlessPods := getServiceHosts(opt.Get().Global.Namespace, true)
        if err := dns.DumpHosts(svcToIp, ""); err != nil {
            return err
        }
        watchServicesAndPods(opt.Get().Global.Namespace, svcToIp, headlessPods, true)

        forwardedPodPort := util.GetRandomTcpPort()
        if _, err := transmission.SetupPortForwardToLocal(shadowPodName, common.StandardDnsPort, forwardedPodPort); err != nil {
            return err
        }

        dnsPort := util.AlternativeDnsPort
        if util.IsWindows() {
            dnsPort = common.StandardDnsPort
        } else if util.IsMacos() {
            dnsPort = opt.Get().Connect.DnsPort
        }
        // must set up name server before change dns config
        // otherwise the upstream name server address will be incorrect in linux
        if err := dns.SetupLocalDns(forwardedPodPort, dnsPort, getDnsOrder(opt.Get().Connect.DnsMode)); err != nil {
            log.Error().Err(err).Msgf("Failed to setup local dns server")
            return err
        }
        return dns.SetNameServer(fmt.Sprintf("%s:%d", common.Localhost, dnsPort))
    } else {
        return fmt.Errorf("invalid dns mode: '%s', supportted mode are %s, %s, %s", opt.Get().Connect.DnsMode,
            util.DnsModeLocalDns, util.DnsModePodDns, util.DnsModeHosts)
    }
    return nil
}

func getDnsOrder(dnsMode string) []string {
    if ! strings.Contains(dnsMode, ":") {
        return []string{ util.DnsOrderCluster, util.DnsOrderUpstream }
    }
    return strings.Split(strings.SplitN(dnsMode, ":", 2)[1], ",")
}

func watchServicesAndPods(namespace string, svcToIp map[string]string, headlessPods []string, shortDomainOnly bool) {
    setupTime := time.Now().Unix()
    go cluster.Ins().WatchService("", namespace,
        func(svc *coreV1.Service) {
            // ignore add service event during watch setup
            if time.Now().Unix() - setupTime > 3 {
                svcToIp, headlessPods = getServiceHosts(namespace, shortDomainOnly)
                _ = dns.DumpHosts(svcToIp, namespace)
            }
        },
        func(svc *coreV1.Service) {
            svcToIp, headlessPods = getServiceHosts(namespace, shortDomainOnly)
            _ = dns.DumpHosts(svcToIp, namespace)
        }, nil)
    go cluster.Ins().WatchPod("", namespace, nil, func(pod *coreV1.Pod) {
        if util.Contains(headlessPods, pod.Name) {
            // it may take some time for new pod get assign an ip
            time.Sleep(5 * time.Second)
            svcToIp, headlessPods = getServiceHosts(namespace, shortDomainOnly)
            _ = dns.DumpHosts(svcToIp, namespace)
        }
    }, nil)
}

func dumpToHost(targetNamespaces string) error {
    namespacesToDump := []string{opt.Get().Global.Namespace}
    if targetNamespaces != "" {
        namespacesToDump = []string{}
        for _, ns := range strings.Split(targetNamespaces, ",") {
            namespacesToDump = append(namespacesToDump, ns)
        }
    }
    hosts := map[string]string{}
    for _, namespace := range namespacesToDump {
        log.Debug().Msgf("Search service in %s namespace ...", namespace)
        svcToIp, headlessPods := getServiceHosts(namespace, false)
        watchServicesAndPods(namespace, svcToIp, headlessPods, false)
        for svc, ip := range svcToIp {
            hosts[svc] = ip
        }
    }
    return dns.DumpHosts(hosts, "")
}

func getServiceHosts(namespace string, shortDomainOnly bool) (map[string]string, []string) {
    hosts := make(map[string]string)
    podNames := make([]string, 0)
    services, err := cluster.Ins().GetAllServiceInNamespace(namespace)
    if err == nil {
        for _, service := range services.Items {
            ip := service.Spec.ClusterIP
            if ip == "" || ip == "None" {
                pods, err2 := cluster.Ins().GetPodsByLabel(service.Spec.Selector, namespace)
                if err2 != nil || len(pods.Items) == 0 {
                    continue
                }
                for _, p := range pods.Items {
                    ip = p.Status.PodIP
                    if ip != "" {
                        podNames = append(podNames, p.Name)
                        break
                    }
                }
                log.Debug().Msgf("Headless service found: %s.%s %s", service.Name, namespace, ip)
            } else {
                log.Debug().Msgf("Service found: %s.%s %s", service.Name, namespace, ip)
            }
            if shortDomainOnly {
                hosts[service.Name] = ip
            } else {
                if namespace == opt.Get().Global.Namespace {
                    hosts[service.Name] = ip
                }
                hosts[fmt.Sprintf("%s.%s", service.Name, namespace)] = ip
                hosts[fmt.Sprintf("%s.%s.svc.%s", service.Name, namespace, opt.Get().Connect.ClusterDomain)] = ip
            }
        }
    }
    return hosts, podNames
}

func getOrCreateShadow() (string, string, string, error) {
    shadowPodName := fmt.Sprintf("kt-connect-shadow-%s", strings.ToLower(util.RandomString(5)))
    if opt.Get().Connect.ShareShadow {
        shadowPodName = fmt.Sprintf("kt-connect-shadow-daemon")
    }

    endPointIP, podName, privateKeyPath, err := cluster.Ins().GetOrCreateShadow(shadowPodName, getLabels(),
        make(map[string]string), getEnvs(), "", map[int]string{})
    if err != nil {
        return "", "", "", err
    }

    return endPointIP, podName, privateKeyPath, nil
}

func getEnvs() map[string]string {
    envs := make(map[string]string)
    localDomains := dns.GetLocalDomains()
    if localDomains != "" {
        log.Debug().Msgf("Found local domains: %s", localDomains)
        envs[common.EnvVarLocalDomains] = localDomains
    }
    if strings.HasPrefix(opt.Get().Connect.DnsMode, util.DnsModeLocalDns) {
        envs[common.EnvVarDnsProtocol] = "tcp"
    } else {
        envs[common.EnvVarDnsProtocol] = "udp"
    }
    if opt.Get().Global.Debug {
        envs[common.EnvVarLogLevel] = "debug"
    } else {
        envs[common.EnvVarLogLevel] = "info"
    }
    return envs
}

func getLabels() map[string]string {
    labels := map[string]string{
        util.KtRole:    util.RoleConnectShadow,
    }
    if opt.Get().Global.UseShadowDeployment {
        labels[util.KtTarget] = util.RandomString(20)
    }
    return labels
}