use ip as partition key

.github/workflows/main.yml

@@ -29,15 +29,15 @@ jobs:
   # This workflow contains a single job called "build"
   build:
     # The type of runner that the job will run on
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-22.04
 
     # Steps represent a sequence of tasks that will be executed as part of the job
     steps:
       # Checks-out your repository under $GITHUB_WORKSPACE, so your job can access it
       - uses: actions/checkout@v2
       - uses: actions/setup-go@v2
         with:
-          go-version: '^1.16.1' # The Go version to download (if necessary) and use.
+          go-version: '^1.22' # The Go version to download (if necessary) and use.
       - name: install required packages
         run: sudo apt install libpcap-dev expect
       - run: go build -o ./mirroring-api-logging
@@ -68,7 +68,7 @@ jobs:
           # Build a docker container and push it to DockerHub 
           docker buildx create --use
           echo "Building and Pushing image to ECR..."
-          docker buildx build --platform linux/arm64/v8,linux/amd64,linux/arm/v7 -t $ECR_REGISTRY/$REGISTRY_ALIAS/mirror-api-logging:$IMAGE_TAG . --push
+          docker buildx build --platform linux/arm64/v8,linux/amd64 -t $ECR_REGISTRY/$REGISTRY_ALIAS/mirror-api-logging:$IMAGE_TAG . --push
           echo "::set-output name=image::$ECR_REGISTRY/$REGISTRY_ALIAS/mirror-api-logging:$IMAGE_TAG"
 
       - name: DockerHub login
@@ -89,6 +89,6 @@ jobs:
           # Build a docker container and push it to DockerHub 
           docker buildx create --use
           echo "Building and Pushing image to DockerHub..."
-          docker buildx build --platform linux/arm64/v8,linux/amd64,linux/arm/v7 -t $ECR_REGISTRY/mirror-api-logging:$IMAGE_TAG . --push
+          docker buildx build --platform linux/arm64/v8,linux/amd64 -t $ECR_REGISTRY/mirror-api-logging:$IMAGE_TAG . --push
           echo "::set-output name=image::$ECR_REGISTRY/mirror-api-logging:$IMAGE_TAG"
 

Dockerfile

@@ -1,6 +1,7 @@
 FROM golang:1.22-alpine
 RUN apk add build-base
 RUN apk add libpcap-dev
+RUN apk add tcpdump
 
 WORKDIR /app
 

go.mod

@@ -7,6 +7,7 @@ require (
 	github.com/google/uuid v1.6.0
 	github.com/segmentio/kafka-go v0.4.25
 	go.mongodb.org/mongo-driver v1.11.3
+	golang.org/x/net v0.30.0
 	google.golang.org/protobuf v1.36.1
 )
 
@@ -23,7 +24,6 @@ require (
 	github.com/xdg-go/stringprep v1.0.4 // indirect
 	github.com/youmark/pkcs8 v0.0.0-20181117223130-1be2e3e5546d // indirect
 	golang.org/x/crypto v0.28.0 // indirect
-	golang.org/x/net v0.30.0 // indirect
 	golang.org/x/sync v0.8.0 // indirect
 	golang.org/x/sys v0.26.0 // indirect
 	golang.org/x/text v0.19.0 // indirect

kafka.go

@@ -6,42 +6,80 @@ import (
 	trafficpb "github.com/akto-api-security/mirroring-api-logging/protobuf/traffic_payload"
 	"github.com/segmentio/kafka-go"
 	"google.golang.org/protobuf/proto"
-	"log"
+	"log/slog"
+	"strings"
 	"time"
 )
 
+var CLIENT_IP_HEADERS = []string{
+	"x-forwarded-for",
+	"x-real-ip",
+	"x-cluster-client-ip",
+	"true-client-ip",
+	"x-original-forwarded-for",
+	"x-client-ip",
+	"client-ip",
+}
+
 func Produce(kafkaWriter *kafka.Writer, ctx context.Context, value *trafficpb.HttpResponseParam) error {
 	// intialize the writer with the broker addresses, and the topic
 	protoBytes, err := proto.Marshal(value)
 	if err != nil {
-		log.Println("Failed to serialize protobuf message: ", err)
+		slog.Error("Failed to serialize protobuf message", "error", err)
 		return err
 	}
-
+
+	if value.Ip == "" {
+		slog.Warn("ip is empty, avoiding kafka push")
+		return nil
+	}
 	// Send serialized message to Kafka
+	topic := "akto.api.logs2"
 	msg := kafka.Message{
-		Topic: "akto.api.logs2",
-		Key:   []byte("testkey"),
+		Topic: topic,
+		Key:   []byte(value.Ip), // what to do when ip is empty?
 		Value: protoBytes,
 	}
 
 	err = kafkaWriter.WriteMessages(ctx, msg)
 	if err != nil {
-		log.Println("ERROR while writing messages: ", err)
+		slog.Error("Kafka write for threat failed", "topic", topic, "error", err)
 	}
 	return err
 }
 
+func GetSourceIp(reqHeaders map[string]*trafficpb.StringList, packetIp string) string {
+
+	for _, header := range CLIENT_IP_HEADERS {
+		if headerValues, exists := reqHeaders[header]; exists {
+			for _, headerValue := range headerValues.Values {
+				parts := strings.Split(headerValue, ",")
+				for _, part := range parts {
+					ip := strings.TrimSpace(part)
+					if ip != "" {
+						slog.Debug("Ip found in", "the header",  header)
+						return ip
+					}
+				}
+			}
+		}
+	}
+
+	slog.Debug("No ip found in headers returning", "packetIp", packetIp)
+	return packetIp
+}
+
 func ProduceStr(kafkaWriter *kafka.Writer, ctx context.Context, message string) error {
 	// intialize the writer with the broker addresses, and the topic
+	topic := "akto.api.logs"
 	msg := kafka.Message{
-		Topic: "akto.api.logs",
+		Topic: topic,
 		Value: []byte(message),
 	}
 	err := kafkaWriter.WriteMessages(ctx, msg)
 
 	if err != nil {
-		log.Println("ERROR while writing messages: ", err)
+		slog.Error("Kafka write for runtime failed", "topic", topic, "error", err)
 		return err
 	}
 
@@ -57,6 +95,8 @@ func GetKafkaWriter(kafkaURL, topic string, batchSize int, batchTimeout time.Dur
 		MaxAttempts:  1,
 		ReadTimeout:  batchTimeout,
 		WriteTimeout: batchTimeout,
+		Balancer:     &kafka.Hash{},
+		Compression: kafka.Zstd,
 	}
 }
 
@@ -76,7 +116,7 @@ func GetCredential(kafkaURL string, groupID string, topic string) Credential {
 	defer func(r *kafka.Reader) {
 		err := r.Close()
 		if err != nil {
-			log.Printf("could not close reader: %v", err)
+			slog.Error("could not close kafka reader", "error", err)
 		}
 	}(r)
 
@@ -86,25 +126,25 @@ func GetCredential(kafkaURL string, groupID string, topic string) Credential {
 	for {
 		select {
 		case <-ctx.Done():
-			log.Println("Timeout reached, no message received.")
+			slog.Error("Timeout reached, no message received.")
 			return msg // Return empty Credential if timeout occurs
 		default:
 			// Attempt to read a message from the Kafka topic
 			m, err := r.ReadMessage(ctx)
 			if err != nil {
 				if err == context.DeadlineExceeded {
-					log.Println("Timeout reached, no message received.")
+					slog.Error("Timeout reached, no message received.")
 					return msg
 				}
-				log.Printf("could not read message: %v", err)
+				slog.Error("Kafka Read failed for", "topic", topic, "error", err)
 				return msg // Return empty Credential on read error
 			}
 
-			log.Println("Found message: " + string(m.Value))
+			slog.Debug("Found message: " + string(m.Value))
 
 			err = json.Unmarshal(m.Value, &msg)
 			if err != nil {
-				log.Printf("could not unmarshal message: %v", err)
+				slog.Error("could not unmarshal kafka message", "error", err)
 				return msg // Return empty Credential on unmarshal error
 			}