godns/main.go

package main

import (
	"flag"
	"github.com/go-redis/redis/v7"
	log "github.com/sirupsen/logrus"
	"github.com/spf13/viper"
	"io"
	"meow.tf/joker/godns/api"
	"meow.tf/joker/godns/cache"
	"meow.tf/joker/godns/hosts"
	"meow.tf/joker/godns/resolver"
	"meow.tf/joker/godns/settings"
	"os"
	"os/signal"
	"runtime"
	"runtime/pprof"
	"time"
)

const (
	Version = "0.3.0"
)

var (
	cfgFile string
)

func init() {
	flag.StringVar(&cfgFile, "config", "/etc/godns/godns.conf", "")
}

func main() {
	initLogger()

	viper.SetConfigFile(cfgFile)

	viper.AutomaticEnv()

	if err := viper.ReadInConfig(); err == nil {
		log.WithField("file", viper.ConfigFileUsed()).Info("Using configuration from file")
	}

	server := &Server{
		host:     viper.GetString("server.host"),
		networks: viper.GetStringSlice("server.nets"),
		rTimeout: 5 * time.Second,
		wTimeout: 5 * time.Second,
	}

	var resolverSettings resolver.Settings

	/*
	resolver.Settings{
			Timeout:        viper.GetInt("resolv.timeout"),
			Interval:       viper.GetInt("resolv.interval"),
			SetEDNS0:       viper.GetBool("resolv.edns0"),
			ServerListFile: viper.GetStringSlice("resolv.server-list"),
			ResolvFile:     viper.GetString("resolv.file"),
		}
	 */

	viper.UnmarshalKey("resolv", &resolverSettings)

	var resolverCache, negCache cache.Cache
	r := resolver.NewResolver(resolverSettings)

	cacheDuration := viper.GetDuration("cache.expire")

	backend := viper.GetString("cache.backend")

	var redisConfig settings.RedisSettings

	viper.UnmarshalKey("cache.redis", &redisConfig)

	switch backend {
	case "memory":
		cacheMaxCount := viper.GetInt("cache.memory.maxCount")

		negCache = cache.NewMemoryCache(cacheDuration/2, cacheMaxCount)
		resolverCache = cache.NewMemoryCache(cacheDuration, cacheMaxCount)
	case "memcached":
		servers := viper.GetStringSlice("cache.memcached.servers")

		resolverCache = cache.NewMemcachedCache(servers, int32(cacheDuration.Seconds()))
		negCache = cache.NewMemcachedCache(servers, int32(cacheDuration.Seconds()/2))
	case "redis":

		resolverCache = cache.NewRedisCache(redisConfig, cacheDuration)
		negCache = cache.NewRedisCache(redisConfig, cacheDuration/2)
	default:
		log.WithField("backend", backend).Fatalln("Invalid cache backend")
	}

	providers := make([]hosts.Provider, 0)

	if viper.GetBool("hosts.file.enable") {
		providers = append(providers, hosts.NewFileProvider(viper.GetString("hosts.file.file"), viper.GetDuration("hosts.file.ttl")))
	}

	if viper.GetBool("hosts.bolt.enable") {
		providers = append(providers, hosts.NewBoltProvider(viper.GetString("hosts.bolt.file")))
	}

	if viper.GetBool("hosts.redis.enable") {
		rc := redis.NewClient(&redis.Options{Addr: redisConfig.Addr(), DB: redisConfig.DB, Password: redisConfig.Password})

		providers = append(providers, hosts.NewRedisProvider(rc, viper.GetString("hosts.redis.key")))
	}

	h := hosts.NewHosts(providers)

	a := api.New()

	hosts.EnableAPI(h, a.Router())

	if viper.GetBool("api.enabled") {
		go func() {
			err := a.Start()

			if err != nil {
				log.WithError(err).Fatalln("Unable to bind API")
			}
		}()
	}

	handler := NewHandler(r, resolverCache, negCache, h)

	server.Run(handler)

	log.Infof("joker dns %s (%s)", Version, runtime.Version())

	if viper.GetBool("debug") {
		go profileCPU()
		go profileMEM()
	}

	sig := make(chan os.Signal)
	signal.Notify(sig, os.Interrupt)
	<- sig

	log.Info("signal received, stopping")
}

func profileCPU() {
	f, err := os.Create("godns.cprof")
	if err != nil {
		log.WithError(err).Error("Unable to profile cpu due to error")
		return
	}

	pprof.StartCPUProfile(f)

	time.AfterFunc(6*time.Minute, func() {
		pprof.StopCPUProfile()
		f.Close()
	})
}

func profileMEM() {
	f, err := os.Create("godns.mprof")

	if err != nil {
		log.WithError(err).Error("Unable to profile memory due to error")
		return
	}

	time.AfterFunc(5*time.Minute, func() {
		pprof.WriteHeapProfile(f)
		f.Close()
	})

}

func initLogger() {
	if viper.GetBool("log.stdout") {
		// log.SetLogger("console", nil)
	}

	if file := viper.GetString("log.file"); file != "" {
		f, err := os.Create(file)

		if err != nil {
			return
		}

		log.SetOutput(io.MultiWriter(f, log.StandardLogger().Out))
	}

	level, err := log.ParseLevel(viper.GetString("log.level"))

	if err != nil {
		return
	}

	log.SetLevel(level)
}