alibaba/kt-connect

package dns

import (
    "bufio"
    "bytes"
    "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/util"
    "github.com/rs/zerolog/log"
    "io/ioutil"
    "os"
    "os/exec"
    "os/signal"
    "strconv"
    "strings"
    "syscall"
)

const (
    commentKtAdded   = " # Added by KtConnect"
    commentKtRemoved = " # Removed by KtConnect"
)

// SetNameServer set dns server records
func SetNameServer(dnsServer string) error {
    dnsSignal := make(chan error)
    go func() {
        defer func() {
            restoreResolvConf()
            if strings.HasPrefix(opt.Get().Connect.DnsMode, util.DnsModeLocalDns) {
                restoreIptables()
            }
        }()
        if strings.HasPrefix(opt.Get().Connect.DnsMode, util.DnsModeLocalDns) {
            if err := setupIptables(); err != nil {
                dnsSignal <-err
                return
            }
        }
        dnsSignal <- setupResolvConf(dnsServer)

        sigCh := make(chan os.Signal, 1)
        signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM)
        <-sigCh
    }()
    return <-dnsSignal
}

// HandleExtraDomainMapping handle extra domain change
func HandleExtraDomainMapping(extraDomains map[string]string, localDnsPort int) {
    // pass
}

// RestoreNameServer remove the nameservers added by ktctl
func RestoreNameServer() {
    restoreResolvConf()
    restoreIptables()
}

func setupResolvConf(dnsServer string) error {
    f, err := os.Open(util.ResolvConf)
    if err != nil {
        return err
    }
    defer f.Close()

    var buf bytes.Buffer

    sample := fmt.Sprintf("%s %s ", util.FieldNameserver, strings.Split(dnsServer, ":")[0])
    scanner := bufio.NewScanner(f)
    for scanner.Scan() {
        line := scanner.Text()
        if strings.HasPrefix(line, sample) {
            // required dns server already been added
            return nil
        } else if strings.HasPrefix(line, util.FieldNameserver) {
            buf.WriteString("#")
            buf.WriteString(line)
            buf.WriteString(commentKtRemoved)
            buf.WriteString("\n")
        } else {
            buf.WriteString(line)
            buf.WriteString("\n")
        }
    }

    // Add nameserver and comment to resolv.conf
    nameserverIp := strings.Split(dnsServer, ":")[0]
    buf.WriteString(fmt.Sprintf("%s %s%s\n", util.FieldNameserver, nameserverIp, commentKtAdded))

    stat, _ := f.Stat()
    return ioutil.WriteFile(util.ResolvConf, buf.Bytes(), stat.Mode())
}

func setupIptables() error {
    // run command: iptables --table nat --insert OUTPUT --proto udp --dest 127.0.0.1/32 --dport 53 --jump REDIRECT --to-ports 10053
    if _, _, err := util.RunAndWait(exec.Command("iptables",
        "--table",
        "nat",
        "--insert",
        "OUTPUT",
        "--proto",
        "udp",
        "--dest",
        fmt.Sprintf("%s/32", common.Localhost),
        "--dport",
        strconv.Itoa(common.StandardDnsPort),
        "--jump",
        "REDIRECT",
        "--to-ports",
        strconv.Itoa(util.AlternativeDnsPort),
    )); err != nil {
        log.Error().Msgf("Failed to use local dns server")
        return err
    }
    return nil
}

func restoreResolvConf() {
    f, err := os.Open(util.ResolvConf)
    if err != nil {
        log.Error().Err(err).Msgf("Failed to open resolve.conf during restoring")
        return
    }
    defer f.Close()

    var buf bytes.Buffer

    scanner := bufio.NewScanner(f)
    for scanner.Scan() {
        line := scanner.Text()
        if strings.HasSuffix(line, commentKtRemoved) {
            line = strings.TrimSuffix(line, commentKtRemoved)
            line = strings.TrimPrefix(line, "#")
            buf.WriteString(line)
            buf.WriteString("\n")
        } else if strings.HasSuffix(line, commentKtAdded) {
            log.Debug().Msgf("remove line: %s ", line)
        } else {
            buf.WriteString(line)
            buf.WriteString("\n")
        }
    }

    stat, _ := f.Stat()
    if err = ioutil.WriteFile(util.ResolvConf, buf.Bytes(), stat.Mode()); err != nil {
        log.Error().Err(err).Msgf("Failed to write resolve.conf during restoring")
    }
}

func restoreIptables() {
    for {
        // run command: iptables --table nat --delete OUTPUT --proto udp --dest 127.0.0.1/32 --dport 53 --jump REDIRECT --to-ports 10053
        _, _, err := util.RunAndWait(exec.Command("iptables",
            "--table",
            "nat",
            "--delete",
            "OUTPUT",
            "--proto",
            "udp",
            "--dest",
            fmt.Sprintf("%s/32", common.Localhost),
            "--dport",
            strconv.Itoa(common.StandardDnsPort),
            "--jump",
            "REDIRECT",
            "--to-ports",
            strconv.Itoa(util.AlternativeDnsPort),
        ))
        if err != nil {
            // no more rule left
            break
        }
    }
}