internal/check/check.go

// Package check runs connectivity checks against a set of targets.
//
// See mercemay.top/src/portr/ for context.
package check

import (
	"context"
	"net"
	"sort"
	"time"

	"github.com/mercemay/portr/internal/check/rate"
	"github.com/mercemay/portr/internal/check/result"
	"github.com/mercemay/portr/internal/check/retry"
	"github.com/mercemay/portr/internal/check/worker"
	"github.com/mercemay/portr/internal/config"
	"github.com/mercemay/portr/internal/config/target"
)

// Report is the full outcome of a single run.
type Report struct {
	Started  time.Time
	Finished time.Time
	Results  []result.Result
}

// Summary counts open/closed targets.
func (r Report) Summary() (open, closed int) {
	for _, res := range r.Results {
		if res.Open {
			open++
			continue
		}
		closed++
	}
	return
}

// Run probes every target using the given config and returns a Report.
// Context cancellation aborts in-flight workers cleanly.
func Run(ctx context.Context, cfg config.Config) (Report, error) {
	dial := (&net.Dialer{Timeout: cfg.Timeout}).DialContext
	return RunWithDial(ctx, cfg, dial)
}

// RunWithDial is Run with an injectable dialer (for tests).
func RunWithDial(ctx context.Context, cfg config.Config, dial worker.DialFunc) (Report, error) {
	lim := rate.New(cfg.RateLimit)
	bo := retry.New(cfg.Retries, 50*time.Millisecond, cfg.Timeout)
	p := worker.NewPool(cfg.Concurrency, dial, lim, bo)

	started := time.Now()
	out := p.Run(ctx, cfg.Targets)

	results := make([]result.Result, 0, len(cfg.Targets))
	for r := range out {
		results = append(results, r)
	}
	sortResults(results)
	return Report{
		Started:  started,
		Finished: time.Now(),
		Results:  results,
	}, ctx.Err()
}

func sortResults(rs []result.Result) {
	sort.Slice(rs, func(i, j int) bool {
		a, b := rs[i], rs[j]
		if a.Target.Host != b.Target.Host {
			return a.Target.Host < b.Target.Host
		}
		return a.Target.Port < b.Target.Port
	})
}

// FilterByTag returns targets whose tag list contains tag.
func FilterByTag(in []target.Target, tag string) []target.Target {
	if tag == "" {
		return in
	}
	out := in[:0:0]
	for _, t := range in {
		if t.HasTag(tag) {
			out = append(out, t)
		}
	}
	return out
}