Merge pull request PelicanPlatform#1708 from patrickbrophy/fix-issue-…

…1635 Fix `proj` label setting bug
jhiemstrawisc · Nov 12, 2024 · 48e8576 · 48e8576
2 parents dbb8da6 + c5d82ad
commit 48e8576
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 28 deletions.
diff --git a/director/director.go b/director/director.go
@@ -294,23 +294,6 @@ func extractVersionAndService(ginCtx *gin.Context) (reqVer *version.Version, ser
 	return reqVer, service, nil
 }
 
-// Helper function to extract project from User-Agent
-// Will return an empty string if no project is found
-func extractProjectFromUserAgent(userAgents []string) string {
-	prefix := "project/"
-
-	for _, userAgent := range userAgents {
-		parts := strings.Split(userAgent, " ")
-		for _, part := range parts {
-			if strings.HasPrefix(part, prefix) {
-				return strings.TrimPrefix(part, prefix)
-			}
-		}
-	}
-
-	return ""
-}
-
 func versionCompatCheck(reqVer *version.Version, service string) error {
 	var minCompatVer *version.Version
 	switch service {
@@ -510,7 +493,7 @@ func redirectToCache(ginCtx *gin.Context) {
 	}
 
 	ctx := context.Background()
-	project := extractProjectFromUserAgent(ginCtx.Request.Header.Values("User-Agent"))
+	project := utils.ExtractProjectFromUserAgent(ginCtx.Request.Header.Values("User-Agent"))
 	ctx = context.WithValue(ctx, ProjectContextKey{}, project)
 	cacheAds, err = sortServerAds(ctx, ipAddr, cacheAds, cachesAvailabilityMap)
 	if err != nil {
@@ -747,7 +730,7 @@ func redirectToOrigin(ginCtx *gin.Context) {
 	}
 
 	ctx := context.Background()
-	project := extractProjectFromUserAgent(ginCtx.Request.Header.Values("User-Agent"))
+	project := utils.ExtractProjectFromUserAgent(ginCtx.Request.Header.Values("User-Agent"))
 	ctx = context.WithValue(ctx, ProjectContextKey{}, project)
 
 	availableAds, err = sortServerAds(ctx, ipAddr, availableAds, nil)

diff --git a/director/director_test.go b/director/director_test.go
@@ -50,6 +50,7 @@ import (
 	"github.com/pelicanplatform/pelican/test_utils"
 	"github.com/pelicanplatform/pelican/token"
 	"github.com/pelicanplatform/pelican/token_scopes"
+	"github.com/pelicanplatform/pelican/utils"
 )
 
 func NamespaceAdContainsPath(ns []server_structs.NamespaceAdV2, path string) bool {
@@ -2013,42 +2014,42 @@ func TestGetFinalRedirectURL(t *testing.T) {
 func TestExtractProjectFromUserAgent(t *testing.T) {
 	t.Run("Single User-Agent with project prefix", func(t *testing.T) {
 		userAgents := []string{"pelican-client/1.0.0 project/test"}
-		result := extractProjectFromUserAgent(userAgents)
+		result := utils.ExtractProjectFromUserAgent(userAgents)
 		assert.Equal(t, "test", result)
 	})
 	t.Run("Singlue User-Agent with swapped order", func(t *testing.T) {
 		userAgents := []string{"project/test pelican-client/1.0.0"}
-		result := extractProjectFromUserAgent(userAgents)
+		result := utils.ExtractProjectFromUserAgent(userAgents)
 		assert.Equal(t, "test", result)
 	})
 
 	t.Run("Single User-Agent with additional segments", func(t *testing.T) {
 		userAgents := []string{"pelican-client/blah project/myproject foo/bar"}
-		result := extractProjectFromUserAgent(userAgents)
+		result := utils.ExtractProjectFromUserAgent(userAgents)
 		assert.Equal(t, "myproject", result)
 	})
 
 	t.Run("Multiple User-Agents with project prefix", func(t *testing.T) {
 		userAgents := []string{"pelican-client/1.0.0 project/test", "pelican-client/1.0.0 project/test2"}
-		result := extractProjectFromUserAgent(userAgents)
+		result := utils.ExtractProjectFromUserAgent(userAgents)
 		assert.Equal(t, "test", result)
 	})
 
 	t.Run("Multiple User-Agents with swapped order", func(t *testing.T) {
 		userAgents := []string{"project/test pelican-client/1.0.0", "project/test2 pelican-client/1.0.0"}
-		result := extractProjectFromUserAgent(userAgents)
+		result := utils.ExtractProjectFromUserAgent(userAgents)
 		assert.Equal(t, "test", result)
 	})
 
 	t.Run("No Project Prefix", func(t *testing.T) {
 		userAgents := []string{"pelican-client/1.0.0"}
-		result := extractProjectFromUserAgent(userAgents)
+		result := utils.ExtractProjectFromUserAgent(userAgents)
 		assert.Equal(t, "", result)
 	})
 
 	t.Run("No User-Agent", func(t *testing.T) {
 		userAgents := []string{}
-		result := extractProjectFromUserAgent(userAgents)
+		result := utils.ExtractProjectFromUserAgent(userAgents)
 		assert.Equal(t, "", result)
 	})
 }
diff --git a/metrics/xrootd_metrics.go b/metrics/xrootd_metrics.go
@@ -904,13 +904,14 @@ func HandlePacket(packet []byte) error {
 		if xrdUserId, appinfo, err := GetSIDRest(packet[12 : 12+infoSize]); err == nil {
 			if userids.Has(xrdUserId) {
 				userId := userids.Get(xrdUserId).Value()
+				project := utils.ExtractProjectFromUserAgent([]string{appinfo})
 				if sessions.Has(userId) {
 					existingRec := sessions.Get(userId).Value()
-					existingRec.Project = appinfo
+					existingRec.Project = project
 					existingRec.Host = xrdUserId.Host
 					sessions.Set(userId, existingRec, ttlcache.DefaultTTL)
 				} else {
-					sessions.Set(userId, UserRecord{Project: appinfo, Host: xrdUserId.Host}, ttlcache.DefaultTTL)
+					sessions.Set(userId, UserRecord{Project: project, Host: xrdUserId.Host}, ttlcache.DefaultTTL)
 				}
 			}
 		} else {

diff --git a/utils/utils.go b/utils/utils.go
@@ -127,3 +127,20 @@ func ExtractVersionAndServiceFromUserAgent(userAgent string) (reqVer, service st
 	service = (strings.Split(userAgentSplit[0], "-"))[1]
 	return reqVer, service
 }
+
+// Helper function to extract project from User-Agent
+// Will return an empty string if no project is found
+func ExtractProjectFromUserAgent(userAgents []string) string {
+	prefix := "project/"
+
+	for _, userAgent := range userAgents {
+		parts := strings.Split(userAgent, " ")
+		for _, part := range parts {
+			if strings.HasPrefix(part, prefix) {
+				return strings.TrimPrefix(part, prefix)
+			}
+		}
+	}
+
+	return ""
+}