Add Kafka notifier

Author: Daniel Jimenez

cmd/vulnerability-db-consumer/config.go

@@ -19,6 +19,7 @@ type config struct {
 	DB          dbConfig
 	SQS         sqsConfig
 	SNS         snsConfig
+	Kafka       kafkaConfig
 	Report      reportConfig
 	Maintenance maintenanceConfig
 }
@@ -46,11 +47,19 @@ type sqsConfig struct {
 }
 
 type snsConfig struct {
-	TopicARN string `toml:"topic_arn"`
 	Enabled  bool
+	TopicARN string `toml:"topic_arn"`
 	Endpoint string `toml:"endpoint"`
 }
 
+type kafkaConfig struct {
+	Enabled   bool   `toml:"enabled"`
+	User      string `toml:"user"`
+	Pass      string `toml:"password"`
+	BrokerURL string `toml:"broker_url"`
+	Topic     string `toml:"topic"`
+}
+
 type reportConfig struct {
 	URLReplace string `toml:"url_replace"`
 }

cmd/vulnerability-db-consumer/main.go

@@ -11,6 +11,7 @@ import (
 	"os"
 	"sync"
 
+	"github.com/adevinta/vulnerability-db/pkg/asyncapi/kafka"
 	"github.com/adevinta/vulnerability-db/pkg/maintenance"
 	"github.com/adevinta/vulnerability-db/pkg/notify"
 	"github.com/adevinta/vulnerability-db/pkg/processor"
@@ -43,14 +44,9 @@ func main() {
 	}
 
 	// Build notifier.
-	snsConf := notify.SNSConfig{
-		TopicArn: conf.SNS.TopicARN,
-		Enabled:  conf.SNS.Enabled,
-		Endpoint: conf.SNS.Endpoint,
-	}
-	snsNotifier, err := notify.NewSNSNotifier(snsConf, logger)
+	notifier, err := buildNotifier(conf, logger)
 	if err != nil {
-		log.Fatalf("Error creating notifier: %v", err)
+		log.Fatalf("Error building notifier: %v", err)
 	}
 
 	// Build processor.
@@ -59,7 +55,7 @@ func main() {
 		log.Fatalf("Error creating results client: %v", err)
 	}
 
-	processor, err := processor.NewCheckProcessor(snsNotifier, db, resultsClient, conf.Report.URLReplace, conf.MaxEventAge, logger)
+	processor, err := processor.NewCheckProcessor(notifier, db, resultsClient, conf.Report.URLReplace, conf.MaxEventAge, logger)
 	if err != nil {
 		log.Fatalf("Error creating queue processor: %v", err)
 	}
@@ -97,6 +93,48 @@ func main() {
 	wg.Wait()
 }
 
+// buildNotifier builds the appropiate notifier given the defined configuration.
+// TODO: Once the integrations dependent on the old notification format have been
+// deprecated or updated to comply with the new format through Kafka topic channel
+// we can get rid of SNS and multi implementations of notifier and just use Kafka.
+func buildNotifier(conf *config, logger *log.Logger) (notify.Notifier, error) {
+	if !conf.SNS.Enabled && !conf.Kafka.Enabled {
+		return notify.NewNoopNotifier(), nil
+	}
+	if conf.SNS.Enabled && !conf.Kafka.Enabled {
+		return buildSNSNotifier(conf, logger)
+	}
+	if !conf.SNS.Enabled && conf.Kafka.Enabled {
+		return buildKafkaNotifier(conf, logger)
+	}
+	// Multi Notifier
+	k, err := buildKafkaNotifier(conf, logger)
+	if err != nil {
+		return nil, err
+	}
+	s, err := buildSNSNotifier(conf, logger)
+	if err != nil {
+		return nil, err
+	}
+	return notify.NewMultiNotifier(k, s), nil
+}
+
+func buildSNSNotifier(conf *config, logger *log.Logger) (*notify.SNSNotifier, error) {
+	return notify.NewSNSNotifier(notify.SNSConfig{
+		TopicArn: conf.SNS.TopicARN,
+		Endpoint: conf.SNS.Endpoint,
+	}, logger)
+}
+
+func buildKafkaNotifier(conf *config, logger *log.Logger) (*notify.KafkaNotifier, error) {
+	kafkaCli, err := kafka.NewClient(conf.Kafka.User, conf.Kafka.Pass,
+		conf.Kafka.BrokerURL, conf.Kafka.Topic)
+	if err != nil {
+		return nil, err
+	}
+	return notify.NewKafkaNotifier(&kafkaCli, logger), nil
+}
+
 func setupLogger(cfg config) *log.Logger {
 	var logger = log.New()
 

go.mod

@@ -6,6 +6,7 @@ require (
 	github.com/BurntSushi/toml v0.4.1
 	github.com/adevinta/vulcan-report v1.0.0
 	github.com/aws/aws-sdk-go v1.44.98
+	github.com/confluentinc/confluent-kafka-go v1.9.2
 	github.com/google/go-cmp v0.5.9
 	github.com/jmoiron/sqlx v1.3.4
 	github.com/lib/pq v1.10.3
@@ -14,7 +15,6 @@ require (
 
 require (
 	github.com/jmespath/go-jmespath v0.4.0 // indirect
-	github.com/stretchr/testify v1.7.0 // indirect
 	golang.org/x/sys v0.0.0-20220913175220-63ea55921009 // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect
 )

go.sum

@@ -0,0 +1,85 @@
+/*
+Copyright 2022 Adevinta
+*/
+
+package kafka
+
+import (
+	"errors"
+	"fmt"
+
+	"github.com/confluentinc/confluent-kafka-go/kafka"
+)
+
+var (
+	// ErrEmptyPayload is returned by the Push method of the [Client] when the
+	// given payload is empty.
+	ErrEmptyPayload = errors.New("payload can't be empty")
+)
+
+const (
+	kafkaSecurityProtocol = "sasl_ssl"
+	kafkaSaslMechanisms   = "SCRAM-SHA-256"
+)
+
+// Client implements an EventStreamClient using Kafka as the event stream
+// system.
+type Client struct {
+	producer *kafka.Producer
+	topic    string
+}
+
+// NewClient creates a new Kafka client connected to the a broker using the
+// given credentials and setting the mapping between all the entities and their
+// corresponding topics.
+func NewClient(user string, password string, broker string, topic string) (Client, error) {
+	config := kafka.ConfigMap{
+		"bootstrap.servers": broker,
+	}
+	if password != "" {
+		config.SetKey("security.protocol", kafkaSecurityProtocol)
+		config.SetKey("sasl.mechanisms", kafkaSaslMechanisms)
+		config.SetKey("sasl.username", user)
+		config.SetKey("sasl.password", password)
+	}
+	p, err := kafka.NewProducer(&config)
+	if err != nil {
+		return Client{}, err
+	}
+	return Client{p, topic}, nil
+}
+
+// Push sends the payload of an entity, with the specified id, to corresponding
+// topic according to the specified entity, using the kafka broker the client
+// is connected to. The method waits until kafka confirms the message has been
+// stored in the topic.
+func (c *Client) Push(id string, payload []byte, metadata map[string][]byte) error {
+	delivered := make(chan kafka.Event)
+	defer close(delivered)
+	var headers []kafka.Header
+	for k, v := range metadata {
+		headers = append(headers, kafka.Header{
+			Key:   k,
+			Value: v,
+		})
+	}
+	msg := kafka.Message{
+		TopicPartition: kafka.TopicPartition{
+			Topic:     &c.topic,
+			Partition: kafka.PartitionAny,
+		},
+		Key:     []byte(id),
+		Value:   []byte(payload),
+		Headers: headers,
+	}
+	err := c.producer.Produce(&msg, delivered)
+	if err != nil {
+		return fmt.Errorf("error producing message: %w", err)
+	}
+	e := <-delivered
+	m := e.(*kafka.Message)
+	if m.TopicPartition.Error != nil {
+		return fmt.Errorf("error delivering message: %w", m.TopicPartition.Error)
+	}
+	return nil
+}

pkg/asyncapi/kafka/client.go

@@ -0,0 +1,50 @@
+package notify
+
+import (
+	"encoding/json"
+
+	log "github.com/sirupsen/logrus"
+)
+
+const (
+	version = "v0.0.1" // Defines the notification schema version
+)
+
+// EventStreamClient represent a client of an event stream system, like Kafka
+// or AWS FIFO SQS queues.
+type EventStreamClient interface {
+	Push(id string, payload []byte, metadata map[string][]byte) error
+}
+
+// KafkaNotifier represents a Notifier implementation to send notifications to
+// a Kafka topic.
+type KafkaNotifier struct {
+	c EventStreamClient
+	l *log.Logger
+}
+
+// NewKafkaNotifier creates a new KafkaNotifier.
+func NewKafkaNotifier(c EventStreamClient, l *log.Logger) *KafkaNotifier {
+	return &KafkaNotifier{
+		c,
+		l,
+	}
+}
+
+// PushFinding sends the given FindingNotification to the configured Kafka topic
+// in the wrapped EventStreamClient.
+func (k *KafkaNotifier) PushFinding(f FindingNotification) error {
+	k.l.WithFields(log.Fields{
+		"notifier": "kafka",
+		"id":       f.ID,
+	}).Debug("pushing finding notification")
+
+	payload, err := json.Marshal(f)
+	if err != nil {
+		return err
+	}
+	meta := map[string][]byte{
+		"version": []byte(version),
+	}
+	return k.c.Push(f.ID, payload, meta)
+}

pkg/notify/kafka.go

@@ -0,0 +1,28 @@
+package notify
+
+// MultiNotifier represents a Notifier which delegates the
+// notification delivery into multiple notifier implementations.
+type MultiNotifier struct {
+	notifiers []Notifier
+}
+
+// NewMultiNotifier creates a new MultiNotifier.
+func NewMultiNotifier(notifiers ...Notifier) *MultiNotifier {
+	return &MultiNotifier{
+		notifiers: notifiers,
+	}
+}
+
+func (m *MultiNotifier) PushFinding(f FindingNotification) error {
+	// For every notifier wrapped by the multi implementation, push
+	// the given FindingNotification. Return an error on the first
+	// notifier that fails to deliver the notification, so possible
+	// repeated sending events might happen. This should be handled
+	// by the consumers complying with at-least-once semantics.
+	for iN := range m.notifiers {
+		if err := m.notifiers[iN].PushFinding(f); err != nil {
+			return err
+		}
+	}
+	return nil
+}

pkg/notify/multi.go

@@ -0,0 +1,13 @@
+package notify
+
+// NoopNotifier represents a notifier which performs no action.
+type NoopNotifier struct{}
+
+// NoopNotifier creates a new NoopNotifier.
+func NewNoopNotifier() *NoopNotifier {
+	return &NoopNotifier{}
+}
+
+func (n *NoopNotifier) PushFinding(f FindingNotification) error {
+	return nil // Do nothing
+}