Added dns search code.

config.yaml

@@ -1,5 +1,6 @@
 domain: your-name-server-domain (eg) test.example.com .. and so on)
 port: tcp, udp open port number(default->53)
+email: your-email(responsible name)
 aws:
   enable: true or false
   access_key: your-aws-access-key

go.mod

@@ -4,6 +4,7 @@ go 1.12
 
 require (
 	github.com/aws/aws-sdk-go v1.19.49
+	github.com/davecgh/go-spew v1.1.1
 	github.com/gjbae1212/go-module v0.4.8
 	github.com/logrusorgru/aurora v0.0.0-20190428105938-cea283e61946
 	github.com/miekg/dns v1.1.14

server/config.go

@@ -18,7 +18,8 @@ import (
 )
 
 const (
-	defaultPort = "53"
+	defaultPort  = "53"
+	defaultRName = "gjbae1212.gmail.com."
 )
 
 type AwsConfig struct {
@@ -31,7 +32,7 @@ type GcpConfig struct {
 	client    *compute.Service
 }
 
-func ParseConfig(config map[interface{}]interface{}) (domain string, port string, awsConfig *AwsConfig, gcpConfig *GcpConfig, err error) {
+func ParseConfig(config map[interface{}]interface{}) (domain, port, rname string, awsConfig *AwsConfig, gcpConfig *GcpConfig, err error) {
 	if config == nil {
 		err = fmt.Errorf("[err] ParseConfig empty params")
 		return
@@ -61,6 +62,12 @@ func ParseConfig(config map[interface{}]interface{}) (domain string, port string
 		}
 	}
 
+	if v, ok := config["email"]; !ok {
+		rname = defaultRName
+	} else {
+		rname = strings.Replace(v.(string), "@", ".", -1) + "."
+	}
+
 	for name, v := range config {
 		switch name.(string) {
 		case "aws":

server/config_test.go

@@ -28,7 +28,7 @@ func TestParseConfig(t *testing.T) {
 	}
 
 	for _, t := range tests {
-		do, _, ac, gc, err := ParseConfig(t.input)
+		do, _, _, ac, gc, err := ParseConfig(t.input)
 		assert.Equal(t.domain, do)
 		assert.Equal(t.awsConfig, ac)
 		assert.Equal(t.gcpConfig, gc)
@@ -48,10 +48,11 @@ func TestParseConfig(t *testing.T) {
 		assert.NoError(err)
 
 		// both enable
-		do, po, ac, gc, err := ParseConfig(config)
+		do, po, rn, ac, gc, err := ParseConfig(config)
 		assert.NoError(err)
 		assert.NotEqual("", do)
 		assert.NotEqual("", po)
+		assert.NotEqual("", rn)
 		assert.True(strings.HasSuffix(do, "."))
 		assert.NotEmpty(ac)
 		assert.NotEmpty(gc)
@@ -62,7 +63,7 @@ func TestParseConfig(t *testing.T) {
 
 		// gcp enable off
 		config["gcp"].(map[interface{}]interface{})["enable"] = "false"
-		do, _, ac, gc, err = ParseConfig(config)
+		do, _, _, ac, gc, err = ParseConfig(config)
 		assert.NoError(err)
 		assert.NotEqual("", do)
 		assert.NotEmpty(ac)
@@ -71,7 +72,7 @@ func TestParseConfig(t *testing.T) {
 		// aws enable off
 		config["gcp"].(map[interface{}]interface{})["enable"] = "true"
 		config["aws"].(map[interface{}]interface{})["enable"] = "false"
-		do, _, ac, gc, err = ParseConfig(config)
+		do, _, _, ac, gc, err = ParseConfig(config)
 		assert.NoError(err)
 		assert.NotEqual("", do)
 		assert.Empty(ac)

server/server.go

@@ -6,6 +6,9 @@ import (
 	"io/ioutil"
 	"log"
 	"net"
+	"strconv"
+	"strings"
+	"time"
 
 	goip "github.com/gjbae1212/go-module/ip"
 	"github.com/logrusorgru/aurora"
@@ -17,13 +20,13 @@ const (
 )
 
 type Server interface {
-	Lookup(search string) ([]*Record, error)
 	Start()
 }
 
 type server struct {
 	domain     string
 	port       string
+	rname      string
 	nameserver string
 	publicIP   string
 	store      *Store
@@ -33,16 +36,138 @@ func (s *server) Start() {
 	udpServer := &dns.Server{Addr: ":" + s.port, Net: "udp"}
 	go udpServer.ListenAndServe()
 	tcpServer := &dns.Server{Addr: ":" + s.port, Net: "tcp"}
-	log.Printf("[start] listen(%s) nameserver(%s) publicIP(%s)\n", aurora.Green(fmt.Sprintf("%s:%s", s.domain, s.port)),
-		aurora.Yellow(s.nameserver), aurora.Cyan(s.publicIP))
+	log.Printf("%s listen(%s) nameserver(%s) domain(%s)\n", aurora.Green("[start]"), aurora.Blue(fmt.Sprintf("%s:%s", s.publicIP, s.port)),
+		aurora.Yellow(s.nameserver), aurora.Cyan(s.domain))
 	tcpServer.ListenAndServe()
 }
 
 func (s *server) Lookup(search string) ([]*Record, error) {
-	return nil, nil
+	search = strings.ToLower(search)
+	seps := strings.Split(search, ".")
+	if len(seps) == 1 {
+		return s.store.Lookup(seps[0])
+	}
+
+	prefix := seps[0]
+	suffix := seps[len(seps)-1]
+	num := 0
+	checkVendor := UNKNOWN
+	if ix, err := strconv.Atoi(prefix); err == nil {
+		num = ix
+	}
+	if suffix == "aws" {
+		checkVendor = AWS
+	} else if suffix == "gcp" {
+		checkVendor = GCP
+	}
+
+	var query string
+	if checkVendor != UNKNOWN && num > 0 { // if prefix is number and suffix is aws or gcp
+		if len(seps) == 2 {
+			query = strings.Join(seps[0:(len(seps)-1)], ".")
+		} else {
+			query = strings.Join(seps[1:(len(seps)-1)], ".")
+		}
+	} else if checkVendor != UNKNOWN && num == 0 { // if prefix is string and suffix is aws or gcp
+		query = strings.Join(seps[0:(len(seps)-1)], ".")
+	} else if checkVendor == UNKNOWN && num > 0 { // if prefix is number and suffix is not aws or gcp
+		query = strings.Join(seps[1:(len(seps))], ".")
+	} else { // etc
+		query = strings.Join(seps[0:(len(seps))], ".")
+	}
+
+	allRecords, err := s.store.Lookup(query)
+	if err != nil {
+		return nil, err
+	}
+
+	var filter []*Record
+	if checkVendor != UNKNOWN {
+		for _, record := range allRecords {
+			if record.Vendor == checkVendor {
+				filter = append(filter, record)
+			}
+		}
+	} else {
+		filter = allRecords
+	}
+
+	var records []*Record
+	if num > 0 {
+		if len(filter) >= num {
+			records = append(records, filter[num-1])
+		}
+	} else {
+		records = filter
+	}
+	return records, nil
 }
 
 func (s *server) dnsRequest(w dns.ResponseWriter, r *dns.Msg) {
+	m := new(dns.Msg)
+	m.SetReply(r)
+	m.Compress = false
+	m.Authoritative = true
+
+	for _, msg := range m.Question {
+		switch msg.Qtype {
+		case dns.TypeNS: // dns nameserver
+			if msg.Name == s.domain {
+				m.Answer = append(m.Answer, s.ns())
+			}
+		case dns.TypeSOA: // dns info
+			if msg.Name == s.domain {
+				m.Answer = append(m.Answer, s.soa())
+			}
+		case dns.TypeA: // ipv4
+			if strings.HasSuffix(msg.Name, s.domain) {
+				prefix := strings.TrimSpace(strings.TrimSuffix(msg.Name, "."+s.domain))
+				records, err := s.Lookup(prefix)
+				if err != nil {
+					log.Printf("[err] lookup %+v\n", err)
+				} else {
+					for _, record := range records {
+						m.Answer = append(m.Answer, &dns.A{
+							Hdr: dns.RR_Header{
+								Name:   msg.Name,
+								Rrtype: dns.TypeA,
+								Class:  dns.ClassINET,
+								Ttl:    uint32(record.TTL() / time.Second),
+							},
+							A: record.PublicIP,
+						})
+					}
+				}
+			}
+		}
+	}
+
+	// if response is not exist.
+	if len(m.Answer) == 0 {
+		m.Ns = append(m.Ns, s.soa())
+	}
+
+	w.WriteMsg(m)
+}
+
+func (s *server) ns() *dns.NS {
+	return &dns.NS{
+		Hdr: dns.RR_Header{Name: s.domain, Rrtype: dns.TypeNS, Class: dns.ClassINET, Ttl: uint32(TTL / time.Second)},
+		Ns:  s.nameserver,
+	}
+}
+
+func (s *server) soa() *dns.SOA {
+	return &dns.SOA{
+		Hdr:     dns.RR_Header{Name: s.domain, Rrtype: dns.TypeSOA, Class: dns.ClassINET, Ttl: uint32(TTL / time.Second)},
+		Ns:      s.nameserver,
+		Mbox:    s.rname,
+		Serial:  uint32(s.store.cacheUpdatedAt.Unix()), // cache updatedAt
+		Refresh: uint32((6 * time.Hour) / time.Second),
+		Retry:   uint32((30 * time.Minute) / time.Second),
+		Expire:  uint32((24 * time.Hour) / time.Second),
+		Minttl:  uint32((2 * time.Minute) / time.Second),
+	}
 }
 
 func NewServer(yamlPath string) (Server, error) {
@@ -59,7 +184,7 @@ func NewServer(yamlPath string) (Server, error) {
 		return nil, err
 	}
 
-	domain, port, awsconfig, gcpconfig, err := ParseConfig(config)
+	domain, port, rname, awsconfig, gcpconfig, err := ParseConfig(config)
 	if err != nil {
 		return nil, err
 	}
@@ -104,7 +229,7 @@ func NewServer(yamlPath string) (Server, error) {
 	if err != nil {
 		return nil, err
 	}
-	s := &server{domain: domain, port: port, nameserver: nameserver, publicIP: publicIP, store: store}
+	s := &server{domain: domain, port: port, nameserver: nameserver, rname: rname, publicIP: publicIP, store: store}
 
 	// register handler
 	dns.HandleFunc(s.domain, s.dnsRequest)

server/store.go

@@ -2,14 +2,15 @@ package server
 
 import (
 	"fmt"
-	"github.com/logrusorgru/aurora"
 	"log"
 	"net"
 	"strconv"
 	"strings"
 	"sync"
 	"time"
 
+	"github.com/logrusorgru/aurora"
+
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/ec2"
 )
@@ -24,14 +25,16 @@ const (
 
 const (
 	CacheName             = "CLOUD-NAME-SERVER"
+	UNKNOWN   CloudVendor = "UNKNOWN"
 	AWS       CloudVendor = "AWS"
 	GCP       CloudVendor = "GCP"
 )
 
 type Store struct {
-	cache   *sync.Map
-	awsconf *AwsConfig
-	gcpconf *GcpConfig
+	awsconf        *AwsConfig
+	gcpconf        *GcpConfig
+	cache          *sync.Map
+	cacheUpdatedAt time.Time
 }
 
 type Record struct {
@@ -141,7 +144,8 @@ func (s *Store) renewal() error {
 		}
 	}
 	s.cache.Store(CacheName, table)
-	log.Printf("[%s][%d] cache table %s\n", aurora.Yellow("update"), count, time.Now().String())
+	s.cacheUpdatedAt = time.Now()
+	log.Printf("%s[%d] cache table %s\n", aurora.Yellow("[update]"), count, time.Now().String())
 	return nil
 }
 

server/store_test.go

@@ -20,7 +20,7 @@ func TestNewStore(t *testing.T) {
 		assert.NoError(err)
 		yaml.Unmarshal(bys, &config)
 		assert.NoError(err)
-		_, _, awsconfig, gcpconfig, err := ParseConfig(config)
+		_, _, _, awsconfig, gcpconfig, err := ParseConfig(config)
 		assert.NoError(err)
 
 		store, err := NewStore(awsconfig, gcpconfig)
@@ -40,7 +40,7 @@ func TestStore_Lookup(t *testing.T) {
 		assert.NoError(err)
 		yaml.Unmarshal(bys, &config)
 		assert.NoError(err)
-		_, _, awsconfig, gcpconfig, err := ParseConfig(config)
+		_, _, _, awsconfig, gcpconfig, err := ParseConfig(config)
 		assert.NoError(err)
 
 		store, err := NewStore(awsconfig, gcpconfig)
@@ -86,7 +86,7 @@ func TestRecord_TTL(t *testing.T) {
 		assert.NoError(err)
 		yaml.Unmarshal(bys, &config)
 		assert.NoError(err)
-		_, _, awsconfig, gcpconfig, err := ParseConfig(config)
+		_, _, _, awsconfig, gcpconfig, err := ParseConfig(config)
 		assert.NoError(err)
 
 		store, err := NewStore(awsconfig, gcpconfig)
@@ -109,7 +109,7 @@ func BenchmarkStore_Lookup(b *testing.B) {
 		config := make(map[interface{}]interface{})
 		bys, _ := ioutil.ReadFile(yamlPath)
 		yaml.Unmarshal(bys, &config)
-		_, _, awsconfig, gcpconfig, _ := ParseConfig(config)
+		_, _, _, awsconfig, gcpconfig, _ := ParseConfig(config)
 		store, _ := NewStore(awsconfig, gcpconfig)
 		for i := 0; i < b.N; i++ {
 			store.Lookup(os.Getenv("TEST_AWS_1"))