From 757a4f63e5ff68a358df99c6b86b94bc1d8d32a8 Mon Sep 17 00:00:00 2001
From: Kevin Franklin Kim <kevin.kim@bestbytes.com>
Date: Thu, 21 Nov 2024 12:36:05 +0100
Subject: [PATCH] feat: add middlewares

---
 pkg/collect/collect.go      | 16 ++++++------
 pkg/http/mpv2/middleware.go | 49 +++++++++++++++++++++++++++++++++++++
 2 files changed, 56 insertions(+), 9 deletions(-)

diff --git a/pkg/collect/collect.go b/pkg/collect/collect.go
index 99e08c0..8746184 100644
--- a/pkg/collect/collect.go
+++ b/pkg/collect/collect.go
@@ -17,12 +17,11 @@ import (
 
 type (
 	Collect struct {
-		l                      *zap.Logger
-		taggingURL             string
-		taggingClient          *http.Client
-		taggingMaxResponseCode int
-		gtagHTTPMiddlewares    []gtaghttp.Middleware
-		mpv2HTTPMiddlewares    []mpv2http.Middleware
+		l                   *zap.Logger
+		taggingURL          string
+		taggingClient       *http.Client
+		gtagHTTPMiddlewares []gtaghttp.Middleware
+		mpv2HTTPMiddlewares []mpv2http.Middleware
 	}
 	Option func(*Collect) error
 )
@@ -65,9 +64,8 @@ func WithMPv2HTTPMiddlewares(v ...mpv2http.Middleware) Option {
 
 func New(l *zap.Logger, opts ...Option) (*Collect, error) {
 	inst := &Collect{
-		l:                      l,
-		taggingClient:          http.DefaultClient,
-		taggingMaxResponseCode: http.StatusBadRequest,
+		l:             l,
+		taggingClient: http.DefaultClient,
 	}
 
 	for _, opt := range opts {
diff --git a/pkg/http/mpv2/middleware.go b/pkg/http/mpv2/middleware.go
index 0172342..4251f90 100644
--- a/pkg/http/mpv2/middleware.go
+++ b/pkg/http/mpv2/middleware.go
@@ -94,6 +94,55 @@ func MiddlewareTimestamp(next MiddlewareHandler) MiddlewareHandler {
 	}
 }
 
+func MiddlewareUserAgentEventParam(next MiddlewareHandler) MiddlewareHandler {
+	return func(l *zap.Logger, w http.ResponseWriter, r *http.Request, payload *mpv2.Payload[any]) error {
+		if userAgent := r.Header.Get("User-Agent"); userAgent != "" {
+			for i, event := range payload.Events {
+				if value, ok := event.Params.(map[string]any); ok {
+					value["user_agent"] = userAgent
+				}
+				payload.Events[i] = event
+			}
+		}
+		return next(l, w, r, payload)
+	}
+}
+
+func MiddlewareIPOverrideEventParam(next MiddlewareHandler) MiddlewareHandler {
+	return func(l *zap.Logger, w http.ResponseWriter, r *http.Request, payload *mpv2.Payload[any]) error {
+		var ipOverride string
+		for _, key := range []string{"X-Original-Forwarded-For", "X-Forwarded-For", "X-Real-Ip"} {
+			if value := r.Header.Get(key); value != "" {
+				ipOverride = value
+				break
+			}
+		}
+		if ipOverride != "" {
+			for i, event := range payload.Events {
+				if value, ok := event.Params.(map[string]any); ok {
+					value["ip_override"] = ipOverride
+				}
+				payload.Events[i] = event
+			}
+		}
+		return next(l, w, r, payload)
+	}
+}
+
+func MiddlewarePageLocationEventParam(next MiddlewareHandler) MiddlewareHandler {
+	return func(l *zap.Logger, w http.ResponseWriter, r *http.Request, payload *mpv2.Payload[any]) error {
+		if referrer := r.Header.Get("Referer"); referrer != "" {
+			for i, event := range payload.Events {
+				if value, ok := event.Params.(map[string]any); ok {
+					value["page_location"] = referrer
+				}
+				payload.Events[i] = event
+			}
+		}
+		return next(l, w, r, payload)
+	}
+}
+
 func MiddlewareLogger(next MiddlewareHandler) MiddlewareHandler {
 	return func(l *zap.Logger, w http.ResponseWriter, r *http.Request, payload *mpv2.Payload[any]) error {
 		eventNames := make([]string, len(payload.Events))