This repository has been archived on 2023-08-14. You can view files and clone it, but cannot push or open issues or pull requests.
goredns/main.go

116 lines
2.5 KiB
Go
Raw Normal View History

2021-03-02 22:58:47 +00:00
package main
import (
2021-04-03 19:39:56 +00:00
"context"
"fmt"
"log"
"os"
"strconv"
"strings"
"time"
2021-03-02 22:58:47 +00:00
2021-04-03 19:39:56 +00:00
"github.com/miekg/dns"
"go.mongodb.org/mongo-driver/bson"
"go.mongodb.org/mongo-driver/mongo"
"go.mongodb.org/mongo-driver/mongo/options"
2021-05-27 10:02:09 +00:00
"go.mongodb.org/mongo-driver/x/mongo/driver/connstring"
2021-03-02 22:58:47 +00:00
)
2021-05-27 10:02:09 +00:00
type Elem struct {
2021-04-03 19:39:56 +00:00
Ip []string
2021-03-02 22:58:47 +00:00
}
2021-04-03 19:39:56 +00:00
var mongoUri string = os.Getenv("MONGO_URI")
var collectionName string = os.Getenv("GOREDNS_COLLECTION")
2021-04-03 19:39:56 +00:00
func appendResults(etype string, name string, m *dns.Msg, cur *mongo.Cursor) int {
count := 0
for cur.Next(context.TODO()) {
2021-05-27 10:02:09 +00:00
var elem Elem
2021-04-03 19:39:56 +00:00
err := cur.Decode(&elem)
if err != nil {
log.Fatal(err)
}
for _, ip := range elem.Ip {
tp := "A"
if strings.Contains(ip, ":") {
tp = "AAAA"
}
2021-03-02 22:58:47 +00:00
2021-04-03 19:39:56 +00:00
if etype != tp {
continue
}
log.Printf("Appending: %s %s %s\n", name, tp, ip)
rr, err := dns.NewRR(fmt.Sprintf("%s. %s %s", name, tp, ip))
if err == nil {
m.Answer = append(m.Answer, rr)
count += 1
}
}
}
return count
}
2021-03-02 22:58:47 +00:00
2021-04-03 19:39:56 +00:00
func query(tp string, name string, m *dns.Msg, coll *mongo.Collection) {
// TODO: Validate `name` against RE_FQDN
log.Printf("Query %s for %s\n", tp, name)
cur, err := coll.Find(context.TODO(), bson.M{"dns.fqdn": name})
if err != nil {
log.Fatal(err)
}
2021-05-27 10:02:09 +00:00
2021-04-03 19:39:56 +00:00
if appendResults(tp, name, m, cur) == 0 {
cur, err := coll.Find(context.TODO(), bson.M{"dns.san": name})
if err != nil {
log.Fatal(err)
}
appendResults(tp, name, m, cur)
}
}
func wrapper(coll *mongo.Collection) func(dns.ResponseWriter, *dns.Msg) {
return func(w dns.ResponseWriter, r *dns.Msg) {
m := new(dns.Msg)
m.SetReply(r)
m.Compress = false
switch r.Opcode {
case dns.OpcodeQuery:
for _, q := range m.Question {
switch q.Qtype {
case dns.TypeA:
query("A", q.Name[:len(q.Name)-1], m, coll)
case dns.TypeAAAA:
query("AAAA", q.Name[:len(q.Name)-1], m, coll)
}
}
}
w.WriteMsg(m)
}
2021-03-02 22:58:47 +00:00
}
func main() {
2021-05-27 10:02:09 +00:00
cs, err := connstring.ParseAndValidate(mongoUri)
2021-04-03 19:39:56 +00:00
client, err := mongo.NewClient(options.Client().ApplyURI(mongoUri))
if err != nil {
log.Fatal(err)
}
ctx, _ := context.WithTimeout(context.Background(), 10*time.Second)
err = client.Connect(ctx)
if err != nil {
log.Fatal(err)
}
2021-03-02 22:58:47 +00:00
2021-05-27 10:02:09 +00:00
coll := client.Database(cs.Database).Collection(collectionName)
2021-04-03 19:39:56 +00:00
defer client.Disconnect(ctx)
dns.HandleFunc(".", wrapper(coll))
port := 53
server := &dns.Server{Addr: ":" + strconv.Itoa(port), Net: "udp"}
log.Printf("Starting at %d\n", port)
err2 := server.ListenAndServe()
defer server.Shutdown()
if err2 != nil {
log.Fatalf("Failed to start server: %s\n ", err2.Error())
}
2021-03-02 22:58:47 +00:00
}