// Copyright Martin Dosch.
// Use of this source code is governed by the BSD-2-clause
// license that can be found in the LICENSE file.

package xmppsrv

import (
	"context"
	"crypto/tls"
	"errors"
	"fmt"
	"math/rand"
	"net"
	"sort"
	"strings"
)

// The Config type is used to configure which resolver is used.
// Set Resolver to an IP address and DoT to false to use the
// DNS resolver on port 53 on the given IP address.
// Set Resolver to an IP and domain name (e.g. "5.1.66.255#dot.ffmuc.net")
// and DoT to true to use DNSoverTLS (DoT) on that resolver.
type Config struct {
	Resolver string
	DoT      bool
}

// The SRV type is basically the same as net.SRV (see
// https://pkg.go.dev/net#SRV but it adds the string Type
// which contains either xmpp-client, xmpps-client,
// xmpp-server or xmpps-server.
// This is especially useful for the function LookupXmppClient
// which returns a mix of xmpp-client and xmpps-client
// records.
type SRV struct {
	Type     string
	Target   string
	Port     uint16
	Priority uint16
	Weight   uint16
}
type byPriority []SRV

func (o byPriority) Len() int      { return len(o) }
func (o byPriority) Swap(i, j int) { o[i], o[j] = o[j], o[i] }
func (o byPriority) Less(i, j int) bool {
	return o[i].Priority < o[j].Priority
}

func getCNAME(server string) (string, error) {
	var cname string
	// Lookup CNAME
	for i := 0; i <= 5; i++ {
		cname, err := net.LookupCNAME(server)
		// Return the requested domain unchanged if the
		// CNAME record can't be resolved. This can
		// happen if there is no CNAME record and no
		// A/AAAA record. But there could still be SRV
		// records.
		if cname == "" {
			return server, nil
		}
		if cname == server {
			return cname, err
		}
		server = cname
	}
	return cname, errors.New("CNAME loop detected")
}

func getCNAMEDoT(server string, resolver string) (string, error) {
	var cname string
	var err error
	var d tls.Dialer
	s := strings.Split(resolver, "#")
	if len(s) != 2 {
		return cname, errors.New("wrong DoT server syntax")
	}
	ip := net.ParseIP(s[0])
	if ip == nil {
		return cname, errors.New("invalid resolver IP")
	}
	dnsServer := net.JoinHostPort(s[0], "853")
	d.Config = &tls.Config{
		ServerName: s[1],
	}
	r := &net.Resolver{
		Dial: func(cont context.Context, network, address string) (net.Conn, error) {
			return d.DialContext(cont, "tcp", dnsServer)
		},
	}
	for i := 0; i <= 5; i++ {
		cname, err = r.LookupCNAME(context.Background(), server)
		// Return the requested domain unchanged if the
		// CNAME record can't be resolved. This can
		// happen if there is no CNAME record and no
		// A/AAAA record. But there could still be SRV
		// records.
		if cname == "" {
			return server, nil
		}
		if cname == server {
			return cname, err
		}
		server = cname
	}
	return cname, errors.New("CNAME loop detected")
}

func getCNAMECustomResolver(server string, resolver string) (string, error) {
	var cname string
	var err error
	var d net.Dialer
	ip := net.ParseIP(resolver)
	if ip == nil {
		return cname, errors.New("invalid resolver IP")
	}
	dnsServer := net.JoinHostPort(resolver, "53")
	r := &net.Resolver{
		Dial: func(cont context.Context, network, address string) (net.Conn, error) {
			return d.DialContext(cont, "udp", dnsServer)
		},
	}
	for i := 0; i <= 5; i++ {
		cname, err = r.LookupCNAME(context.Background(), server)
		// Return the requested domain unchanged if the
		// CNAME record can't be resolved. This can
		// happen if there is no CNAME record and no
		// A/AAAA record. But there could still be SRV
		// records.
		if cname == "" {
			return server, nil
		}
		if cname == server {
			return cname, err
		}
		server = cname
	}
	return cname, errors.New("CNAME loop detected")
}

func getIPDoT(server string, resolver string) ([]net.IP, error) {
	var err error
	var d tls.Dialer
	var addr []net.IP
	s := strings.Split(resolver, "#")
	if len(s) != 2 {
		return addr, errors.New("wrong DoT server syntax")
	}
	ip := net.ParseIP(s[0])
	if ip == nil {
		return addr, errors.New("invalid resolver IP")
	}
	dnsServer := net.JoinHostPort(s[0], "853")
	d.Config = &tls.Config{
		ServerName: s[1],
	}
	r := &net.Resolver{
		Dial: func(cont context.Context, network, address string) (net.Conn, error) {
			return d.DialContext(cont, "tcp", dnsServer)
		},
	}
	addr, err = r.LookupIP(context.Background(), "ip", server)
	return addr, err
}

func getIPCustomResolver(server string, resolver string) ([]net.IP, error) {
	var addr []net.IP
	var err error
	var d net.Dialer
	ip := net.ParseIP(resolver)
	if ip == nil {
		return addr, errors.New("invalid resolver IP")
	}
	dnsServer := net.JoinHostPort(resolver, "53")
	r := &net.Resolver{
		Dial: func(cont context.Context, network, address string) (net.Conn, error) {
			return d.DialContext(cont, "tcp", dnsServer)
		},
	}
	addr, err = r.LookupIP(context.Background(), "ip", server)
	return addr, err
}

func getSRVDoT(server string, srvType string, resolver string,
) ([]SRV, error) {
	var records []SRV
	var err error
	var d tls.Dialer
	s := strings.Split(resolver, "#")
	if len(s) != 2 {
		return records, errors.New("wrong DoT server syntax")
	}
	ip := net.ParseIP(s[0])
	if ip == nil {
		return records, errors.New("invalid resolver IP")
	}
	dnsServer := net.JoinHostPort(s[0], "853")
	d.Config = &tls.Config{
		ServerName: s[1],
	}
	r := &net.Resolver{
		Dial: func(cont context.Context, network, address string) (net.Conn, error) {
			return d.DialContext(cont, "tcp", dnsServer)
		},
	}
	_, addr, err := r.LookupSRV(context.Background(), srvType, "tcp", server)
	if err != nil {
		return records, err
	}
	if len(addr) > 0 {
		for _, adr := range addr {
			records = append(records, SRV{
				srvType,
				adr.Target, adr.Port, adr.Priority,
				adr.Weight,
			})
		}
	}
	return records, err
}

func getSRVCustomResolver(server string, srvType string, resolver string,
) ([]SRV, error) {
	var records []SRV
	var err error
	var d net.Dialer
	ip := net.ParseIP(resolver)
	if ip == nil {
		return records, errors.New("invalid resolver IP")
	}
	dnsServer := net.JoinHostPort(resolver, "53")
	r := &net.Resolver{
		Dial: func(cont context.Context, network, address string) (net.Conn, error) {
			return d.DialContext(cont, "udp", dnsServer)
		},
	}
	_, addr, err := r.LookupSRV(context.Background(), srvType, "tcp", server)
	if err != nil {
		return records, err
	}
	if len(addr) > 0 {
		for _, adr := range addr {
			records = append(records, SRV{
				srvType,
				adr.Target, adr.Port, adr.Priority,
				adr.Weight,
			})
		}
	}
	return records, err
}

func getSRV(server string, srvType string) ([]SRV, error) {
	var records []SRV
	var err error

	// Look up SRV records.
	_, addr, err := net.LookupSRV(srvType, "tcp", server)
	if err != nil {
		return records, err
	}
	if len(addr) > 0 {
		for _, adr := range addr {
			records = append(records, SRV{
				srvType,
				adr.Target, adr.Port, adr.Priority,
				adr.Weight,
			})
		}
	}
	return records, err
}

// LookupXmppServer returns the xmpp-server SRV records.
func LookupXmppServer(server string) ([]SRV, error) {
	// Look up xmpp-server SRV
	records, err := getSRV(server, "xmpp-server")
	return records, err
}

// GetIP returns the IP records for a given domain.
func (c *Config) GetIP(server string) ([]net.IP, error) {
	var addr []net.IP
	var err error
	// Look up IPs
	switch {
	case c.Resolver == "":
		addr, err = net.LookupIP(server)
	case c.Resolver != "" && !c.DoT:
		addr, err = getIPCustomResolver(server, c.Resolver)
	case c.Resolver != "" && c.DoT:
		addr, err = getIPDoT(server, c.Resolver)
	default:
		return addr, errors.New("xmppsrv: invalid resolver setting")
	}
	return addr, err
}

// LookupXmppServer returns the xmpp-server SRV records.
func (c *Config) LookupXmppServer(server string) ([]SRV, error) {
	var records []SRV
	var err error
	// Look up xmpp-server SRV
	switch {
	case c.Resolver == "":
		records, err = getSRV(server, "xmpp-server")
	case c.Resolver != "" && !c.DoT:
		records, err = getSRVCustomResolver(server, "xmpp-server", c.Resolver)
	case c.Resolver != "" && c.DoT:
		records, err = getSRVDoT(server, "xmpp-server", c.Resolver)
	default:
		return records, errors.New("xmppsrv: invalid resolver setting")
	}

	return records, err
}

// LookupXmppsServer returns the xmpps-server SRV records.
func LookupXmppsServer(server string) ([]SRV, error) {
	// Look up xmpps-server SRV
	records, err := getSRV(server, "xmpps-server")
	return records, err
}

// LookupCNAME returns the final target name.
func (c *Config) LookupCNAME(server string) (string, error) {
	var cname string
	var err error
	switch {
	case c.Resolver == "":
		// Look up CNAME
		cname, err = getCNAME(server)
	case c.Resolver != "" && !c.DoT:
		// Look up CNAME
		cname, err = getCNAMECustomResolver(server, c.Resolver)
	case c.Resolver != "" && c.DoT:
		// Look up CNMAE
		cname, err = getCNAMEDoT(server, c.Resolver)
	default:
		return cname, errors.New("xmppsrv: invalid resolver setting")
	}
	return cname, err
}

// LookupXmppsServer returns the xmpps-server SRV records.
func (c *Config) LookupXmppsServer(server string) ([]SRV, error) {
	var records []SRV
	var err error
	switch {
	case c.Resolver == "":
		// Look up xmpps-server SRV
		records, err = getSRV(server, "xmpps-server")
	case c.Resolver != "" && !c.DoT:
		// Look up xmpps-server SRV
		records, err = getSRVCustomResolver(server, "xmpps-server",
			c.Resolver)
	case c.Resolver != "" && c.DoT:
		// Look up xmpps-server SRV
		records, err = getSRVDoT(server, "xmpps-server", c.Resolver)
	default:
		return records, errors.New("xmppsrv: invalid resolver setting")
	}
	return records, err
}

// LookupXmppClient returns the xmpp-server SRV records.
func LookupXmppClient(server string) ([]SRV, error) {
	// Look up xmpp-client SRV
	records, err := getSRV(server, "xmpp-client")
	return records, err
}

// LookupXmppClient returns the xmpp-server SRV records.
func (c *Config) LookupXmppClient(server string) ([]SRV, error) {
	var records []SRV
	var err error
	switch {
	case c.Resolver == "":
		// Look up xmpp-client SRV
		records, err = getSRV(server, "xmpp-client")
	case c.Resolver != "" && !c.DoT:
		// Look up xmpp-client SRV
		records, err = getSRVCustomResolver(server, "xmpp-client",
			c.Resolver)
	case c.Resolver != "" && c.DoT:
		// Look up xmpp-client SRV
		records, err = getSRVDoT(server, "xmpp-client", c.Resolver)
	default:
		return records, errors.New("xmppsrv: invalid resolver setting")
	}
	return records, err
}

// LookupXmppsClient returns the xmpp-server SRV records.
func LookupXmppsClient(server string) ([]SRV, error) {
	// Look up xmpps-client SRV
	records, err := getSRV(server, "xmpps-client")
	return records, err
}

// LookupXmppsClient returns the xmpp-server SRV records.
func (c *Config) LookupXmppsClient(server string) ([]SRV, error) {
	var records []SRV
	var err error
	switch {
	case c.Resolver == "":
		// Look up xmpps-client SRV
		records, err = getSRV(server, "xmpps-client")
	case c.Resolver != "" && !c.DoT:
		// Look up xmpps-client SRV
		records, err = getSRVCustomResolver(server, "xmpps-client",
			c.Resolver)
	case c.Resolver != "" && c.DoT:
		// Look up xmpps-client SRV
		records, err = getSRVDoT(server, "xmpps-client", c.Resolver)
	default:
		return records, errors.New("xmppsrv: invalid resolver setting")
	}
	return records, err
}

// LookupClient returns the xmpp-client and xmpps-client SRV records sorted by
// priority and weight.
func LookupClient(server string) ([]SRV, error) {
	var records, records2 []SRV
	var err, err2 error
	// Look up xmpp-client SRV
	records, err = getSRV(server, "xmpp-client")
	// Look up xmpps-client SRV records.
	records2, err2 = getSRV(server, "xmpps-client")
	switch rand.Intn(2) {
	case 0:
		records = append(records, records2...)
	case 1:
		records = append(records2, records...)
	default:
		records = append(records, records2...)
	}
	if err != nil && err2 != nil {
		return records, fmt.Errorf("failure in xmpp-client lookup: %v\n"+
			"failure in xmpps-client lookup: %v", err, err2)
	}
	switch len(records) {
	case 0:
		return records, errors.New("no client records found")
	case 1:
		return records, nil
	default:
		// Sort xmpp- and xmpps-client SRV records according to the priority
		// and weight.
		sort.Sort(byPriority(records))
	}
	return records, nil
}

// LookupClient returns the xmpp-client and xmpps-client SRV records sorted by
// priority and weight.
func (c *Config) LookupClient(server string) ([]SRV, error) {
	var records, records2 []SRV
	var err, err2 error
	switch {
	case c.Resolver == "":
		// Look up xmpp-client SRV
		records, err = getSRV(server, "xmpp-client")
		// Look up xmpps-client SRV records.
		records2, err2 = getSRV(server, "xmpps-client")
	case c.Resolver != "" && !c.DoT:
		// Look up xmpp-client SRV
		records, err = getSRVCustomResolver(server, "xmpp-client", c.Resolver)
		// Look up xmpps-client SRV records.
		records2, err2 = getSRVCustomResolver(server, "xmpps-client", c.Resolver)
	case c.Resolver != "" && c.DoT:
		// Look up xmpp-client SRV
		records, err = getSRVDoT(server, "xmpp-client", c.Resolver)
		// Look up xmpps-client SRV records.
		records2, err2 = getSRVDoT(server, "xmpps-client", c.Resolver)
	default:
		return records, errors.New("xmppsrv: invalid resolver setting")
	}
	switch rand.Intn(2) {
	case 0:
		records = append(records, records2...)
	case 1:
		records = append(records2, records...)
	default:
		records = append(records, records2...)
	}
	if err != nil && err2 != nil {
		return records, fmt.Errorf("failure in xmpp-client lookup: %v\n"+
			"failure in xmpps-client lookup: %v", err, err2)
	}
	switch len(records) {
	case 0:
		return records, errors.New("no client records found")
	case 1:
		return records, nil
	default:
		// Sort xmpp- and xmpps-client SRV records according to the priority
		// and weight.
		sort.Sort(byPriority(records))
	}
	return records, nil
}

// LookupServer returns the xmpp-server and xmpps-server SRV records sorted by
// priority and weight.
func LookupServer(server string) ([]SRV, error) {
	var err, err2 error
	// Look up xmpp-client SRV
	records, err := getSRV(server, "xmpp-server")
	// Look up xmpps-client SRV records.
	records2, err2 := getSRV(server, "xmpps-server")
	switch rand.Intn(2) {
	case 0:
		records = append(records, records2...)
	case 1:
		records = append(records2, records...)
	default:
		records = append(records, records2...)
	}
	if err != nil && err2 != nil {
		return records, fmt.Errorf("failure in xmpp-server lookup: %v\n"+
			"failure in xmpps-server lookup: %v", err, err2)
	}
	switch len(records) {
	case 0:
		return records, errors.New("no server records found")
	case 1:
		return records, nil
	default:
		// Sort xmpp- and xmpps-server SRV records according to the priority
		// and weight.
		sort.Sort(byPriority(records))
	}
	return records, nil
}

// LookupServer returns the xmpp-server and xmpps-server SRV records sorted by
// priority and weight.
func (c *Config) LookupServer(server string) ([]SRV, error) {
	var records, records2 []SRV
	var err, err2 error
	switch {
	case c.Resolver == "":
		// Look up xmpp-client SRV
		records, err = getSRV(server, "xmpp-server")
		// Look up xmpps-client SRV records.
		records2, err2 = getSRV(server, "xmpps-server")
	case c.Resolver != "" && !c.DoT:
		// Look up xmpp-client SRV
		records, err = getSRVCustomResolver(server, "xmpp-server",
			c.Resolver)
		// Look up xmpps-client SRV records.
		records2, err2 = getSRVCustomResolver(server, "xmpps-server",
			c.Resolver)
	case c.Resolver != "" && c.DoT:
		// Look up xmpp-client SRV
		records, err = getSRVDoT(server, "xmpp-server", c.Resolver)
		// Look up xmpps-client SRV records.
		records2, err2 = getSRVDoT(server, "xmpps-server", c.Resolver)
	default:
		return records, errors.New("xmppsrv: invalid resolver setting")
	}
	switch rand.Intn(2) {
	case 0:
		records = append(records, records2...)
	case 1:
		records = append(records2, records...)
	default:
		records = append(records, records2...)
	}
	if err != nil && err2 != nil {
		return records, fmt.Errorf("failure in xmpp-server lookup: %v\n"+
			"failure in xmpps-server lookup: %v", err, err2)
	}
	switch len(records) {
	case 0:
		return records, errors.New("no server records found")
	case 1:
		return records, nil
	default:
		// Sort xmpp- and xmpps-server SRV records according to the priority
		// and weight.
		sort.Sort(byPriority(records))
	}
	return records, nil
}