135 lines
4.3 KiB
Go
135 lines
4.3 KiB
Go
package dns
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"net"
|
|
"path/filepath"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
|
|
miekgDNS "github.com/miekg/dns"
|
|
"github.com/spf13/viper"
|
|
)
|
|
|
|
type ResponseWriter = miekgDNS.ResponseWriter
|
|
type Msg = miekgDNS.Msg
|
|
type HandlerFunc = miekgDNS.HandlerFunc
|
|
|
|
// DNS struct represents the DNS server with its configuration, server, mux, records, and resolver.
|
|
type DNS struct {
|
|
Config Config
|
|
Server Server
|
|
Mux Mux
|
|
Records Records
|
|
Resolver Resolver
|
|
}
|
|
|
|
// Config struct holds the configuration settings.
|
|
type Config struct {
|
|
*viper.Viper
|
|
}
|
|
|
|
// Server struct wraps the miekgDNS.Server.
|
|
type Server struct {
|
|
*miekgDNS.Server
|
|
}
|
|
|
|
// Mux struct wraps the miekgDNS.ServeMux.
|
|
type Mux struct {
|
|
*miekgDNS.ServeMux
|
|
}
|
|
|
|
// Records struct holds the DNS records.
|
|
type Records struct {
|
|
TypeA map[string]string
|
|
TypeNS map[string]string // Add NS records
|
|
}
|
|
|
|
// Resolver struct wraps the net.Resolver.
|
|
type Resolver struct {
|
|
*net.Resolver
|
|
}
|
|
|
|
// New initializes a new DNS instance with the given configuration file path.
|
|
func New(filePath string) *DNS {
|
|
app := &DNS{}
|
|
app.ConfigInit(filePath) // Initialize configuration
|
|
app.Records.TypeA = app.Config.GetStringMapString("records.type_a") // Load Type A records
|
|
alternate_resolver_ip := app.Config.GetString("alternate_resolver.ip") // Get alternate resolver IP
|
|
alternate_resolver_port := strconv.Itoa(app.Config.GetInt("alternate_resolver.port")) // Get alternate resolver port
|
|
app.ServerInit() // Initialize server
|
|
app.MuxInit() // Initialize multiplexer
|
|
app.Resolver = NewResolver(alternate_resolver_ip + ":" + alternate_resolver_port) // Initialize resolver
|
|
app.Server.Handler = app.Mux // Set handler
|
|
|
|
return app
|
|
}
|
|
|
|
// Run starts the DNS server.
|
|
func (a *DNS) Run() {
|
|
err := a.Server.ListenAndServe() // Start server
|
|
if err != nil {
|
|
fmt.Println(err.Error()) // Print error if any
|
|
}
|
|
}
|
|
|
|
// ConfigInit initializes the configuration from the given file path.
|
|
func (a *DNS) ConfigInit(filePath string) {
|
|
dir, file := filepath.Split(filePath) // Split file path
|
|
fe := filepath.Ext(file) // Get file extension
|
|
ext, _ := strings.CutPrefix(fe, ".") // Remove leading dot from extension
|
|
name, _ := strings.CutSuffix(file, fe) // Get file name without extension
|
|
|
|
a.Config.Viper = viper.New() // Create new Viper instance
|
|
a.Config.AddConfigPath(dir) // Add config path
|
|
a.Config.SetConfigName(name) // Set config name
|
|
a.Config.SetConfigType(ext) // Set config type
|
|
a.Config.ReadInConfig() // Read config
|
|
}
|
|
|
|
// ServerInit initializes the DNS server with the configuration settings.
|
|
func (a *DNS) ServerInit() {
|
|
port := a.Config.GetInt("port") // Get port from config
|
|
a.Server.Server = &miekgDNS.Server{
|
|
Addr: ":" + strconv.Itoa(port), // Set server address
|
|
Net: a.Config.GetString("network"), // Set network type
|
|
Handler: nil, // Handler will be set later
|
|
}
|
|
}
|
|
|
|
// MuxInit initializes the DNS request multiplexer.
|
|
func (a *DNS) MuxInit() {
|
|
a.Mux.ServeMux = miekgDNS.NewServeMux() // Create new ServeMux
|
|
a.Mux.HandleFunc(".", a.HandleTypeA) // Handle all requests with HandleTypeA
|
|
a.Mux.HandleFunc(".", a.HandleTypeNS) // Handle all requests with HandleTypeNS
|
|
}
|
|
|
|
// NewResolver creates a new DNS resolver with the given DNS server address.
|
|
func NewResolver(DNSserverAddr string) Resolver {
|
|
return Resolver{
|
|
&net.Resolver{
|
|
PreferGo: true, // Prefer Go's resolver
|
|
Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
|
|
d := net.Dialer{Timeout: 5 * time.Second} // Set dial timeout
|
|
return d.DialContext(ctx, "udp", DNSserverAddr) // Dial DNS server
|
|
},
|
|
},
|
|
}
|
|
}
|
|
|
|
// Lookup performs a DNS lookup for the given address using the resolver.
|
|
func (r *Resolver) Lookup(lookupAddr string) string {
|
|
var resp []string
|
|
var err error
|
|
|
|
ctx := context.Background() // Create background context
|
|
resp, err = r.LookupHost(ctx, lookupAddr) // Perform lookup
|
|
|
|
if err != nil {
|
|
fmt.Println(err.Error()) // Print error if any
|
|
}
|
|
|
|
return resp[0] // Return first result
|
|
}
|