// Package monitor implements codit's external host/service monitoring: a
// background engine that checks registered targets on an interval and records
// heartbeats. See docs/monitoring-design.md.
package monitor

import "context"
import "errors"
import "fmt"
import "net"
import "time"

// ErrBlockedEgress is returned when a target resolves to an address the egress
// policy forbids (loopback, private, link-local, metadata, etc.).
var ErrBlockedEgress = errors.New("target address is not permitted by the egress policy")

// EgressPolicy decides whether the monitor engine may connect to a resolved IP.
// The default policy denies internal ranges so a user-registered target cannot
// be used to probe codit's own network (SSRF). AllowInternal disables that.
type EgressPolicy struct {
	AllowInternal bool
}

func DefaultEgressPolicy() EgressPolicy {
	// denies internal ranges.
	//return EgressPolicy{AllowInternal: false}

	// allows internal ranges.
	return EgressPolicy{AllowInternal: true}
}

// IPAllowed reports whether the policy permits connecting to ip.
func (p EgressPolicy) IPAllowed(ip net.IP) bool {
	if ip == nil {
		return false
	}
	if p.AllowInternal {
		return true
	}
	if ip.IsLoopback() || ip.IsUnspecified() || ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() || ip.IsMulticast() || ip.IsInterfaceLocalMulticast() {
		return false
	}
	if ip.IsPrivate() {
		return false
	}
	return true
}

// resolveAllowedIP resolves host and returns the first IP the policy permits,
// pinning the connection to that IP (which guards against DNS-rebinding: the
// address actually dialed is the one validated). It returns ErrBlockedEgress if
// the host resolves only to forbidden addresses.
func (p EgressPolicy) resolveAllowedIP(ctx context.Context, host string) (net.IP, error) {
	var ips []net.IP
	var ip net.IP
	var err error

	// A literal IP needs no DNS lookup.
	ip = net.ParseIP(host)
	if ip != nil {
		if !p.IPAllowed(ip) {
			return nil, ErrBlockedEgress
		}
		return ip, nil
	}

	ips, err = net.DefaultResolver.LookupIP(ctx, "ip", host)
	if err != nil {
		return nil, err
	}
	for _, ip = range ips {
		if p.IPAllowed(ip) {
			return ip, nil
		}
	}
	return nil, ErrBlockedEgress
}

// ResolveValidated resolves host and returns every resolved IP (for display)
// plus the first IP the policy permits (for dialing). A literal IP resolves to
// itself. Returns ErrBlockedEgress if no resolved address is permitted.
func (p EgressPolicy) ResolveValidated(ctx context.Context, host string) ([]net.IP, net.IP, error) {
	var ips []net.IP
	var ip net.IP
	var chosen net.IP
	var err error

	ip = net.ParseIP(host)
	if ip != nil {
		if !p.IPAllowed(ip) {
			return []net.IP{ip}, nil, ErrBlockedEgress
		}
		return []net.IP{ip}, ip, nil
	}
	ips, err = net.DefaultResolver.LookupIP(ctx, "ip", host)
	if err != nil {
		return nil, nil, err
	}
	for _, ip = range ips {
		if chosen == nil && p.IPAllowed(ip) {
			chosen = ip
		}
	}
	if chosen == nil {
		return ips, nil, ErrBlockedEgress
	}
	return ips, chosen, nil
}

// DialContext is a net.Dialer-compatible dial function that enforces the policy:
// it resolves the host, picks a permitted IP, and dials that IP directly. Used
// as the Transport.DialContext for HTTP/TCP checks.
func (p EgressPolicy) DialContext(ctx context.Context, network string, address string) (net.Conn, error) {
	var host string
	var port string
	var ip net.IP
	var dialer net.Dialer
	var err error

	host, port, err = net.SplitHostPort(address)
	if err != nil {
		return nil, err
	}
	ip, err = p.resolveAllowedIP(ctx, host)
	if err != nil {
		return nil, err
	}
	dialer = net.Dialer{Timeout: 10 * time.Second}
	return dialer.DialContext(ctx, network, net.JoinHostPort(ip.String(), port))
}

// CheckHostPort validates that host (and its resolved IP) is permitted; used by
// checkers that dial directly (TLS, TCP, ping) rather than through an
// http.Transport. Returns the permitted IP to dial.
func (p EgressPolicy) CheckHostPort(ctx context.Context, host string) (net.IP, error) {
	if host == "" {
		return nil, fmt.Errorf("empty host")
	}
	return p.resolveAllowedIP(ctx, host)
}