diff --git a/dns/dns_proxy.go b/dns/dns_proxy.go index 9451ba8..40afb44 100644 --- a/dns/dns_proxy.go +++ b/dns/dns_proxy.go @@ -755,6 +755,12 @@ func (p *DNSProxy) RemoveDNSRecord(domain string, ip net.IP) { p.recordStore.RemoveRecord(domain, ip) } +// RemoveDNSRecordForSite removes DNS records for a domain that are owned by a specific site. +// If ip is nil, removes all records for the domain that are owned by that site. +func (p *DNSProxy) RemoveDNSRecordForSite(domain string, ip net.IP, siteId int) { + p.recordStore.RemoveRecordForSite(domain, ip, siteId) +} + // GetDNSRecords returns all IP addresses for a domain and record type. // The second return value indicates whether the domain exists. func (p *DNSProxy) GetDNSRecords(domain string, recordType RecordType) ([]net.IP, bool) { diff --git a/dns/dns_records.go b/dns/dns_records.go index 270bae6..68be007 100644 --- a/dns/dns_records.go +++ b/dns/dns_records.go @@ -23,6 +23,7 @@ type recordSet struct { A []net.IP AAAA []net.IP SiteId int + owners map[string]map[int]bool // IP string -> owning site IDs } // DNSRecordStore manages local DNS records for A, AAAA, and PTR queries. @@ -71,9 +72,17 @@ func (s *DNSRecordStore) AddRecord(domain string, ip net.IP, siteId int) error { } if m[domain] == nil { - m[domain] = &recordSet{SiteId: siteId} + m[domain] = &recordSet{SiteId: siteId, owners: make(map[string]map[int]bool)} } rs := m[domain] + if rs.owners == nil { + rs.owners = make(map[string]map[int]bool) + } + ipKey := ip.String() + if rs.owners[ipKey] == nil { + rs.owners[ipKey] = make(map[int]bool) + } + rs.owners[ipKey][siteId] = true if isV4 { for _, existing := range rs.A { if existing.Equal(ip) { @@ -97,7 +106,6 @@ func (s *DNSRecordStore) AddRecord(domain string, ip net.IP, siteId int) error { return nil } - // AddPTRRecord adds a PTR record mapping an IP address to a domain name // ip should be a valid IPv4 or IPv6 address // domain should be in FQDN format (e.g., "example.com.") @@ -123,6 +131,16 @@ func (s *DNSRecordStore) AddPTRRecord(ip net.IP, domain string) error { // If ip is nil, removes all records for the domain (including wildcards) // Automatically removes corresponding PTR records for non-wildcard domains func (s *DNSRecordStore) RemoveRecord(domain string, ip net.IP) { + s.removeRecord(domain, ip, 0, false) +} + +// RemoveRecordForSite removes DNS records owned by a specific site. +// If ip is nil, it removes all records for the domain owned by that site. +func (s *DNSRecordStore) RemoveRecordForSite(domain string, ip net.IP, siteId int) { + s.removeRecord(domain, ip, siteId, true) +} + +func (s *DNSRecordStore) removeRecord(domain string, ip net.IP, siteId int, bySite bool) { s.mu.Lock() defer s.mu.Unlock() @@ -142,8 +160,20 @@ func (s *DNSRecordStore) RemoveRecord(domain string, ip net.IP) { if rs == nil { return } + if rs.owners == nil { + rs.owners = make(map[string]map[int]bool) + } if ip == nil { + if bySite { + rs.A = s.removeOwnedIPs(rs, rs.A, siteId, !isWildcard, domain) + rs.AAAA = s.removeOwnedIPs(rs, rs.AAAA, siteId, !isWildcard, domain) + if len(rs.A) == 0 && len(rs.AAAA) == 0 { + delete(m, domain) + } + return + } + // Remove all records for this domain if !isWildcard { for _, ipAddr := range rs.A { @@ -162,6 +192,19 @@ func (s *DNSRecordStore) RemoveRecord(domain string, ip net.IP) { } // Remove specific IP + ipKey := ip.String() + if bySite { + owners := rs.owners[ipKey] + if len(owners) == 0 { + return + } + delete(owners, siteId) + if len(owners) > 0 { + return + } + delete(rs.owners, ipKey) + } + if ip.To4() != nil { rs.A = removeIP(rs.A, ip) if !isWildcard { @@ -177,6 +220,7 @@ func (s *DNSRecordStore) RemoveRecord(domain string, ip net.IP) { } } } + delete(rs.owners, ipKey) // Clean up empty record sets if len(rs.A) == 0 && len(rs.AAAA) == 0 { @@ -184,6 +228,33 @@ func (s *DNSRecordStore) RemoveRecord(domain string, ip net.IP) { } } +func (s *DNSRecordStore) removeOwnedIPs(rs *recordSet, ips []net.IP, siteId int, removePTR bool, domain string) []net.IP { + kept := make([]net.IP, 0, len(ips)) + for _, ipAddr := range ips { + ipKey := ipAddr.String() + owners := rs.owners[ipKey] + if len(owners) == 0 { + kept = append(kept, ipAddr) + continue + } + + delete(owners, siteId) + if len(owners) > 0 { + kept = append(kept, ipAddr) + continue + } + + delete(rs.owners, ipKey) + if removePTR { + if ptrDomain, exists := s.ptrRecords[ipKey]; exists && ptrDomain == domain { + delete(s.ptrRecords, ipKey) + } + } + } + + return kept +} + // RemovePTRRecord removes a PTR record for an IP address func (s *DNSRecordStore) RemovePTRRecord(ip net.IP) { s.mu.Lock() diff --git a/dns/dns_records_test.go b/dns/dns_records_test.go index 0b4481d..191db6c 100644 --- a/dns/dns_records_test.go +++ b/dns/dns_records_test.go @@ -782,6 +782,37 @@ func TestAutomaticPTRRecordOnRemove(t *testing.T) { } } +func TestRemoveRecordForSiteKeepsSharedAliasIP(t *testing.T) { + store := NewDNSRecordStore() + + domain := "shared.example.com." + ip := net.ParseIP("192.168.1.100") + + if err := store.AddRecord(domain, ip, 10); err != nil { + t.Fatalf("Failed to add record for site 10: %v", err) + } + if err := store.AddRecord(domain, ip, 20); err != nil { + t.Fatalf("Failed to add record for site 20: %v", err) + } + + store.RemoveRecordForSite(domain, ip, 10) + + ips, exists := store.GetRecords(domain, RecordTypeA) + if !exists { + t.Fatal("Expected shared record to still exist after removing one site owner") + } + if len(ips) != 1 || !ips[0].Equal(ip) { + t.Fatalf("Expected shared IP to remain after first owner removal, got %v", ips) + } + + store.RemoveRecordForSite(domain, ip, 20) + + ips, exists = store.GetRecords(domain, RecordTypeA) + if exists { + t.Fatalf("Expected domain to be removed after last owner removal, got %v", ips) + } +} + func TestAutomaticPTRRecordOnRemoveAll(t *testing.T) { store := NewDNSRecordStore() diff --git a/peers/manager.go b/peers/manager.go index 17e57cf..f043373 100644 --- a/peers/manager.go +++ b/peers/manager.go @@ -250,9 +250,6 @@ func (pm *PeerManager) RemovePeer(siteId int) error { // For aliases for _, alias := range peer.Aliases { - if pm.isAliasInUseByOtherPeer(siteId, alias) { - continue - } address := net.ParseIP(alias.AliasAddress) if address == nil { continue @@ -315,11 +312,8 @@ func (pm *PeerManager) UpdatePeer(siteConfig SiteConfig) error { } // Update aliases - // Remove old aliases that are not present anymore and not used by other peers. + // Remove old aliases for _, alias := range oldPeer.Aliases { - if hasAlias(siteConfig.Aliases, alias) || pm.isAliasInUseByOtherPeer(siteConfig.SiteId, alias) { - continue - } address := net.ParseIP(alias.AliasAddress) if address == nil { continue @@ -534,30 +528,6 @@ func (pm *PeerManager) getOwnedAllowedIPs(siteId int) []string { return owned } -// hasAlias returns true when aliases contains the same alias name and alias address. -func hasAlias(aliases []Alias, target Alias) bool { - for _, alias := range aliases { - if strings.EqualFold(alias.Alias, target.Alias) && alias.AliasAddress == target.AliasAddress { - return true - } - } - return false -} - -// isAliasInUseByOtherPeer returns true when another peer has the same alias name and address. -// Must be called with lock held. -func (pm *PeerManager) isAliasInUseByOtherPeer(excludeSiteId int, target Alias) bool { - for otherSiteId, otherPeer := range pm.peers { - if otherSiteId == excludeSiteId { - continue - } - if hasAlias(otherPeer.Aliases, target) { - return true - } - } - return false -} - // addAllowedIp adds an IP (subnet) to the allowed IPs list of a peer // and updates WireGuard configuration if this peer owns the IP. // Must be called with lock held. @@ -795,14 +765,10 @@ func (pm *PeerManager) RemoveAlias(siteId int, aliasName string) error { newAliases = append(newAliases, a) } - if aliasToRemove == nil { - return nil - } - - if !pm.isAliasInUseByOtherPeer(siteId, *aliasToRemove) { + if aliasToRemove != nil { address := net.ParseIP(aliasToRemove.AliasAddress) if address != nil { - pm.dnsProxy.RemoveDNSRecord(aliasName, address) + pm.dnsProxy.RemoveDNSRecordForSite(aliasName, address, siteId) } }