package mock import ( "context" "time" "github.com/miekg/dns" "gitlab.com/inetmock/inetmock/internal/endpoint" "gitlab.com/inetmock/inetmock/pkg/logging" "go.uber.org/zap" ) type dnsHandler struct { logger logging.Logger dnsServer *dns.Server } func (d *dnsHandler) Start(lifecycle endpoint.Lifecycle) (err error) { var options dnsOptions if options, err = loadFromConfig(lifecycle); err != nil { return } d.logger = lifecycle.Logger().With( zap.String("handler_name", lifecycle.Name()), zap.String("address", lifecycle.Uplink().Addr().String()), ) handler := ®exHandler{ handlerName: lifecycle.Name(), fallback: options.Fallback, logger: lifecycle.Logger(), auditEmitter: lifecycle.Audit(), } for _, rule := range options.Rules { d.logger.Info( "register DNS rule", zap.String("pattern", rule.pattern.String()), zap.String("response", rule.response.String()), ) handler.AddRule(rule) } if lifecycle.Uplink().Listener != nil { d.dnsServer = &dns.Server{ Listener: lifecycle.Uplink().Listener, Handler: handler, } } else { d.dnsServer = &dns.Server{ PacketConn: lifecycle.Uplink().PacketConn, Handler: handler, } } go d.startServer() return } func (d *dnsHandler) startServer() { if err := d.dnsServer.ActivateAndServe(); err != nil { d.logger.Error( "failed to start DNS server listener", zap.Error(err), ) } } func (d *dnsHandler) shutdownOnEnd(ctx context.Context) { <-ctx.Done() shutdownCtx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) defer cancel() if err := d.dnsServer.ShutdownContext(shutdownCtx); err != nil { d.logger.Error("failed to shutdown DNS server", zap.Error(err)) } }