tatsuyadns.go

package main

import (
	"io"
	"log"
	"net/http"
	"os"
	"strings"
	"sync"

	"github.com/miekg/dns"
)

var ipAddr = []byte{} // redacted
var domain = "u.isucon.dev."
var subdomains = sync.Map{}

func isTargetDomain(name string) bool {
	if name == domain {
		return true
	} else {
		if strings.HasSuffix(name, "."+domain) {
			_, ok := subdomains.Load(strings.TrimSuffix(name, "."+domain))
			if ok {
				return true
			}
		}
	}
	return false
}

func handleReflect(w dns.ResponseWriter, r *dns.Msg) {
	// log.Printf("%+v", *r)
	m := new(dns.Msg)
	m.SetReply(r)

	name := r.Question[0].Name
	if isTargetDomain(name) {
		rr := &dns.A{
			Hdr: dns.RR_Header{
				Name:   r.Question[0].Name,
				Rrtype: dns.TypeA,
				Class:  dns.ClassINET,
				Ttl:    0,
			},
			A: ipAddr,
		}
		m.Answer = append(m.Answer, rr)
	}

	w.WriteMsg(m)
}

func handleAdd(w http.ResponseWriter, r *http.Request) {
	defer r.Body.Close()

	if r.Method != http.MethodPost {
		w.WriteHeader(http.StatusMethodNotAllowed)
		return
	}

	body, err := io.ReadAll(r.Body)
	if err != nil {
		log.Printf("ERROR: %s", err)
		w.WriteHeader(http.StatusInternalServerError)
		return
	}
	subdomains.Store(string(body), struct{}{})
}

func handleReset(w http.ResponseWriter, r *http.Request) {
	defer r.Body.Close()

	if r.Method != http.MethodPost {
		w.WriteHeader(http.StatusMethodNotAllowed)
		return
	}

	subdomains = sync.Map{}
	for d := range initSubdomains {
		subdomains.Store(d, struct{}{})
	}
}

func main() {
	for d := range initSubdomains {
		subdomains.Store(d, struct{}{})
	}

	dnsListen := os.Getenv("DNS_LISTEN")

	dns.HandleFunc("u.isucon.dev.", handleReflect)
	server := &dns.Server{Addr: dnsListen, Net: "udp"}
	errCh := make(chan error)
	go func() {
		log.Printf("DNS start")
		errCh <- server.ListenAndServe()
	}()
	go func() {
		log.Printf("HTTP start")
		http.HandleFunc("/add", handleAdd)
		http.HandleFunc("/reset", handleReset)
		errCh <- http.ListenAndServe(":10000", nil)
	}()
	log.Fatal(<-errCh)
}