MagiTrickle/records.go

199 lines
4.0 KiB
Go
Raw Normal View History

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
}
}