alibaba/kt-connect

package dns

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/util"
    "github.com/miekg/dns"
    "github.com/rs/zerolog/log"
    "net"
    "regexp"
    "strconv"
    "strings"
    "time"
)

type DnsServer struct {
    dnsAddresses []string
    extraDomains map[string]string
}

func SetupLocalDns(remoteDnsPort, localDnsPort int, dnsOrder []string) error {
    var res = make(chan error)
    go func() {
        upstreamDnsAddresses := getDnsAddresses(dnsOrder, GetNameServer(), remoteDnsPort)
        // domain-name -> ip
        extraDomains := getIngressDomains()
        log.Info().Msgf("Setup local DNS with upstream %v", upstreamDnsAddresses)
        HandleExtraDomainMapping(extraDomains, localDnsPort)
        res <-common.SetupDnsServer(&DnsServer{upstreamDnsAddresses, extraDomains}, localDnsPort, "udp")
    }()
    select {
    case err := <-res:
        return err
    case <-time.After(1 * time.Second):
        return nil
    }
}

func getIngressDomains() map[string]string {
    if opt.Get().Connect.IngressIp == "" {
        return map[string]string{}
    }
    if !util.IsValidIp(opt.Get().Connect.IngressIp) {
        log.Warn().Msgf("Ingress Ip '" + opt.Get().Connect.IngressIp + "' is invalid")
        return map[string]string{}
    }
    ingressDomains := make(map[string]string)
    if ingresses, err := cluster.Ins().GetAllIngressInNamespace(opt.Get().Global.Namespace); err != nil {
        log.Warn().Err(err).Msgf("Failed to found ingress instances")
    } else {
        for _, ingress := range ingresses.Items {
            for _, rule := range ingress.Spec.Rules {
                if rule.Host != "" {
                    log.Debug().Msgf("Find ingress domain " + rule.Host)
                    ingressDomains[rule.Host] = opt.Get().Connect.IngressIp
                }
            }
        }
    }
    return ingressDomains
}

func getDnsAddresses(dnsOrder []string, upstreamDns string, clusterDnsPort int) []string {
    upstreamPattern := fmt.Sprintf("^([cdptu]{3}:)?%s(:[0-9]+)?$", util.DnsOrderUpstream)
    var dnsAddresses []string
    for _, dnsAddr := range dnsOrder {
        if dnsAddr == util.DnsOrderCluster {
            dnsAddresses = append(dnsAddresses, fmt.Sprintf("tcp:%s:%d", common.Localhost, clusterDnsPort))
        } else if ok, err := regexp.MatchString(upstreamPattern, dnsAddr); err == nil && ok {
            upstreamParts := strings.Split(dnsAddr, ":")
            if upstreamDns != "" {
                switch strings.Count(dnsAddr, ":") {
                case 0:
                    dnsAddresses = append(dnsAddresses, fmt.Sprintf("udp:%s:%d", upstreamDns, common.StandardDnsPort))
                case 1:
                    if _, err = strconv.Atoi(upstreamParts[1]); err == nil {
                        dnsAddresses = append(dnsAddresses, fmt.Sprintf("udp:%s:%s", upstreamDns, upstreamParts[1]))
                    } else {
                        dnsAddresses = append(dnsAddresses, fmt.Sprintf("%s:%s:%d", upstreamParts[0], upstreamDns, common.StandardDnsPort))
                    }
                case 2:
                    dnsAddresses = append(dnsAddresses, fmt.Sprintf("%s:%s:%s", upstreamParts[0], upstreamDns, upstreamParts[2]))
                default:
                    log.Warn().Msgf("Skip invalid upstream dns server %s", dnsAddr)
                }
            }
        } else {
            switch strings.Count(dnsAddr, ":") {
            case 0:
                dnsAddresses = append(dnsAddresses, fmt.Sprintf("udp:%s:%d", dnsAddr, common.StandardDnsPort))
            case 1:
                if _, err = strconv.Atoi(strings.Split(dnsAddr, ":")[1]); err == nil {
                    dnsAddresses = append(dnsAddresses, fmt.Sprintf("udp:%s", dnsAddr))
                } else {
                    dnsAddresses = append(dnsAddresses, fmt.Sprintf("%s:%d", dnsAddr, common.StandardDnsPort))
                }
            case 2:
                dnsAddresses = append(dnsAddresses, dnsAddr)
            default:
                log.Warn().Msgf("Skip invalid dns server %s", dnsAddr)
            }
        }
    }
    return dnsAddresses
}

// ServeDNS query DNS record
func (s *DnsServer) ServeDNS(w dns.ResponseWriter, req *dns.Msg) {
    msg := (&dns.Msg{}).SetReply(req)
    msg.Authoritative = true
    msg.Answer = query(req, s.dnsAddresses, s.extraDomains)
    if err := w.WriteMsg(msg); err != nil {
        log.Warn().Err(err).Msgf("Failed to reply dns request")
    }
}

func query(req *dns.Msg, dnsAddresses []string, extraDomains map[string]string) []dns.RR {
    domain := req.Question[0].Name
    qtype := req.Question[0].Qtype

    answer := common.ReadCache(domain, qtype, int64(opt.Get().Connect.DnsCacheTtl))
    if answer != nil {
        log.Debug().Msgf("Found domain %s (%d) in cache", domain, qtype)
        return answer
    }

    for host, ip := range extraDomains {
        if wildcardMatch(host, domain) {
            return []dns.RR{toARecord(domain, ip)}
        }
    }

    for _, dnsAddr := range dnsAddresses {
        dnsParts := strings.SplitN(dnsAddr, ":", 3)
        protocol := dnsParts[0]
        ip := dnsParts[1]
        port, err := strconv.Atoi(dnsParts[2])
        if ip == "" || err != nil || (protocol != "tcp" && protocol != "udp") {
            // skip invalid dns address
            continue
        }
        res, err := common.NsLookup(domain, qtype, protocol, fmt.Sprintf("%s:%d", ip, port))
        if res != nil && len(res.Answer) > 0 {
            // only record none-empty result of cluster dns
            log.Debug().Msgf("Found domain %s (%d) in dns (%s:%d)", domain, qtype, ip, port)
            common.WriteCache(domain, qtype, res.Answer, time.Now().Unix())
            return res.Answer
        } else if err != nil && !common.IsDomainNotExist(err) {
            // usually io timeout error
            log.Warn().Err(err).Msgf("Failed to lookup %s (%d) in dns (%s:%d)", domain, qtype, ip, port)
        }
    }
    log.Debug().Msgf("Empty answer for domain lookup %s (%d)", domain, qtype)
    common.WriteCache(domain, qtype, []dns.RR{}, time.Now().Unix()-int64(opt.Get().Connect.DnsCacheTtl)/2)
    return []dns.RR{}
}

func wildcardMatch(pattenDomain, targetDomain string) bool {
    if !strings.HasSuffix(pattenDomain, ".") {
        pattenDomain = pattenDomain + "."
    }
    if strings.Contains(pattenDomain, "*") {
        ok, err := regexp.MatchString("^" + strings.ReplaceAll(strings.ReplaceAll(pattenDomain, ".", "\\."), "*", ".*") + "$", targetDomain)
        return ok && err == nil
    } else {
        return pattenDomain == targetDomain
    }
}

func toARecord(domain, ip string) dns.RR {
    return &dns.A {
        Hdr: dns.RR_Header {
            Name: domain,
            Rrtype: dns.TypeA,
            Class: dns.ClassINET,
            Ttl: 5,
            Rdlength: 4,
        },
        A: net.ParseIP(ip),
    }
}