2021-01-13 17:07:04 +00:00
|
|
|
package mock
|
2020-04-01 22:58:44 +00:00
|
|
|
|
|
|
|
import (
|
2020-10-02 09:56:48 +00:00
|
|
|
"context"
|
2021-02-10 20:26:45 +00:00
|
|
|
"time"
|
2020-12-26 13:11:49 +00:00
|
|
|
|
2020-04-01 22:58:44 +00:00
|
|
|
"github.com/miekg/dns"
|
2021-02-10 20:26:45 +00:00
|
|
|
"gitlab.com/inetmock/inetmock/internal/endpoint"
|
2020-12-26 13:11:49 +00:00
|
|
|
"gitlab.com/inetmock/inetmock/pkg/logging"
|
2020-04-01 22:58:44 +00:00
|
|
|
"go.uber.org/zap"
|
|
|
|
)
|
|
|
|
|
|
|
|
type dnsHandler struct {
|
2020-04-13 22:14:56 +00:00
|
|
|
logger logging.Logger
|
2021-02-10 20:26:45 +00:00
|
|
|
dnsServer *dns.Server
|
2020-04-01 22:58:44 +00:00
|
|
|
}
|
|
|
|
|
2021-02-10 20:26:45 +00:00
|
|
|
func (d *dnsHandler) Start(lifecycle endpoint.Lifecycle) (err error) {
|
2020-10-02 09:56:48 +00:00
|
|
|
var options dnsOptions
|
2021-02-10 20:26:45 +00:00
|
|
|
if options, err = loadFromConfig(lifecycle); err != nil {
|
2020-10-02 09:56:48 +00:00
|
|
|
return
|
|
|
|
}
|
|
|
|
|
2021-02-10 20:26:45 +00:00
|
|
|
d.logger = lifecycle.Logger().With(
|
|
|
|
zap.String("handler_name", lifecycle.Name()),
|
|
|
|
zap.String("address", lifecycle.Uplink().Addr().String()),
|
2020-10-02 09:56:48 +00:00
|
|
|
)
|
2020-04-01 22:58:44 +00:00
|
|
|
|
|
|
|
handler := ®exHandler{
|
2021-02-10 20:26:45 +00:00
|
|
|
handlerName: lifecycle.Name(),
|
2021-01-07 21:00:12 +00:00
|
|
|
fallback: options.Fallback,
|
2021-02-10 20:26:45 +00:00
|
|
|
logger: lifecycle.Logger(),
|
|
|
|
auditEmitter: lifecycle.Audit(),
|
2020-04-01 22:58:44 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
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)
|
|
|
|
}
|
|
|
|
|
2021-02-10 20:26:45 +00:00
|
|
|
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,
|
|
|
|
}
|
2020-04-01 22:58:44 +00:00
|
|
|
}
|
|
|
|
|
2021-02-10 20:26:45 +00:00
|
|
|
go d.startServer()
|
2020-04-13 22:14:56 +00:00
|
|
|
return
|
2020-04-01 22:58:44 +00:00
|
|
|
}
|
|
|
|
|
2021-02-10 20:26:45 +00:00
|
|
|
func (d *dnsHandler) startServer() {
|
|
|
|
if err := d.dnsServer.ActivateAndServe(); err != nil {
|
2020-04-01 22:58:44 +00:00
|
|
|
d.logger.Error(
|
|
|
|
"failed to start DNS server listener",
|
|
|
|
zap.Error(err),
|
|
|
|
)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2021-02-10 20:26:45 +00:00
|
|
|
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))
|
2020-04-01 22:58:44 +00:00
|
|
|
}
|
|
|
|
}
|