MagiTrickle/records.go

278 lines
5.8 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"
"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 {
2024-08-27 01:53:30 +03:00
i := 0
2024-08-26 21:07:33 +03:00
for _, record := range r.ARecords {
2024-08-27 01:53:30 +03:00
if time.Now().Before(record.Deadline) {
r.ARecords[i] = record
i++
2024-08-24 19:47:10 +03:00
}
2024-08-26 21:07:33 +03:00
}
2024-08-27 01:53:30 +03:00
r.ARecords = r.ARecords[:i]
2024-08-26 21:07:33 +03:00
2024-08-27 01:53:30 +03:00
i = 0
2024-08-26 21:07:33 +03:00
for _, record := range r.CNameRecords {
2024-08-27 01:53:30 +03:00
if time.Now().Before(record.Deadline) {
r.CNameRecords[i] = record
i++
2024-08-26 21:07:33 +03:00
}
}
2024-08-27 01:53:30 +03:00
r.CNameRecords = r.CNameRecords[:i]
2024-08-26 21:07:33 +03:00
2024-08-27 01:53:30 +03:00
return len(r.ARecords) == 0 && len(r.CNameRecords) == 0
2024-08-26 21:07:33 +03:00
}
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
}
2024-08-27 01:44:17 +03:00
func (r *Records) getCNames(domainName string, recursive bool, reversive bool) []string {
2024-08-26 21:07:33 +03:00
record, ok := r.Records[domainName]
if !ok {
return nil
}
if record.Cleanup() {
delete(r.Records, domainName)
return nil
}
2024-08-27 01:44:17 +03:00
excludedFromCNameList := map[string]struct{}{
domainName: {},
}
cNameList := make([]string, 0)
for _, cnameRecord := range record.CNameRecords {
if _, exists := excludedFromCNameList[cnameRecord.CName]; !exists {
cNameList = append(cNameList, cnameRecord.CName)
excludedFromCNameList[cnameRecord.CName] = struct{}{}
}
2024-08-24 19:47:10 +03:00
}
2024-08-25 00:18:15 +03:00
if recursive {
2024-08-27 01:44:17 +03:00
excludedFromProcess := map[string]struct{}{
domainName: {},
}
processingList := cNameList
for len(processingList) > 0 {
newProcessingList := []string{}
for _, cname := range processingList {
if _, exists := excludedFromProcess[cname]; exists {
continue
}
record, ok := r.Records[cname]
if !ok {
continue
}
if record.Cleanup() {
delete(r.Records, cname)
continue
}
for _, cNameRecord := range record.CNameRecords {
if _, exists := excludedFromCNameList[cNameRecord.CName]; !exists {
cNameList = append(cNameList, cNameRecord.CName)
excludedFromCNameList[cNameRecord.CName] = struct{}{}
}
newProcessingList = append(newProcessingList, cNameRecord.CName)
}
2024-08-26 21:07:33 +03:00
}
2024-08-27 01:44:17 +03:00
processingList = newProcessingList
}
}
2024-08-26 21:07:33 +03:00
2024-08-27 01:44:17 +03:00
if reversive {
excludedFromProcess := make(map[string]struct{})
processingList := []string{domainName}
for len(processingList) > 0 {
nextProcessingList := make([]string, 0)
for _, target := range processingList {
if _, exists := excludedFromProcess[target]; exists {
continue
}
for cname, record := range r.Records {
if record.Cleanup() {
delete(r.Records, cname)
continue
}
for _, cnameRecord := range record.CNameRecords {
if cnameRecord.CName != target {
continue
}
if _, exists := excludedFromCNameList[record.Name]; !exists {
cNameList = append(cNameList, record.Name)
excludedFromCNameList[record.Name] = struct{}{}
}
nextProcessingList = append(nextProcessingList, record.Name)
break
}
}
excludedFromProcess[target] = struct{}{}
2024-08-25 00:18:15 +03:00
}
2024-08-27 01:44:17 +03:00
processingList = nextProcessingList
2024-08-24 19:47:10 +03:00
}
}
return cNameList
}
2024-08-27 01:44:17 +03:00
func (r *Records) GetCNameRecords(domainName string, recursive bool, reversive bool) []string {
2024-08-24 19:47:10 +03:00
r.mutex.RLock()
defer r.mutex.RUnlock()
2024-08-27 01:44:17 +03:00
return r.getCNames(domainName, recursive, reversive)
2024-08-25 00:18:15 +03:00
}
2024-08-27 01:44:17 +03:00
func (r *Records) GetARecords(domainName string, recursive bool, reversive bool) []net.IP {
2024-08-25 00:18:15 +03:00
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 {
2024-08-27 01:44:17 +03:00
cNameList = append(cNameList, r.getCNames(domainName, true, reversive)...)
2024-08-25 00:18:15 +03:00
}
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-09-06 06:34:52 +03:00
func (r *Records) ListKnownDomains() []string {
r.mutex.Lock()
defer r.mutex.Unlock()
domains := make([]string, 0)
for name, record := range r.Records {
if record.Cleanup() {
delete(r.Records, name)
continue
}
domains = append(domains, name)
}
return domains
}
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
}
}