feat: implement Prometheus metrics [WIP]

common/options.go

@@ -3,7 +3,10 @@ package common
 import "github.com/charmbracelet/log"
 
 type Options struct {
-	Port        int
+	*Port
+	// NOTE(dwisiswant0): I think it would be fine if we just added
+	// our own metrics route (w/ --metrics flag) for the Prometheus
+	// handler instead of running a new server (--metrics-port)
 	Destination string
 	*Config
 	*TLS

common/port.go

@@ -0,0 +1,5 @@
+package common
+
+type Port struct {
+	Server, Metrics int
+}

common/validator_test.go

@@ -4,7 +4,7 @@ import "testing"
 
 func TestOptions_Validate_ValidConfig(t *testing.T) {
 	opt := &Options{
-		Port:        8080,
+		Port:        &Port{Server: 8080},
 		Destination: "example.com",
 		Config: &Config{
 			Path:   "config.yaml",
@@ -21,7 +21,7 @@ func TestOptions_Validate_ValidConfig(t *testing.T) {
 
 func TestOptions_Validate_EmptyDestination(t *testing.T) {
 	opt := &Options{
-		Port:        8080,
+		Port:        &Port{Server: 8080},
 		Destination: "",
 		Config: &Config{
 			Path:   "config.yaml",
@@ -38,7 +38,7 @@ func TestOptions_Validate_EmptyDestination(t *testing.T) {
 
 func TestOptions_Validate_InvalidConfigFormat(t *testing.T) {
 	opt := &Options{
-		Port:        8080,
+		Port:        &Port{Server: 8080},
 		Destination: "example.com",
 		Config: &Config{
 			Path:   "config.json",
@@ -55,7 +55,7 @@ func TestOptions_Validate_InvalidConfigFormat(t *testing.T) {
 
 func TestOptions_Validate_MissingConfigPathAndFormat(t *testing.T) {
 	opt := &Options{
-		Port:        8080,
+		Port:        &Port{Server: 8080},
 		Destination: "example.com",
 		Config:      &Config{},
 		TLS:         nil,

common/vars.go

@@ -21,6 +21,7 @@ Options:
   -d, --dest <ADDR>:<PORT>     Set the destination address for forwarding requests
   -c, --conf <FILE>            Specify the path to the teler WAF configuration file
   -f, --format <FORMAT>        Specify the configuration file format (json/yaml) (default: yaml)
+      --metrics-port <PORT>    Specify the port for exposing metrics data
       --cert <FILE>            Specify the path to the SSL certificate file
       --key <FILE>             Specify the path to the SSL private key file
   -V, --version                Display the current teler-proxy version

go.mod

@@ -9,6 +9,7 @@ require (
 	github.com/go-co-op/gocron v1.34.0
 	github.com/kitabisa/teler-waf v1.2.2
 	github.com/mattn/go-colorable v0.1.13
+	github.com/prometheus/client_golang v1.17.0
 )
 
 require (
@@ -21,8 +22,10 @@ require (
 	github.com/aws/aws-sdk-go v1.44.298 // indirect
 	github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect
 	github.com/aymerick/douceur v0.2.0 // indirect
+	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/bgentry/go-netrc v0.0.0-20140422174119-9fd32a8b3d3d // indirect
 	github.com/bitfield/script v0.22.0 // indirect
+	github.com/cespare/xxhash/v2 v2.2.0 // indirect
 	github.com/daniel-hutao/spinlock v0.1.0 // indirect
 	github.com/dwisiswant0/clientip v0.3.0 // indirect
 	github.com/gabriel-vasile/mimetype v1.4.2 // indirect
@@ -50,6 +53,7 @@ require (
 	github.com/lucasb-eyer/go-colorful v1.2.0 // indirect
 	github.com/mattn/go-isatty v0.0.18 // indirect
 	github.com/mattn/go-runewidth v0.0.15 // indirect
+	github.com/matttproud/golang_protobuf_extensions v1.0.4 // indirect
 	github.com/microcosm-cc/bluemonday v1.0.25 // indirect
 	github.com/mitchellh/go-homedir v1.1.0 // indirect
 	github.com/mitchellh/go-testing-interface v1.14.1 // indirect
@@ -60,6 +64,9 @@ require (
 	github.com/projectdiscovery/blackrock v0.0.1 // indirect
 	github.com/projectdiscovery/mapcidr v1.1.2 // indirect
 	github.com/projectdiscovery/utils v0.0.41 // indirect
+	github.com/prometheus/client_model v0.4.1-0.20230718164431-9a2bf3000d16 // indirect
+	github.com/prometheus/common v0.44.0 // indirect
+	github.com/prometheus/procfs v0.11.1 // indirect
 	github.com/rivo/uniseg v0.4.4 // indirect
 	github.com/robfig/cron/v3 v3.0.1 // indirect
 	github.com/saintfish/chardet v0.0.0-20230101081208-5e3ef4b5456d // indirect

go.sum

@@ -198,13 +198,17 @@ github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiE
 github.com/aymanbagabas/go-osc52/v2 v2.0.1/go.mod h1:uYgXzlJ7ZpABp8OJ+exZzJJhRNQ2ASbcXHWsFqH8hp8=
 github.com/aymerick/douceur v0.2.0 h1:Mv+mAeH1Q+n9Fr+oyamOlAkUNPWPlA8PPGR0QAaYuPk=
 github.com/aymerick/douceur v0.2.0/go.mod h1:wlT5vV2O3h55X9m7iVYN0TBM0NH/MmbLnd30/FjWUq4=
+github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
+github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw=
 github.com/bgentry/go-netrc v0.0.0-20140422174119-9fd32a8b3d3d h1:xDfNPAt8lFiC1UJrqV3uuy861HCTo708pDMbjHHdCas=
 github.com/bgentry/go-netrc v0.0.0-20140422174119-9fd32a8b3d3d/go.mod h1:6QX/PXZ00z/TKoufEY6K/a0k6AhaJrQKdFe6OfVXsa4=
 github.com/bitfield/script v0.22.0 h1:LA7QHuEsXMPD52YLtxWrlqCCy+9FOpzNYfsRHC5Gsrc=
 github.com/bitfield/script v0.22.0/go.mod h1:ms4w+9B8f2/W0mbsgWDVTtl7K94bYuZc3AunnJC4Ebs=
 github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU=
 github.com/cespare/xxhash v1.1.0/go.mod h1:XrSqR1VqqWfGrhpAt58auRo0WTKS1nRRg3ghfAqPWnc=
 github.com/cespare/xxhash/v2 v2.1.1/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
+github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44=
+github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
 github.com/charmbracelet/lipgloss v0.9.1 h1:PNyd3jvaJbg4jRHKWXnCj1akQm4rh8dbEzN1p/u1KWg=
 github.com/charmbracelet/lipgloss v0.9.1/go.mod h1:1mPmG4cxScwUQALAAnacHaigiiHB9Pmr+v1VEawJl6I=
 github.com/charmbracelet/log v0.2.4 h1:3pKtq5/Y5QMKtcZt7kDqD1p9w7lICzHYQACBFY4ocHA=
@@ -419,6 +423,8 @@ github.com/mattn/go-runewidth v0.0.12/go.mod h1:RAqKPSqVFrSLVXbA8x7dzmKdmGzieGRC
 github.com/mattn/go-runewidth v0.0.14/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w=
 github.com/mattn/go-runewidth v0.0.15 h1:UNAjwbU9l54TA3KzvqLGxwWjHmMgBUVhBiTjelZgg3U=
 github.com/mattn/go-runewidth v0.0.15/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w=
+github.com/matttproud/golang_protobuf_extensions v1.0.4 h1:mmDVorXM7PCGKw94cs5zkfA9PSy5pEvNWRP0ET0TIVo=
+github.com/matttproud/golang_protobuf_extensions v1.0.4/go.mod h1:BSXmuO+STAnVfrANrmjBb36TMTDstsz7MSK+HVaYKv4=
 github.com/microcosm-cc/bluemonday v1.0.25 h1:4NEwSfiJ+Wva0VxN5B8OwMicaJvD8r9tlJWm9rtloEg=
 github.com/microcosm-cc/bluemonday v1.0.25/go.mod h1:ZIOjCQp1OrzBBPIJmfX4qDYFuhU02nx4bn030ixfHLE=
 github.com/mitchellh/go-homedir v1.1.0 h1:lukF9ziXFxDFPkA1vsr5zpc1XuPDn/wFntq5mG+4E0Y=
@@ -442,7 +448,15 @@ github.com/projectdiscovery/mapcidr v1.1.2 h1:Mmq/nPqvVc7fjvH/kJVK0IBOny/LrJIxZ4
 github.com/projectdiscovery/mapcidr v1.1.2/go.mod h1:Aoq0x/wJl6KDbtQ8OcPkjIDCqx2iEyx5ty1nzso8wXM=
 github.com/projectdiscovery/utils v0.0.41 h1:GY0/EhZR2DYne20eS8KApzcQpYJ5rWw8bSFYBDvJH5Q=
 github.com/projectdiscovery/utils v0.0.41/go.mod h1:e+67VwXkS9o+dnRU0jM1BRdjm5GvNwkG/B2yhpZh42s=
+github.com/prometheus/client_golang v1.17.0 h1:rl2sfwZMtSthVU752MqfjQozy7blglC+1SOtjMAMh+Q=
+github.com/prometheus/client_golang v1.17.0/go.mod h1:VeL+gMmOAxkS2IqfCq0ZmHSL+LjWfWDUmp1mBz9JgUY=
 github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=
+github.com/prometheus/client_model v0.4.1-0.20230718164431-9a2bf3000d16 h1:v7DLqVdK4VrYkVD5diGdl4sxJurKJEMnODWRJlxV9oM=
+github.com/prometheus/client_model v0.4.1-0.20230718164431-9a2bf3000d16/go.mod h1:oMQmHW1/JoDwqLtg57MGgP/Fb1CJEYF2imWWhWtMkYU=
+github.com/prometheus/common v0.44.0 h1:+5BrQJwiBB9xsMygAB3TNvpQKOwlkc25LbISbrdOOfY=
+github.com/prometheus/common v0.44.0/go.mod h1:ofAIvZbQ1e/nugmZGz4/qCb9Ap1VoSTIO7x0VV9VvuY=
+github.com/prometheus/procfs v0.11.1 h1:xRC8Iq1yyca5ypa9n1EZnWZkt7dwcoRPQwX/5gwaUuI=
+github.com/prometheus/procfs v0.11.1/go.mod h1:eesXgaPo1q7lBpVMoMy0ZOFTth9hBn4W/y0/p/ScXhY=
 github.com/rivo/uniseg v0.1.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc=
 github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc=
 github.com/rivo/uniseg v0.4.4 h1:8TfxU8dW6PdqD27gjM8MVNuicgxIjxpm4K7x4jp8sis=
@@ -1031,8 +1045,8 @@ gopkg.in/cheggaaa/pb.v1 v1.0.27/go.mod h1:V/YB90LKu/1FcN3WVnfiiE5oMCibMjukxqG/qS
 gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI=
 gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
 gopkg.in/yaml.v2 v2.2.3/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
-gopkg.in/yaml.v2 v2.2.8 h1:obN1ZagJSUGI0Ek/LBmuj4SNLPfIny3KsKFopxRdj10=
 gopkg.in/yaml.v2 v2.2.8/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
+gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

internal/metrics/metrics.go

@@ -0,0 +1,44 @@
+package metrics
+
+import (
+	"github.com/kitabisa/teler-waf/threat"
+	"github.com/prometheus/client_golang/prometheus"
+)
+
+type Metrics struct {
+	Threats  *prometheus.GaugeVec
+	Events   *prometheus.CounterVec
+	Duration *prometheus.HistogramVec
+}
+
+func New(registry prometheus.Registerer) *Metrics {
+	m := new(Metrics)
+
+	var threats []string
+	for _, t := range threat.List() {
+		threats = append(threats, t.String())
+	}
+
+	m.Threats = prometheus.NewGaugeVec(prometheus.GaugeOpts{
+		Namespace: namespace,
+		Name:      "threat_datasets",
+		Help:      "Number of avaiable threat datasets.",
+	}, threats)
+
+	m.Events = prometheus.NewCounterVec(prometheus.CounterOpts{
+		Namespace: namespace,
+		Name:      "teler_event",
+		Help:      "Number of incoming teler WAF events.",
+	}, []string{"rule", "threat"})
+
+	m.Duration = prometheus.NewHistogramVec(prometheus.HistogramOpts{
+		Namespace: namespace,
+		Name:      "request_duration_seconds",
+		Help:      "Request duration times in seconds.",
+		Buckets:   []float64{0.1, 0.15, 0.2, 0.25, 0.3},
+	}, []string{"status", "method", "path"})
+
+	registry.MustRegister(m.Threats, m.Events, m.Duration)
+
+	return m
+}

internal/metrics/vars.go

@@ -0,0 +1,5 @@
+package metrics
+
+import "github.com/kitabisa/teler-proxy/common"
+
+var namespace = common.App

internal/runner/const.go

@@ -1,5 +1,6 @@
 package runner
 
 const (
-	errSomething = "Something went wrong"
+	errSomething   = "Something went wrong"
+	errInitMetrics = "Error while metrics initialization"
 )

internal/runner/options.go

@@ -9,21 +9,24 @@ import (
 )
 
 func ParseOptions() *common.Options {
-	opt := &common.Options{}
-	cfg := &common.Config{}
-	tls := &common.TLS{}
+	opt := new(common.Options)
+	cfg := new(common.Config)
+	tls := new(common.TLS)
+	port := new(common.Port)
 
 	opt.Logger = logger.New()
 
-	flag.IntVar(&opt.Port, "p", 1337, "")
-	flag.IntVar(&opt.Port, "port", 1337, "")
+	flag.IntVar(&port.Server, "p", 1337, "")
+	flag.IntVar(&port.Server, "port", 1337, "")
 
 	flag.StringVar(&opt.Destination, "d", "", "")
 	flag.StringVar(&opt.Destination, "dest", "", "")
 
 	flag.StringVar(&cfg.Path, "c", "", "")
 	flag.StringVar(&cfg.Path, "conf", "", "")
 
+	flag.IntVar(&port.Metrics, "metrics-port", 0, "")
+
 	flag.StringVar(&cfg.Format, "f", "yaml", "")
 	flag.StringVar(&cfg.Format, "format", "yaml", "")
 
@@ -53,6 +56,7 @@ func ParseOptions() *common.Options {
 
 	opt.Config = cfg
 	opt.TLS = tls
+	opt.Port = port
 
 	return opt
 }

internal/runner/prometheus.go

@@ -0,0 +1,75 @@
+package runner
+
+import (
+	"strconv"
+	"time"
+
+	"net/http"
+
+	"github.com/kitabisa/teler-proxy/internal/metrics"
+	"github.com/kitabisa/teler-waf/threat"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/collectors"
+)
+
+type Prometheus struct {
+	*metrics.Metrics
+	*prometheus.Registry
+}
+
+type customResponseWriter struct {
+	http.ResponseWriter
+	statusCode int
+}
+
+func (r *Runner) initMetrics() error {
+	reg := prometheus.NewRegistry()
+	reg.MustRegister(collectors.NewGoCollector())
+
+	m := metrics.New(reg)
+
+	threatLabels := make(prometheus.Labels)
+	for _, t := range threat.List() {
+		count, err := t.Count()
+		if err != nil {
+			return err
+		}
+
+		threatLabels[t.String()] = strconv.Itoa(count)
+	}
+
+	m.Threats.With(threatLabels).Set(1)
+
+	r.Prometheus.Metrics = m
+	r.Prometheus.Registry = reg
+
+	return nil
+}
+
+func (run *Runner) promMiddleware(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
+		now := time.Now()
+
+		customResponseWriter := &customResponseWriter{ResponseWriter: w}
+		next.ServeHTTP(customResponseWriter, req)
+
+		run.Prometheus.Metrics.Duration.With(prometheus.Labels{
+			"method": req.Method,
+			"status": customResponseWriter.StatusString(),
+			"path":   req.URL.Path,
+		}).Observe(time.Since(now).Seconds())
+	})
+}
+
+func (w *customResponseWriter) WriteHeader(code int) {
+	w.statusCode = code
+	w.ResponseWriter.WriteHeader(code)
+}
+
+func (w *customResponseWriter) Status() int {
+	return w.statusCode
+}
+
+func (w *customResponseWriter) StatusString() string {
+	return strconv.Itoa(w.Status())
+}

internal/runner/runner.go

@@ -24,6 +24,7 @@ type Runner struct {
 	*common.Options
 	*cron.Cron
 	*http.Server
+	*Prometheus
 
 	shutdown  shutdown
 	telerOpts teler.Options
@@ -54,21 +55,41 @@ func New(opt *common.Options) error {
 		run.watcher.config = new(fsnotify.Watcher)
 	}
 
+	if opt.Port.Metrics > 0 {
+		if err := run.initMetrics(); err != nil {
+			opt.Logger.Error(errInitMetrics, "err", err)
+		}
+	}
+
 	dest := buildDest(opt.Destination)
 	writer := writer.New()
+	if run.Prometheus.Metrics != nil {
+		writer.Metrics = run.Prometheus.Metrics
+	}
 
 	tun, err := run.createTunnel(dest, writer)
 	if err != nil {
 		return err
 	}
 
+	handler := tun.ReverseProxy
+	// TODO(dwisisant0):
+	//   1. wrap Tunnel server
+	//   2. implement Prometheus.Registry in single metrics route
+	//   3. implement promMiddleware
+	// if run.Prometheus.Metrics != nil {
+	// 	handler = http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
+	// 		promMiddleware(tun).ServeHTTP(w, req)
+	// 	})
+	// }
+
 	logger := log.StandardLog(log.StandardLogOptions{
 		ForceLevel: log.ErrorLevel,
 	})
 
 	server := &http.Server{
-		Addr:     fmt.Sprintf(":%d", opt.Port),
-		Handler:  tun,
+		Addr:     fmt.Sprintf(":%d", opt.Port.Server),
+		Handler:  handler,
 		ErrorLog: logger,
 	}
 
@@ -128,7 +149,7 @@ func (r *Runner) start() error {
 
 	r.Options.Logger.Info(
 		"Server started!",
-		"port", r.Options.Port,
+		"port", r.Options.Port.Server,
 		"tls", tls,
 		"pid", os.Getpid(),
 	)

internal/runner/utils.go

@@ -74,5 +74,10 @@ func (r *Runner) shouldCron() bool {
 func (r *Runner) createTunnel(dest string, writer io.Writer) (*tunnel.Tunnel, error) {
 	opt := r.Options
 
-	return tunnel.NewTunnel(opt.Port, dest, opt.Config.Path, opt.Config.Format, writer)
+	return tunnel.NewTunnel(
+		opt.Port.Server,
+		dest, opt.Config.Path,
+		opt.Config.Format,
+		writer,
+	)
 }