Make it protocol aware

dns/dns_records.go

@@ -0,0 +1,166 @@
+package dns
+
+import (
+	"net"
+	"sync"
+
+	"github.com/miekg/dns"
+)
+
+// RecordType represents the type of DNS record
+type RecordType uint16
+
+const (
+	RecordTypeA    RecordType = RecordType(dns.TypeA)
+	RecordTypeAAAA RecordType = RecordType(dns.TypeAAAA)
+)
+
+// DNSRecordStore manages local DNS records for A and AAAA queries
+type DNSRecordStore struct {
+	mu          sync.RWMutex
+	aRecords    map[string][]net.IP // domain -> list of IPv4 addresses
+	aaaaRecords map[string][]net.IP // domain -> list of IPv6 addresses
+}
+
+// NewDNSRecordStore creates a new DNS record store
+func NewDNSRecordStore() *DNSRecordStore {
+	return &DNSRecordStore{
+		aRecords:    make(map[string][]net.IP),
+		aaaaRecords: make(map[string][]net.IP),
+	}
+}
+
+// AddRecord adds a DNS record mapping (A or AAAA)
+// domain should be in FQDN format (e.g., "example.com.")
+// ip should be a valid IPv4 or IPv6 address
+func (s *DNSRecordStore) AddRecord(domain string, ip net.IP) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	// Ensure domain ends with a dot (FQDN format)
+	if len(domain) == 0 || domain[len(domain)-1] != '.' {
+		domain = domain + "."
+	}
+
+	// Normalize domain to lowercase
+	domain = dns.Fqdn(domain)
+
+	if ip.To4() != nil {
+		// IPv4 address
+		s.aRecords[domain] = append(s.aRecords[domain], ip)
+	} else if ip.To16() != nil {
+		// IPv6 address
+		s.aaaaRecords[domain] = append(s.aaaaRecords[domain], ip)
+	} else {
+		return &net.ParseError{Type: "IP address", Text: ip.String()}
+	}
+
+	return nil
+}
+
+// RemoveRecord removes a specific DNS record mapping
+// If ip is nil, removes all records for the domain
+func (s *DNSRecordStore) RemoveRecord(domain string, ip net.IP) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	// Ensure domain ends with a dot (FQDN format)
+	if len(domain) == 0 || domain[len(domain)-1] != '.' {
+		domain = domain + "."
+	}
+
+	// Normalize domain to lowercase
+	domain = dns.Fqdn(domain)
+
+	if ip == nil {
+		// Remove all records for this domain
+		delete(s.aRecords, domain)
+		delete(s.aaaaRecords, domain)
+		return
+	}
+
+	if ip.To4() != nil {
+		// Remove specific IPv4 address
+		if ips, ok := s.aRecords[domain]; ok {
+			s.aRecords[domain] = removeIP(ips, ip)
+			if len(s.aRecords[domain]) == 0 {
+				delete(s.aRecords, domain)
+			}
+		}
+	} else if ip.To16() != nil {
+		// Remove specific IPv6 address
+		if ips, ok := s.aaaaRecords[domain]; ok {
+			s.aaaaRecords[domain] = removeIP(ips, ip)
+			if len(s.aaaaRecords[domain]) == 0 {
+				delete(s.aaaaRecords, domain)
+			}
+		}
+	}
+}
+
+// GetRecords returns all IP addresses for a domain and record type
+func (s *DNSRecordStore) GetRecords(domain string, recordType RecordType) []net.IP {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+
+	// Normalize domain to lowercase FQDN
+	domain = dns.Fqdn(domain)
+
+	var records []net.IP
+	switch recordType {
+	case RecordTypeA:
+		if ips, ok := s.aRecords[domain]; ok {
+			// Return a copy to prevent external modifications
+			records = make([]net.IP, len(ips))
+			copy(records, ips)
+		}
+	case RecordTypeAAAA:
+		if ips, ok := s.aaaaRecords[domain]; ok {
+			// Return a copy to prevent external modifications
+			records = make([]net.IP, len(ips))
+			copy(records, ips)
+		}
+	}
+
+	return records
+}
+
+// HasRecord checks if a domain has any records of the specified type
+func (s *DNSRecordStore) HasRecord(domain string, recordType RecordType) bool {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+
+	// Normalize domain to lowercase FQDN
+	domain = dns.Fqdn(domain)
+
+	switch recordType {
+	case RecordTypeA:
+		_, ok := s.aRecords[domain]
+		return ok
+	case RecordTypeAAAA:
+		_, ok := s.aaaaRecords[domain]
+		return ok
+	}
+
+	return false
+}
+
+// Clear removes all records from the store
+func (s *DNSRecordStore) Clear() {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	s.aRecords = make(map[string][]net.IP)
+	s.aaaaRecords = make(map[string][]net.IP)
+}
+
+// removeIP is a helper function to remove a specific IP from a slice
+func removeIP(ips []net.IP, toRemove net.IP) []net.IP {
+	result := make([]net.IP, 0, len(ips))
+	for _, ip := range ips {
+		if !ip.Equal(toRemove) {
+			result = append(result, ip)
+		}
+	}
+	return result
+}