mtg/stats/prometheus_test.go

package stats_test

import (
	"fmt"
	"io"
	"net"
	"net/http"
	"testing"
	"time"

	"github.com/9seconds/mtg/v2/events"
	"github.com/9seconds/mtg/v2/mtglib"
	"github.com/9seconds/mtg/v2/stats"
	"github.com/stretchr/testify/suite"
)

type PrometheusTestSuite struct {
	suite.Suite

	httpListener net.Listener
	factory      *stats.PrometheusFactory
	prometheus   events.Observer
}

func (suite *PrometheusTestSuite) Get() (string, error) {
	addr := fmt.Sprintf("http://%s/", suite.httpListener.Addr().String())

	resp, err := http.Get(addr) // nolint: noctx
	if err != nil {
		return "", err // nolint: wrapcheck
	}

	defer resp.Body.Close()

	data, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", err // nolint: wrapcheck
	}

	return string(data), nil
}

func (suite *PrometheusTestSuite) SetupTest() {
	suite.httpListener, _ = net.Listen("tcp", "127.0.0.1:0")
	suite.factory = stats.NewPrometheus("mtg", "/")
	suite.prometheus = suite.factory.Make()

	go suite.factory.Serve(suite.httpListener) // nolint: errcheck
}

func (suite *PrometheusTestSuite) TearDownTest() {
	suite.prometheus.Shutdown()
	suite.NoError(suite.factory.Close())
	suite.httpListener.Close()
}

func (suite *PrometheusTestSuite) TestTelegramPath() {
	suite.prometheus.EventStart(
		mtglib.NewEventStart("connID", net.ParseIP("10.0.0.10")))
	time.Sleep(100 * time.Millisecond)

	data, err := suite.Get()
	suite.NoError(err)
	suite.Contains(data, `mtg_client_connections{ip_family="ipv4"} 1`)

	suite.prometheus.EventConnectedToDC(
		mtglib.NewEventConnectedToDC("connID", net.ParseIP("10.0.0.1"), 4))
	time.Sleep(100 * time.Millisecond)

	data, err = suite.Get()
	suite.NoError(err)
	suite.Contains(data, `mtg_telegram_connections{dc="4",telegram_ip="10.0.0.1"} 1`)

	suite.prometheus.EventTraffic(
		mtglib.NewEventTraffic("connID", 200, true))
	time.Sleep(100 * time.Millisecond)

	data, err = suite.Get()
	suite.NoError(err)
	suite.Contains(data, `mtg_telegram_traffic{dc="4",direction="to_client",telegram_ip="10.0.0.1"} 200`)

	suite.prometheus.EventTraffic(
		mtglib.NewEventTraffic("connID", 100, false))
	time.Sleep(100 * time.Millisecond)

	data, err = suite.Get()
	suite.NoError(err)
	suite.Contains(data, `mtg_telegram_traffic{dc="4",direction="from_client",telegram_ip="10.0.0.1"} 100`)

	suite.prometheus.EventFinish(mtglib.NewEventFinish("connID"))
	time.Sleep(100 * time.Millisecond)

	data, err = suite.Get()
	suite.NoError(err)
	suite.Contains(data, `mtg_client_connections{ip_family="ipv4"} 0`)
	suite.Contains(data, `mtg_telegram_connections{dc="4",telegram_ip="10.0.0.1"} 0`)
}

func (suite *PrometheusTestSuite) TestDomainFrontingPath() {
	suite.prometheus.EventStart(
		mtglib.NewEventStart("connID", net.ParseIP("10.0.0.10")))
	time.Sleep(100 * time.Millisecond)

	data, err := suite.Get()
	suite.NoError(err)
	suite.Contains(data, `mtg_client_connections{ip_family="ipv4"} 1`)

	suite.prometheus.EventDomainFronting(mtglib.NewEventDomainFronting("connID"))
	time.Sleep(100 * time.Millisecond)

	data, err = suite.Get()
	suite.NoError(err)
	suite.Contains(data, `mtg_domain_fronting 1`)
	suite.Contains(data, `mtg_domain_fronting_connections{ip_family="ipv4"} 1`)

	suite.prometheus.EventTraffic(
		mtglib.NewEventTraffic("connID", 200, true))
	time.Sleep(100 * time.Millisecond)

	data, err = suite.Get()
	suite.NoError(err)
	suite.Contains(data, `mtg_domain_fronting_traffic{direction="to_client"} 200`)

	suite.prometheus.EventTraffic(
		mtglib.NewEventTraffic("connID", 100, false))
	time.Sleep(100 * time.Millisecond)

	data, err = suite.Get()
	suite.NoError(err)
	suite.Contains(data, `mtg_domain_fronting_traffic{direction="from_client"} 100`)

	suite.prometheus.EventFinish(mtglib.NewEventFinish("connID"))
	time.Sleep(100 * time.Millisecond)

	data, err = suite.Get()
	suite.NoError(err)
	suite.Contains(data, `mtg_client_connections{ip_family="ipv4"} 0`)
	suite.Contains(data, `mtg_domain_fronting_connections{ip_family="ipv4"} 0`)
}

func (suite *PrometheusTestSuite) TestEventConcurrencyLimited() {
	suite.prometheus.EventConcurrencyLimited(mtglib.NewEventConcurrencyLimited())

	time.Sleep(100 * time.Millisecond)

	data, err := suite.Get()
	suite.NoError(err)
	suite.Contains(data, `mtg_concurrency_limited 1`)
}

func (suite *PrometheusTestSuite) TestEventIPBlocklisted() {
	suite.prometheus.EventIPBlocklisted(
		mtglib.NewEventIPBlocklisted(net.ParseIP("2001:db8::68")))

	time.Sleep(100 * time.Millisecond)

	data, err := suite.Get()
	suite.NoError(err)
	suite.Contains(data, `mtg_ip_blocklisted 1`)
}

func (suite *PrometheusTestSuite) TestEventReplayAttack() {
	suite.prometheus.EventReplayAttack(mtglib.NewEventReplayAttack("connID"))

	time.Sleep(100 * time.Millisecond)

	data, err := suite.Get()
	suite.NoError(err)
	suite.Contains(data, `mtg_replay_attacks 1`)
}

func (suite *PrometheusTestSuite) TestEventIPListSize() {
	suite.prometheus.EventIPListSize(mtglib.NewEventIPListSize(10, false))
	suite.prometheus.EventIPListSize(mtglib.NewEventIPListSize(3, true))

	time.Sleep(100 * time.Millisecond)

	data, err := suite.Get()
	suite.NoError(err)
	suite.Contains(data, `mtg_iplist_size{ip_list="allowlist"} 10`)
	suite.Contains(data, `mtg_iplist_size{ip_list="blocklist"} 3`)
}

func TestPrometheus(t *testing.T) {
	t.Parallel()
	suite.Run(t, &PrometheusTestSuite{})
}