2024-08-26 19:10:40 +03:00
|
|
|
package main
|
2024-08-24 19:47:10 +03:00
|
|
|
|
|
|
|
import (
|
2024-08-26 21:07:33 +03:00
|
|
|
"bytes"
|
2024-08-24 19:47:10 +03:00
|
|
|
"net"
|
2024-08-26 21:07:33 +03:00
|
|
|
"slices"
|
2024-08-24 19:47:10 +03:00
|
|
|
"sync"
|
|
|
|
"time"
|
|
|
|
)
|
|
|
|
|
2024-08-26 21:07:33 +03:00
|
|
|
type ARecord struct {
|
|
|
|
Address net.IP
|
|
|
|
Deadline time.Time
|
2024-08-24 19:47:10 +03:00
|
|
|
}
|
|
|
|
|
2024-08-26 21:07:33 +03:00
|
|
|
func NewARecord(addr net.IP, deadline time.Time) *ARecord {
|
|
|
|
return &ARecord{
|
|
|
|
Address: addr,
|
|
|
|
Deadline: deadline,
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
type CNameRecord struct {
|
|
|
|
CName string
|
|
|
|
Deadline time.Time
|
|
|
|
}
|
|
|
|
|
|
|
|
func NewCNameRecord(domainName string, deadline time.Time) *CNameRecord {
|
|
|
|
return &CNameRecord{
|
|
|
|
CName: domainName,
|
|
|
|
Deadline: deadline,
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
type Record struct {
|
|
|
|
Name string
|
|
|
|
ARecords []*ARecord
|
|
|
|
CNameRecords []*CNameRecord
|
|
|
|
}
|
|
|
|
|
|
|
|
func (r *Record) Cleanup() bool {
|
|
|
|
newARecords := make([]*ARecord, 0)
|
|
|
|
for _, record := range r.ARecords {
|
|
|
|
if time.Now().Sub(record.Deadline).Nanoseconds() <= 0 {
|
|
|
|
newARecords = append(newARecords, record)
|
2024-08-24 19:47:10 +03:00
|
|
|
}
|
2024-08-26 21:07:33 +03:00
|
|
|
}
|
|
|
|
r.ARecords = newARecords
|
|
|
|
|
|
|
|
newCNameRecords := make([]*CNameRecord, 0)
|
|
|
|
for _, record := range r.CNameRecords {
|
|
|
|
if time.Now().Sub(record.Deadline).Nanoseconds() <= 0 {
|
|
|
|
newCNameRecords = append(newCNameRecords, record)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
r.CNameRecords = newCNameRecords
|
|
|
|
|
|
|
|
return len(newARecords) == 0 && len(newCNameRecords) == 0
|
|
|
|
}
|
|
|
|
|
|
|
|
func NewRecord(domainName string) *Record {
|
|
|
|
return &Record{
|
|
|
|
Name: domainName,
|
|
|
|
ARecords: make([]*ARecord, 0),
|
|
|
|
CNameRecords: make([]*CNameRecord, 0),
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
type Records struct {
|
|
|
|
mutex sync.RWMutex
|
|
|
|
Records map[string]*Record
|
|
|
|
}
|
|
|
|
|
|
|
|
func (r *Records) getCNames(domainName string, recursive bool, excludeDomains ...string) []string {
|
|
|
|
record, ok := r.Records[domainName]
|
|
|
|
if !ok {
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
if record.Cleanup() {
|
|
|
|
delete(r.Records, domainName)
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
|
|
|
cNameList := make([]string, len(record.CNameRecords))
|
|
|
|
for idx, cnameRecord := range record.CNameRecords {
|
|
|
|
cNameList[idx] = cnameRecord.CName
|
2024-08-24 19:47:10 +03:00
|
|
|
}
|
|
|
|
|
2024-08-25 00:18:15 +03:00
|
|
|
if recursive {
|
|
|
|
origCNameLen := len(cNameList)
|
|
|
|
for i := 0; i < origCNameLen; i++ {
|
2024-08-26 21:07:33 +03:00
|
|
|
if slices.Contains(excludeDomains, cNameList[i]) {
|
|
|
|
continue
|
|
|
|
}
|
|
|
|
|
|
|
|
excludeDomains = append(excludeDomains, cNameList...)
|
|
|
|
parentList := r.getCNames(cNameList[i], true, excludeDomains...)
|
2024-08-25 00:18:15 +03:00
|
|
|
if parentList != nil {
|
|
|
|
cNameList = append(cNameList, parentList...)
|
|
|
|
}
|
2024-08-24 19:47:10 +03:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
return cNameList
|
|
|
|
}
|
|
|
|
|
2024-08-25 00:18:15 +03:00
|
|
|
func (r *Records) GetCNameRecords(domainName string, recursive bool) []string {
|
2024-08-24 19:47:10 +03:00
|
|
|
r.mutex.RLock()
|
|
|
|
defer r.mutex.RUnlock()
|
|
|
|
|
2024-08-25 00:18:15 +03:00
|
|
|
return r.getCNames(domainName, recursive)
|
|
|
|
}
|
|
|
|
|
|
|
|
func (r *Records) GetARecords(domainName string, recursive bool) []net.IP {
|
|
|
|
r.mutex.RLock()
|
|
|
|
defer r.mutex.RUnlock()
|
2024-08-24 19:47:10 +03:00
|
|
|
|
2024-08-25 00:18:15 +03:00
|
|
|
cNameList := []string{domainName}
|
|
|
|
if recursive {
|
|
|
|
cNameList = append(cNameList, r.getCNames(domainName, true)...)
|
|
|
|
}
|
|
|
|
|
|
|
|
aRecords := make([]net.IP, 0)
|
|
|
|
for _, cName := range cNameList {
|
2024-08-26 21:07:33 +03:00
|
|
|
record, ok := r.Records[cName]
|
|
|
|
if !ok {
|
|
|
|
continue
|
|
|
|
}
|
|
|
|
if record.Cleanup() {
|
|
|
|
delete(r.Records, cName)
|
|
|
|
continue
|
|
|
|
}
|
|
|
|
|
|
|
|
for _, aRecord := range record.ARecords {
|
|
|
|
aRecords = append(aRecords, aRecord.Address)
|
2024-08-24 19:47:10 +03:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2024-08-25 00:18:15 +03:00
|
|
|
return aRecords
|
2024-08-24 19:47:10 +03:00
|
|
|
}
|
|
|
|
|
2024-08-26 19:10:40 +03:00
|
|
|
func (r *Records) PutCNameRecord(domainName string, cName string, ttl time.Duration) {
|
2024-08-24 19:47:10 +03:00
|
|
|
r.mutex.Lock()
|
|
|
|
defer r.mutex.Unlock()
|
|
|
|
|
2024-08-26 21:07:33 +03:00
|
|
|
record, ok := r.Records[domainName]
|
|
|
|
if !ok {
|
|
|
|
record = NewRecord(domainName)
|
|
|
|
r.Records[domainName] = record
|
|
|
|
}
|
|
|
|
record.Cleanup()
|
|
|
|
|
|
|
|
for _, cNameRecord := range record.CNameRecords {
|
|
|
|
if cNameRecord.CName == cName {
|
|
|
|
cNameRecord.Deadline = time.Now().Add(ttl)
|
|
|
|
return
|
|
|
|
}
|
2024-08-24 19:47:10 +03:00
|
|
|
}
|
|
|
|
|
2024-08-26 21:07:33 +03:00
|
|
|
record.CNameRecords = append(record.CNameRecords, NewCNameRecord(cName, time.Now().Add(ttl)))
|
2024-08-24 19:47:10 +03:00
|
|
|
}
|
|
|
|
|
2024-08-26 19:10:40 +03:00
|
|
|
func (r *Records) PutARecord(domainName string, addr net.IP, ttl time.Duration) {
|
2024-08-24 19:47:10 +03:00
|
|
|
r.mutex.Lock()
|
|
|
|
defer r.mutex.Unlock()
|
|
|
|
|
2024-08-26 21:07:33 +03:00
|
|
|
record, ok := r.Records[domainName]
|
|
|
|
if !ok {
|
|
|
|
record = NewRecord(domainName)
|
|
|
|
r.Records[domainName] = record
|
2024-08-26 19:46:44 +03:00
|
|
|
}
|
2024-08-26 21:07:33 +03:00
|
|
|
record.Cleanup()
|
2024-08-26 19:46:44 +03:00
|
|
|
|
2024-08-26 21:07:33 +03:00
|
|
|
for _, aRecord := range record.ARecords {
|
|
|
|
if bytes.Compare(aRecord.Address, addr) == 0 {
|
|
|
|
aRecord.Deadline = time.Now().Add(ttl)
|
|
|
|
return
|
|
|
|
}
|
2024-08-26 19:46:44 +03:00
|
|
|
}
|
2024-08-26 21:07:33 +03:00
|
|
|
record.ARecords = append(record.ARecords, NewARecord(addr, time.Now().Add(ttl)))
|
2024-08-26 19:46:44 +03:00
|
|
|
}
|
|
|
|
|
2024-08-25 00:28:34 +03:00
|
|
|
func (r *Records) Cleanup() {
|
|
|
|
r.mutex.Lock()
|
|
|
|
defer r.mutex.Unlock()
|
|
|
|
|
2024-08-26 21:07:33 +03:00
|
|
|
for domainName, record := range r.Records {
|
|
|
|
if record.Cleanup() {
|
|
|
|
delete(r.Records, domainName)
|
2024-08-25 00:28:34 +03:00
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2024-08-24 19:47:10 +03:00
|
|
|
func NewRecords() *Records {
|
|
|
|
return &Records{
|
2024-08-26 21:07:33 +03:00
|
|
|
Records: make(map[string]*Record),
|
2024-08-24 19:47:10 +03:00
|
|
|
}
|
|
|
|
}
|