This commit is contained in:
LovroG05
2025-03-13 10:47:36 +01:00
parent 496a0c2d7a
commit 9f55f64577
6 changed files with 134 additions and 31 deletions

2
go.mod
View File

@@ -5,6 +5,6 @@ go 1.23.0
toolchain go1.23.4
require (
github.com/robfig/cron/v3 v3.0.1
github.com/patrickmn/go-cache v2.1.0+incompatible
golang.org/x/net v0.37.0
)

4
go.sum
View File

@@ -1,4 +1,4 @@
github.com/robfig/cron/v3 v3.0.1 h1:WdRxkvbJztn8LMz/QEvLN5sBU+xKpSqwwUO1Pjr4qDs=
github.com/robfig/cron/v3 v3.0.1/go.mod h1:eQICP3HwyT7UooqI/z+Ov+PtYAWygg1TEWWzGIFLtro=
github.com/patrickmn/go-cache v2.1.0+incompatible h1:HRMgzkcYKYpi3C8ajMPV8OFXaaRUnok+kx1WdO15EQc=
github.com/patrickmn/go-cache v2.1.0+incompatible/go.mod h1:3Qf8kWWT7OJRJbdiICTKqZju1ZixQ/KpMGzzAfe6+WQ=
golang.org/x/net v0.37.0 h1:1zLorHbz+LYj7MQlSf1+2tPIIgibq2eL5xkrGk6f+2c=
golang.org/x/net v0.37.0/go.mod h1:ivrbrMbzFq5J41QOQh0siUuly180yBYtLp+CKbEaFx8=

11
main.go
View File

@@ -1,6 +1,7 @@
package main
import (
"dnsthingymagik/recordcache"
"dnsthingymagik/resolver"
"dnsthingymagik/resolver/entities"
"fmt"
@@ -22,6 +23,8 @@ func main() {
}
}(udpServer)
cache := recordcache.NewCache()
var wg sync.WaitGroup
for {
buf := make([]byte, 514)
@@ -30,13 +33,13 @@ func main() {
log.Println(err)
}
wg.Add(1)
go process(udpServer, addr, buf[:n], &wg)
go process(udpServer, addr, buf[:n], cache, &wg)
}
//recordCache.cache.PurgeExpiredEntries()
//recordCache.recordcache.PurgeExpiredEntries()
}
func process(udp net.PacketConn, addr net.Addr, buf []byte, wg *sync.WaitGroup) {
func process(udp net.PacketConn, addr net.Addr, buf []byte, cache *recordcache.Cache, wg *sync.WaitGroup) {
defer wg.Done()
msg, err := resolver.PacketParser(buf)
if err != nil {
@@ -46,7 +49,7 @@ func process(udp net.PacketConn, addr net.Addr, buf []byte, wg *sync.WaitGroup)
var faq []entities.Record
for _, q := range msg.Questions {
if q.Type == dnsmessage.TypeA {
nips, err := resolver.ResolveDN(q.Name.String(), msg.Header.ID, dnsmessage.TypeA)
nips, err := resolver.ResolveDN(q.Name, msg.Header.ID, dnsmessage.TypeA, cache)
if err != nil {
log.Println(err)
}

98
recordcache/Cache.go Normal file
View File

@@ -0,0 +1,98 @@
package recordcache
import (
"dnsthingymagik/resolver/entities"
"fmt"
"golang.org/x/net/dns/dnsmessage"
"sync"
"time"
)
type Cache struct {
mu sync.RWMutex
records map[string][]entities.Record
}
func NewCache() *Cache {
c := &Cache{
records: make(map[string][]entities.Record),
}
go c.cleanupExpiredRecords()
return c
}
func generateKey(name dnsmessage.Name, rtype dnsmessage.Type) string {
return fmt.Sprintf("%s:%d", name.String(), rtype)
}
func (c *Cache) Set(record entities.Record) {
c.mu.Lock()
defer c.mu.Unlock()
key := generateKey(record.Name, record.RType)
record.ExpireAt = time.Now().Add(time.Duration(record.TTL) * time.Second)
existingRecords, exists := c.records[key]
if exists {
for _, existingRecord := range existingRecords {
if existingRecord.IP.Equal(record.IP) {
return
}
}
}
c.records[key] = append(c.records[key], record)
}
func (c *Cache) Get(name dnsmessage.Name, rtype dnsmessage.Type) ([]entities.Record, bool) {
c.mu.Lock()
defer c.mu.Unlock()
key := generateKey(name, rtype)
records, exists := c.records[key]
if !exists {
return nil, false
}
validRecords := []entities.Record{}
now := time.Now()
for _, record := range records {
remainingTTL := uint32(record.ExpireAt.Sub(now).Seconds())
if remainingTTL > 0 {
record.TTL = remainingTTL
validRecords = append(validRecords, record)
}
}
if len(validRecords) == 0 {
delete(c.records, key)
return nil, false
}
c.records[key] = validRecords
return validRecords, true
}
func (c *Cache) cleanupExpiredRecords() {
ticker := time.NewTicker(time.Hour) // Run every hour
defer ticker.Stop()
for range ticker.C {
c.mu.Lock()
now := time.Now()
for key, records := range c.records {
validRecords := []entities.Record{}
for _, record := range records {
if now.Before(record.ExpireAt) {
validRecords = append(validRecords, record)
}
}
if len(validRecords) == 0 {
delete(c.records, key)
} else {
c.records[key] = validRecords
}
}
c.mu.Unlock()
}
}

View File

@@ -1,34 +1,38 @@
package resolver
import (
"dnsthingymagik/recordcache"
"dnsthingymagik/resolver/entities"
"dnsthingymagik/resolver/query"
"golang.org/x/net/dns/dnsmessage"
"log"
"net"
"strings"
)
func ResolveDN(domainName string, id uint16, rtype dnsmessage.Type) ([]entities.Record, error) {
func ResolveDN(domainName dnsmessage.Name, id uint16, rtype dnsmessage.Type, cache *recordcache.Cache) ([]entities.Record, error) {
if recs, found := cache.Get(domainName, dnsmessage.TypeA); found {
return recs, nil
}
ROOT_SERVERS := []string{
"198.41.0.4",
}
result, err := resolveFromRoot(domainName, id, rtype, ROOT_SERVERS)
result, err := resolveFromRoot(domainName, id, rtype, ROOT_SERVERS, cache)
if err != nil {
return nil, err
}
for _, rec := range result {
cache.Set(rec)
}
return result, nil
}
func resolveFromRoot(domainName string, id uint16, rtype dnsmessage.Type, ROOT_SERVERS []string) ([]entities.Record, error) {
func resolveFromRoot(domainName dnsmessage.Name, id uint16, rtype dnsmessage.Type, ROOT_SERVERS []string, cache *recordcache.Cache) ([]entities.Record, error) {
var records []entities.Record
if !strings.HasSuffix(domainName, ".") {
domainName += "."
}
for _, rootServer := range ROOT_SERVERS {
q := dnsmessage.Message{
@@ -38,7 +42,7 @@ func resolveFromRoot(domainName string, id uint16, rtype dnsmessage.Type, ROOT_S
},
Questions: []dnsmessage.Question{
{
Name: dnsmessage.MustNewName(domainName),
Name: domainName,
Type: rtype,
Class: dnsmessage.ClassINET,
},
@@ -90,7 +94,7 @@ func resolveFromRoot(domainName string, id uint16, rtype dnsmessage.Type, ROOT_S
//resolve from refferal
for nameserver, nsip := range NSIPmap {
if nsip == nil {
nsips, err := ResolveDN(nameserver, id, dnsmessage.TypeA) // resolve nameserver if no ip
nsips, err := ResolveDN(dnsmessage.MustNewName(nameserver), id, dnsmessage.TypeA, cache) // resolve nameserver if no ip
nsip = nsips[0].IP
if err != nil || len(nsip) == 0 {
log.Printf("DNS server %s not resolved: %s", nameserver, err)
@@ -100,7 +104,7 @@ func resolveFromRoot(domainName string, id uint16, rtype dnsmessage.Type, ROOT_S
NSIPmap[nameserver] = nsip
}
ip, err := resolveFromRefferal(domainName, id, dnsmessage.TypeA, nsip.String()) // resolve refferal using nameserver
ip, err := resolveFromRefferal(domainName, id, dnsmessage.TypeA, nsip.String(), cache) // resolve refferal using nameserver
if err != nil {
log.Printf("DNS server %s not resolving domain %s: %s", nameserver, domainName, err)
continue
@@ -114,11 +118,7 @@ func resolveFromRoot(domainName string, id uint16, rtype dnsmessage.Type, ROOT_S
return records, nil
}
func resolveFromRefferal(domainName string, id uint16, rtype dnsmessage.Type, server string) ([]entities.Record, error) {
if !strings.HasSuffix(domainName, ".") {
domainName += "."
}
func resolveFromRefferal(domainName dnsmessage.Name, id uint16, rtype dnsmessage.Type, server string, cache *recordcache.Cache) ([]entities.Record, error) {
q := dnsmessage.Message{
Header: dnsmessage.Header{
ID: id,
@@ -126,7 +126,7 @@ func resolveFromRefferal(domainName string, id uint16, rtype dnsmessage.Type, se
},
Questions: []dnsmessage.Question{
{
Name: dnsmessage.MustNewName(domainName),
Name: domainName,
Type: rtype,
Class: dnsmessage.ClassINET,
},
@@ -173,7 +173,7 @@ func resolveFromRefferal(domainName string, id uint16, rtype dnsmessage.Type, se
for nameserver, nsip := range NSIPmap {
if nsip == nil {
nsips, err := ResolveDN(nameserver, id, dnsmessage.TypeA) // resolve nameserver if no ip
nsips, err := ResolveDN(dnsmessage.MustNewName(nameserver), id, dnsmessage.TypeA, cache) // resolve nameserver if no ip
nsip = nsips[0].IP
if err != nil || len(nsip) == 0 {
log.Printf("DNS server %s not resolved: %s", nameserver, err)
@@ -183,7 +183,7 @@ func resolveFromRefferal(domainName string, id uint16, rtype dnsmessage.Type, se
NSIPmap[nameserver] = nsip
}
ip, err := resolveFromRefferal(domainName, id, dnsmessage.TypeA, nsip.String()) // resolve refferal using nameserver
ip, err := resolveFromRefferal(domainName, id, dnsmessage.TypeA, nsip.String(), cache) // resolve refferal using nameserver
if err != nil {
log.Printf("DNS server %s not resolving domain %s: %s", nameserver, domainName, err)
continue

View File

@@ -3,6 +3,7 @@ package entities
import (
"golang.org/x/net/dns/dnsmessage"
"net"
"time"
)
type Record struct {
@@ -11,4 +12,5 @@ type Record struct {
TTL uint32
Class dnsmessage.Class
Name dnsmessage.Name
ExpireAt time.Time
}