Merge pull request #9 from d-Rickyy-b/prometheus

Implement prometheus interface

Author: d-Rickyy-b

cmd/main.go

@@ -4,6 +4,7 @@ import (
 	"flag"
 	"go-certstream-server/internal/certificatetransparency"
 	"go-certstream-server/internal/config"
+	"go-certstream-server/internal/prometheus"
 	"go-certstream-server/internal/web"
 	"log"
 )
@@ -22,6 +23,20 @@ func main() {
 	}
 
 	webserver := web.NewWebsocketServer(conf.Webserver.ListenAddr, conf.Webserver.ListenPort, conf.Webserver.CertPath, conf.Webserver.CertKeyPath)
+	if conf.Prometheus.Enabled {
+		// If prometheus is enabled, and interface is either unconfigured or same as webserver config, use existing webserver
+		if (conf.Prometheus.ListenAddr == "" || conf.Prometheus.ListenAddr == conf.Webserver.ListenAddr) &&
+			(conf.Prometheus.ListenPort == 0 || conf.Prometheus.ListenPort == conf.Webserver.ListenPort) {
+			log.Println("Starting prometheus server on same interface as webserver")
+			webserver.RegisterPrometheus(conf.Prometheus.MetricsURL, prometheus.WritePrometheus)
+		} else {
+			log.Println("Starting prometheus server on new interface")
+			metricsServer := web.NewMetricsServer(conf.Prometheus.ListenAddr, conf.Prometheus.ListenPort, conf.Webserver.CertPath, conf.Webserver.CertKeyPath)
+			metricsServer.RegisterPrometheus(conf.Prometheus.MetricsURL, prometheus.WritePrometheus)
+			go metricsServer.Start()
+		}
+	}
+
 	go webserver.Start()
 
 	watcher := certificatetransparency.Watcher{}

config.sample.yaml

@@ -6,3 +6,9 @@ webserver:
   domains_only_url: "/domains-only"
   cert_path: ""
   cert_key_path: ""
+
+prometheus:
+  enabled: true
+  listen_addr: "127.0.0.1"
+  listen_port: 8080
+  metrics_url: "/metrics"

go.mod

@@ -3,6 +3,7 @@ module go-certstream-server
 go 1.18
 
 require (
+	github.com/VictoriaMetrics/metrics v1.22.2
 	github.com/go-chi/chi v1.5.4
 	github.com/google/certificate-transparency-go v1.1.3
 	github.com/gorilla/websocket v1.5.0
@@ -78,18 +79,12 @@ require (
 	go.etcd.io/etcd/server/v3 v3.5.5 // indirect
 	go.etcd.io/etcd/tests/v3 v3.5.5 // indirect
 	go.etcd.io/etcd/v3 v3.5.5 // indirect
-	go.opentelemetry.io/contrib v1.11.0 // indirect
 	go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.36.3 // indirect
 	go.opentelemetry.io/otel v1.11.0 // indirect
-	go.opentelemetry.io/otel/exporters/otlp v0.20.0 // indirect
 	go.opentelemetry.io/otel/exporters/otlp/internal/retry v1.11.0 // indirect
 	go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.11.0 // indirect
 	go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.11.0 // indirect
-	go.opentelemetry.io/otel/internal/metric v0.27.0 // indirect
-	go.opentelemetry.io/otel/metric v0.32.3 // indirect
 	go.opentelemetry.io/otel/sdk v1.11.0 // indirect
-	go.opentelemetry.io/otel/sdk/export/metric v0.28.0 // indirect
-	go.opentelemetry.io/otel/sdk/metric v0.32.3 // indirect
 	go.opentelemetry.io/otel/trace v1.11.0 // indirect
 	go.opentelemetry.io/proto/otlp v0.19.0 // indirect
 	go.uber.org/atomic v1.10.0 // indirect
@@ -103,7 +98,6 @@ require (
 	golang.org/x/text v0.4.0 // indirect
 	golang.org/x/time v0.1.0 // indirect
 	golang.org/x/tools v0.1.12 // indirect
-	golang.org/x/xerrors v0.0.0-20220907171357-04be3eba64a2 // indirect
 	google.golang.org/appengine v1.6.7 // indirect
 	google.golang.org/genproto v0.0.0-20221018160656-63c7b68cfc55 // indirect
 	google.golang.org/grpc v1.50.1 // indirect

go.sum

@@ -17,8 +17,44 @@ import (
 	"net/http"
 	"strings"
 	"sync"
+	"sync/atomic"
 )
 
+var (
+	processedCerts    int64
+	processedPrecerts int64
+	urlMapMutex       sync.RWMutex
+	urlMap            = make(map[string]int64)
+)
+
+func GetProcessedCerts() int64 {
+	return processedCerts
+}
+
+func GetProcessedPrecerts() int64 {
+	return processedPrecerts
+}
+
+func GetCertCountForLog(logname string) int64 {
+	urlMapMutex.RLock()
+	defer urlMapMutex.RUnlock()
+	return urlMap[logname]
+}
+
+func GetLogs() []string {
+	urlMapMutex.RLock()
+	defer urlMapMutex.RUnlock()
+
+	urls := make([]string, len(urlMap))
+
+	counter := 0
+	for key := range urlMap {
+		urls[counter] = key
+		counter++
+	}
+	return urls
+}
+
 // Watcher describes a component that watches for new certificates in a CT log.
 type Watcher struct {
 	Name       string
@@ -132,6 +168,7 @@ func (w *worker) foundCertCallback(rawEntry *ct.RawLogEntry) {
 	}
 	entry.Data.UpdateType = "X509LogEntry"
 	w.entryChan <- entry
+	atomic.AddInt64(&processedCerts, 1)
 }
 
 // foundPrecertCallback is the callback that handles cases where new precerts are found.
@@ -143,9 +180,10 @@ func (w *worker) foundPrecertCallback(rawEntry *ct.RawLogEntry) {
 	}
 	entry.Data.UpdateType = "PrecertLogEntry"
 	w.entryChan <- entry
+	atomic.AddInt64(&processedPrecerts, 1)
 }
 
-// certHandler takes the entries out of the channel and broadcasts them to all clients.
+// certHandler takes the entries out of the entryChan channel and broadcasts them to all clients.
 func certHandler(entryChan chan certstream.Entry) {
 	var processed int64
 	for {
@@ -160,6 +198,12 @@ func certHandler(entryChan chan certstream.Entry) {
 
 		// Run json encoding in the background and send the result to the clients.
 		web.ClientHandler.Broadcast <- entry
+
+		url := normalizeCtlogURL(entry.Data.Source.URL)
+
+		urlMapMutex.Lock()
+		urlMap[url] += 1
+		urlMapMutex.Unlock()
 	}
 }
 
@@ -199,3 +243,10 @@ func getAllLogs() (loglist3.LogList, error) {
 
 	return *allLogs, nil
 }
+
+func normalizeCtlogURL(input string) string {
+	input = strings.TrimPrefix(input, "http://")
+	input = strings.TrimPrefix(input, "https://")
+	input = strings.TrimSuffix(input, "/")
+	return input
+}

internal/certificatetransparency/ct-watcher.go

@@ -20,6 +20,12 @@ type Config struct {
 		CertPath       string `yaml:"cert_path"`
 		CertKeyPath    string `yaml:"cert_key_path"`
 	}
+	Prometheus struct {
+		Enabled    bool   `yaml:"enabled"`
+		MetricsURL string `yaml:"metrics_url"`
+		ListenAddr string `yaml:"listen_addr"`
+		ListenPort int    `yaml:"listen_port"`
+	}
 }
 
 // ReadConfig reads the config file and returns a filled Config struct.
@@ -100,5 +106,15 @@ func validateConfig(config Config) bool {
 		config.Webserver.FullURL = "/domains-only"
 	}
 
+	if config.Prometheus.Enabled {
+		if config.Prometheus.ListenAddr == "" || !IPRegex.MatchString(config.Prometheus.ListenAddr) {
+			log.Fatalln("Prometheus export IP does not match pattern 'x.x.x.x'")
+			return false
+		}
+		if config.Prometheus.ListenPort == 0 {
+			log.Fatalln("Prometheus export port is not set")
+			return false
+		}
+	}
 	return true
 }

internal/config/config.go

@@ -0,0 +1,45 @@
+package prometheus
+
+import (
+	"fmt"
+	"github.com/VictoriaMetrics/metrics"
+	"go-certstream-server/internal/certificatetransparency"
+	"go-certstream-server/internal/web"
+	"io"
+)
+
+var (
+	ctLogsInitialized = false
+	fullClientCount   = metrics.NewGauge("certstreamservergo_clients_total{type=\"full\"}", func() float64 {
+		return float64(web.ClientHandler.ClientFullCount())
+	})
+	liteClientCount = metrics.NewGauge("certstreamservergo_clients_total{type=\"lite\"}", func() float64 {
+		return float64(web.ClientHandler.ClientLiteCount())
+	})
+	domainClientCount = metrics.NewGauge("certstreamservergo_clients_total{type=\"domain\"}", func() float64 {
+		return float64(web.ClientHandler.ClientDomainsCount())
+	})
+	processedCertificates = metrics.NewGauge("certstreamservergo_certificates_total{type=\"regular\"}", func() float64 {
+		return float64(certificatetransparency.GetProcessedCerts())
+	})
+	processedPreCertificates = metrics.NewGauge("certstreamservergo_certificates_total{type=\"precert\"}", func() float64 {
+		return float64(certificatetransparency.GetProcessedPrecerts())
+	})
+)
+
+// WritePrometheus provides an easy way to write metrics to a writer.
+func WritePrometheus(w io.Writer, exposeProcessMetrics bool) {
+	if !ctLogsInitialized {
+		logs := certificatetransparency.GetLogs()
+		for i := 0; i < len(logs); i++ {
+			url := logs[i]
+			metrics.NewGauge(fmt.Sprintf("certstreamservergo_certs_by_log_total{url=\"%s\"}", url), func() float64 {
+				return float64(certificatetransparency.GetCertCountForLog(url))
+			})
+		}
+		if len(logs) > 0 {
+			ctLogsInitialized = true
+		}
+	}
+	metrics.WritePrometheus(w, exposeProcessMetrics)
+}

internal/prometheus/prometheus.go

@@ -38,6 +38,33 @@ func (bm *BroadcastManager) unregisterClient(c *client) {
 	bm.clientLock.Unlock()
 }
 
+// ClientFullCount returns the current number of clients connected to the service on the `full` endpoint.
+func (bm *BroadcastManager) ClientFullCount() (count int64) {
+	return bm.clientCountByType(SubTypeFull)
+}
+
+// ClientLiteCount returns the current number of clients connected to the service on the `lite` endpoint.
+func (bm *BroadcastManager) ClientLiteCount() (count int64) {
+	return bm.clientCountByType(SubTypeLite)
+}
+
+// ClientDomainsCount returns the current number of clients connected to the service on the `domains-only` endpoint.
+func (bm *BroadcastManager) ClientDomainsCount() (count int64) {
+	return bm.clientCountByType(SubTypeDomain)
+}
+
+// clientCountByType returns the current number of clients connected to the service on the endpoint matching the specified SubscriptionType.
+func (bm *BroadcastManager) clientCountByType(subType SubscriptionType) (count int64) {
+	bm.clientLock.RLock()
+	defer bm.clientLock.RUnlock()
+	for _, c := range bm.clients {
+		if c.subType == subType {
+			count++
+		}
+	}
+	return count
+}
+
 // broadcaster is run in a goroutine and handles the dispatching of entries to clients.
 func (bm *BroadcastManager) broadcaster() {
 	for {

internal/web/broadcastmanager.go

@@ -8,6 +8,7 @@ import (
 	"github.com/gorilla/websocket"
 	"go-certstream-server/internal/certstream"
 	"go-certstream-server/internal/config"
+	"io"
 	"log"
 	"net/http"
 	"time"
@@ -25,6 +26,15 @@ type WebServer struct {
 	keyPath   string
 }
 
+// RegisterPrometheus registers a new handler that listens on the given url and calls the given function
+// in order to provide metrics for a prometheus server. This function signature was used, because VictoriaMetrics
+// offers exactly this function signature.
+func (ws *WebServer) RegisterPrometheus(url string, callback func(w io.Writer, exposeProcessMetrics bool)) {
+	ws.routes.HandleFunc(url, func(w http.ResponseWriter, r *http.Request) {
+		callback(w, false)
+	})
+}
+
 // initFullWebsocket is called when a client connects to the /full-stream endpoint.
 // It upgrades the connection to a websocket and starts a goroutine to listen for messages from the client.
 func initFullWebsocket(w http.ResponseWriter, r *http.Request) {
@@ -138,6 +148,20 @@ func (ws *WebServer) initServer() {
 	}
 }
 
+// NewMetricsServer creates a new webserver that listens on the given port and provides metrics for a prometheus server.
+func NewMetricsServer(networkIf string, port int, certPath, keyPath string) *WebServer {
+	server := &WebServer{
+		networkIf: networkIf,
+		port:      port,
+		routes:    chi.NewRouter(),
+		certPath:  certPath,
+		keyPath:   keyPath,
+	}
+	server.initServer()
+	server.routes.Use(middleware.Recoverer)
+	return server
+}
+
 // NewWebsocketServer starts a new webserver and initialized it with the necessary routes.
 // It also starts the broadcaster in ClientHandler as a background job.
 func NewWebsocketServer(networkIf string, port int, certPath, keyPath string) *WebServer {