// Package route53 implements a plugin that returns resource records
// from AWS route53.
package route53

import (
	"context"
	"errors"
	"fmt"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/coredns/coredns/plugin"
	"github.com/coredns/coredns/plugin/file"
	"github.com/coredns/coredns/plugin/pkg/fall"
	"github.com/coredns/coredns/plugin/pkg/upstream"
	"github.com/coredns/coredns/request"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/route53"
	"github.com/aws/aws-sdk-go-v2/service/route53/types"
	"github.com/miekg/dns"
)

// Route53 is a plugin that returns RR from AWS route53.
type Route53 struct {
	Next plugin.Handler
	Fall fall.F

	zoneNames []string
	client    route53Client
	upstream  *upstream.Upstream
	refresh   time.Duration

	zMu   sync.RWMutex
	zones zones
}

type zone struct {
	id  string
	z   *file.Zone
	dns string
}

type zones map[string][]*zone

// New reads from the keys map which uses domain names as its key and hosted
// zone id lists as its values, validates that each domain name/zone id pair
// does exist, and returns a new *Route53. In addition to this, upstream is use
// for doing recursive queries against CNAMEs. Returns error if it cannot
// verify any given domain name/zone id pair.
func New(ctx context.Context, c route53Client, keys map[string][]string, refresh time.Duration) (*Route53, error) {
	zones := make(map[string][]*zone, len(keys))
	zoneNames := make([]string, 0, len(keys))
	for dns, hostedZoneIDs := range keys {
		for _, hostedZoneID := range hostedZoneIDs {
			_, err := c.ListHostedZonesByName(ctx, &route53.ListHostedZonesByNameInput{
				DNSName:      aws.String(dns),
				HostedZoneId: aws.String(hostedZoneID),
			})
			if err != nil {
				return nil, err
			}
			if _, ok := zones[dns]; !ok {
				zoneNames = append(zoneNames, dns)
			}
			zones[dns] = append(zones[dns], &zone{id: hostedZoneID, dns: dns, z: file.NewZone(dns, "")})
		}
	}
	return &Route53{
		client:    c,
		zoneNames: zoneNames,
		zones:     zones,
		upstream:  upstream.New(),
		refresh:   refresh,
	}, nil
}

// Run executes first update, spins up an update forever-loop.
// Returns error if first update fails.
func (h *Route53) Run(ctx context.Context) error {
	if err := h.updateZones(ctx); err != nil {
		return err
	}
	go func() {
		timer := time.NewTimer(h.refresh)
		defer timer.Stop()
		for {
			timer.Reset(h.refresh)
			select {
			case <-ctx.Done():
				log.Debugf("Breaking out of Route53 update loop for %v: %v", h.zoneNames, ctx.Err())
				return
			case <-timer.C:
				if err := h.updateZones(ctx); err != nil && ctx.Err() == nil /* Don't log error if ctx expired. */ {
					log.Errorf("Failed to update zones %v: %v", h.zoneNames, err)
				}
			}
		}
	}()
	return nil
}

// ServeDNS implements the plugin.Handler.ServeDNS.
func (h *Route53) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) {
	state := request.Request{W: w, Req: r}
	qname := state.Name()

	zName := plugin.Zones(h.zoneNames).Matches(qname)
	if zName == "" {
		return plugin.NextOrFailure(h.Name(), h.Next, ctx, w, r)
	}
	z, ok := h.zones[zName]
	if !ok || z == nil {
		return dns.RcodeServerFailure, nil
	}

	m := new(dns.Msg)
	m.SetReply(r)
	m.Authoritative = true
	var result file.Result
	for _, hostedZone := range z {
		h.zMu.RLock()
		m.Answer, m.Ns, m.Extra, result = hostedZone.z.Lookup(ctx, state, qname)
		h.zMu.RUnlock()

		// Take the answer if it's non-empty OR if there is another
		// record type exists for this name (NODATA).
		if len(m.Answer) != 0 || result == file.NoData {
			break
		}
	}

	if len(m.Answer) == 0 && result != file.NoData && h.Fall.Through(qname) {
		return plugin.NextOrFailure(h.Name(), h.Next, ctx, w, r)
	}

	switch result {
	case file.Success:
	case file.NoData:
	case file.NameError:
		m.Rcode = dns.RcodeNameError
	case file.Delegation:
		m.Authoritative = false
	case file.ServerFailure:
		return dns.RcodeServerFailure, nil
	}

	w.WriteMsg(m)
	return dns.RcodeSuccess, nil
}

const escapeSeq = "\\"

// maybeUnescape parses s and converts escaped ASCII codepoints (in octal) back
// to its ASCII representation.
//
// From AWS docs:
//
// "If the domain name includes any characters other than a to z, 0 to 9, -
// (hyphen), or _ (underscore), Route 53 API actions return the characters as
// escape codes."
//
// For our purposes (and with respect to RFC 1035), we'll fish for a-z, 0-9,
// '-', '.' and '*' as the leftmost character (for wildcards) and throw error
// for everything else.
//
// Example:
//
//	`\\052.example.com.` -> `*.example.com`
//	`\\137.example.com.` -> error ('_' is not valid)
func maybeUnescape(s string) (string, error) {
	var outSb strings.Builder
	for {
		i := strings.Index(s, escapeSeq)
		if i < 0 {
			return outSb.String() + s, nil
		}

		outSb.WriteString(s[:i])

		li, ri := i+len(escapeSeq), i+len(escapeSeq)+3
		if ri > len(s) {
			return "", fmt.Errorf("invalid escape sequence: '%s%s'", escapeSeq, s[li:])
		}
		// Parse `\\xxx` in base 8 (2nd arg) and attempt to fit into
		// 8-bit result (3rd arg).
		n, err := strconv.ParseInt(s[li:ri], 8, 8)
		if err != nil {
			return "", fmt.Errorf("invalid escape sequence: '%s%s'", escapeSeq, s[li:ri])
		}

		r := rune(n)
		switch {
		case r >= rune('a') && r <= rune('z'): // Route53 converts everything to lowercase.
		case r >= rune('0') && r <= rune('9'):
		case r == rune('*'):
			if outSb.Len() != 0 {
				return "", errors.New("`*' only supported as wildcard (leftmost label)")
			}
		case r == rune('-'):
		case r == rune('.'):
		default:
			return "", fmt.Errorf("invalid character: %s%#03o", escapeSeq, r)
		}

		outSb.WriteString(string(r))

		s = s[i+len(escapeSeq)+3:]
	}
}

func updateZoneFromRRS(rrs *types.ResourceRecordSet, z *file.Zone) error {
	for _, rr := range rrs.ResourceRecords {
		n, err := maybeUnescape(aws.ToString(rrs.Name))
		if err != nil {
			return fmt.Errorf("failed to unescape `%s' name: %v", aws.ToString(rrs.Name), err)
		}
		v, err := maybeUnescape(aws.ToString(rr.Value))
		if err != nil {
			return fmt.Errorf("failed to unescape `%s' value: %v", aws.ToString(rr.Value), err)
		}

		// Assemble RFC 1035 conforming record to pass into dns scanner.
		rfc1035 := fmt.Sprintf("%s %d IN %s %s", n, aws.ToInt64(rrs.TTL), rrs.Type, v)
		r, err := dns.NewRR(rfc1035)
		if err != nil {
			return fmt.Errorf("failed to parse resource record: %v", err)
		}

		z.Insert(r)
	}
	return nil
}

// updateZones re-queries resource record sets for each zone and updates the
// zone object.
// Returns error if any zones error'ed out, but waits for other zones to
// complete first.
func (h *Route53) updateZones(ctx context.Context) error {
	errc := make(chan error)
	defer close(errc)
	for zName, z := range h.zones {
		go func(zName string, z []*zone) {
			var err error
			defer func() {
				errc <- err
			}()

			for i, hostedZone := range z {
				newZ := file.NewZone(zName, "")
				newZ.Upstream = h.upstream
				in := &route53.ListResourceRecordSetsInput{
					HostedZoneId: aws.String(hostedZone.id),
					MaxItems:     aws.Int32(1000),
				}
				complete := false
				var out *route53.ListResourceRecordSetsOutput
				for out, err = h.client.ListResourceRecordSets(ctx, in); !complete; out, err = h.client.ListResourceRecordSets(ctx, in) {
					if err != nil {
						err = fmt.Errorf("failed to list resource records for %v:%v from route53: %v", zName, hostedZone.id, err)
						return
					}
					for _, rrs := range out.ResourceRecordSets {
						if err := updateZoneFromRRS(&rrs, newZ); err != nil {
							// Maybe unsupported record type. Log and carry on.
							log.Warningf("Failed to process resource record set: %v", err)
						}
					}
					if out.IsTruncated {
						in.StartRecordName = out.NextRecordName
						in.StartRecordType = out.NextRecordType
						in.StartRecordIdentifier = out.NextRecordIdentifier
					} else {
						complete = true
					}
				}
				h.zMu.Lock()
				(*z[i]).z = newZ
				h.zMu.Unlock()
			}
		}(zName, z)
	}
	// Collect errors (if any). This will also sync on all zones updates
	// completion.
	var errs []string
	for range len(h.zones) {
		err := <-errc
		if err != nil {
			errs = append(errs, err.Error())
		}
	}
	if len(errs) != 0 {
		return fmt.Errorf("errors updating zones: %v", errs)
	}
	return nil
}

// Name implements plugin.Handler.Name.
func (h *Route53) Name() string { return "route53" }
