Send X-Agent headers with requests to Aikido

lib/agent/aikido_types/init_data.go

@@ -154,6 +154,9 @@ type ServerData struct {
 	// Logger for the server
 	Logger *log.AikidoLogger
 
+	// Each app gets its own session ID, and a PHP restart (new main PID) gets a new one
+	SessionID string
+
 	// Aikido config that contains info about endpoint, log_level, token, ...
 	AikidoConfig AikidoConfigData
 
@@ -244,6 +247,7 @@ const (
 
 func NewServerData() *ServerData {
 	return &ServerData{
+		SessionID:               newSessionID(),
 		Hostnames:               make(map[string]map[uint32]uint64),
 		HostnamesQueue:          NewQueue[string](MaxNumberOfStoredHostnames),
 		Routes:                  make(map[string]map[string]*Route),

lib/agent/aikido_types/session_id.go

@@ -0,0 +1,26 @@
+package aikido_types
+
+import (
+	"crypto/rand"
+	"fmt"
+	"time"
+)
+
+// The custom Go toolchain we use for Alpine has no uuid package
+func newSessionID() string {
+	var b [16]byte
+	if _, err := rand.Read(b[:]); err != nil {
+		return "unknown"
+	}
+
+	// 48-bit millisecond timestamp, big-endian
+	ms := time.Now().UnixMilli()
+	for i := range 6 {
+		b[i] = byte(ms >> (8 * (5 - i)) & 0xff) //nolint:gosec // masked to a single byte
+	}
+
+	b[6] = (b[6] & 0x0f) | 0x70 // version 7
+	b[8] = (b[8] & 0x3f) | 0x80 // RFC 9562 variant
+
+	return fmt.Sprintf("%x-%x-%x-%x-%x", b[0:4], b[4:6], b[6:8], b[8:10], b[10:16])
+}

lib/agent/aikido_types/session_id_test.go

@@ -0,0 +1,24 @@
+package aikido_types
+
+import (
+	"strconv"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNewSessionID(t *testing.T) {
+	before := time.Now().UnixMilli()
+	id := newSessionID()
+	after := time.Now().UnixMilli()
+
+	assert.Regexp(t, `^[0-9a-f]{8}-[0-9a-f]{4}-7[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$`, id)
+
+	ms, err := strconv.ParseInt(id[0:8]+id[9:13], 16, 64)
+	assert.NoError(t, err)
+	assert.GreaterOrEqual(t, ms, before)
+	assert.LessOrEqual(t, ms, after)
+
+	assert.NotEqual(t, id, newSessionID())
+}

lib/agent/cloud/common.go

@@ -1,12 +1,14 @@
 package cloud
 
 import (
+	"cmp"
 	"encoding/json"
 	. "main/aikido_types"
 	"main/constants"
 	"main/globals"
 	"main/log"
 	"main/utils"
+	"net/http"
 	"regexp"
 	"runtime"
 	"strings"
@@ -34,6 +36,15 @@ func GetAgentInfo(server *ServerData) AgentInfo {
 	}
 }
 
+func setAgentHeaders(req *http.Request, server *ServerData) {
+	req.Header.Set("X-Agent-Platform", "php")
+	req.Header.Set("X-Agent-Library", "firewall-php")
+	req.Header.Set("X-Agent-Version", constants.Version)
+	req.Header.Set("X-Agent-Hostname", cmp.Or(globals.Machine.HostName, "unknown"))
+	req.Header.Set("X-Agent-IP-Address", cmp.Or(globals.Machine.IPAddress, "unknown"))
+	req.Header.Set("X-Agent-Session-Id", server.SessionID)
+}
+
 func ResetHeartbeatTicker(server *ServerData) {
 	if !server.CloudConfig.ReceivedAnyStats {
 		log.Info(server.Logger, "Resetting HeartbeatTicker to 1m!")

lib/agent/cloud/event.go

@@ -57,6 +57,7 @@ func sendCloudRequest(server *ServerData, endpoint string, route string, method
 	req.Header.Set("Authorization", token)
 	req.Header.Set("Content-Type", "application/json")
 	req.Header.Set("Accept-Encoding", "gzip")
+	setAgentHeaders(req, server)
 	// Bounded, so a cloud that accepts the connection but never answers cannot block the calling routine forever
 	client := &http.Client{Timeout: constants.CloudRequestTimeoutInMs * time.Millisecond}
 	resp, err := client.Do(req)

lib/agent/cloud/sse.go

@@ -70,8 +70,7 @@ func connectToConfigStream(ctx context.Context, server *ServerData, readTimeout
 	req.Header.Set("Accept", "text/event-stream")
 	req.Header.Set("Cache-Control", "no-cache")
 	req.Header.Set("Accept-Encoding", "identity")
-	req.Header.Set("X-Agent-Platform", "php")
-	req.Header.Set("X-Agent-Version", constants.Version)
+	setAgentHeaders(req, server)
 
 	log.Debugf(server.Logger, "[%s] Connecting to config stream %s", utils.AnonymizeToken(token), streamEndpoint)
 

lib/agent/cloud/sse_test.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	. "main/aikido_types"
 	"main/constants"
+	"main/globals"
 	"main/log"
 	"net/http"
 	"net/http/httptest"
@@ -36,6 +37,7 @@ type mockCloud struct {
 	streamDisconnects atomic.Int32
 	configRequests    atomic.Int32
 	streamHeaders     atomic.Value
+	configHeaders     atomic.Value
 }
 
 func newMockCloud() *mockCloud {
@@ -78,6 +80,7 @@ func newMockCloud() *mockCloud {
 
 	mux.HandleFunc("/api/runtime/config", func(w http.ResponseWriter, r *http.Request) {
 		cloud.configRequests.Add(1)
+		cloud.configHeaders.Store(r.Header.Clone())
 		w.Header().Set("Content-Type", "application/json")
 		fmt.Fprintf(w, `{"success":true,"serviceId":1,"configUpdatedAt":%d,"heartbeatIntervalInMS":600000,
 			"endpoints":[],"blockedUserIds":[],"allowedIPAddresses":[],"receivedAnyStats":true,"block":true}`,
@@ -120,6 +123,17 @@ func getConfigUpdatedAt(server *ServerData) int64 {
 	return server.CloudConfig.ConfigUpdatedAt
 }
 
+func assertAgentHeaders(t *testing.T, server *ServerData, headers http.Header) {
+	t.Helper()
+	assert.Equal(t, "php", headers.Get("X-Agent-Platform"))
+	assert.Equal(t, "firewall-php", headers.Get("X-Agent-Library"))
+	assert.Equal(t, constants.Version, headers.Get("X-Agent-Version"))
+	assert.Equal(t, "test-host", headers.Get("X-Agent-Hostname"))
+	assert.Equal(t, "unknown", headers.Get("X-Agent-IP-Address"))
+	assert.NotEmpty(t, server.SessionID)
+	assert.Equal(t, server.SessionID, headers.Get("X-Agent-Session-Id"))
+}
+
 func waitFor(timeout time.Duration, condition func() bool) bool {
 	deadline := time.Now().Add(timeout)
 	for time.Now().Before(deadline) {
@@ -133,6 +147,10 @@ func waitFor(timeout time.Duration, condition func() bool) bool {
 
 func TestConfigStreamRoutine(t *testing.T) {
 	t.Run("it applies a config pushed over the stream", func(t *testing.T) {
+		originalMachine := globals.Machine
+		globals.Machine = MachineData{HostName: "test-host"}
+		defer func() { globals.Machine = originalMachine }()
+
 		cloud := newMockCloud()
 		server := newTestServerData(cloud, "AIK_RUNTIME_TEST_TOKEN", true)
 
@@ -152,15 +170,15 @@ func TestConfigStreamRoutine(t *testing.T) {
 		assert.Equal(t, "text/event-stream", streamHeaders.Get("Accept"))
 		assert.Equal(t, "no-cache", streamHeaders.Get("Cache-Control"))
 		assert.Equal(t, "identity", streamHeaders.Get("Accept-Encoding"))
-		assert.Equal(t, "php", streamHeaders.Get("X-Agent-Platform"))
-		assert.Equal(t, constants.Version, streamHeaders.Get("X-Agent-Version"))
+		assertAgentHeaders(t, server, streamHeaders)
 
 		cloud.publishConfig(t, 1000)
 
 		assert.True(t, waitFor(5*time.Second, func() bool {
 			return getConfigUpdatedAt(server) == 1000
 		}), "expected the pushed config to be applied")
 		assert.Equal(t, int32(1), cloud.configRequests.Load())
+		assertAgentHeaders(t, server, cloud.configHeaders.Load().(http.Header))
 
 		// A second event for a config we already have must not trigger another fetch
 		cloud.publishConfig(t, 1000)