Also support routing by headers triple (operationID[alias], taskName, service)

server/pkg/auth.go

@@ -84,21 +84,29 @@ func (s *authServer) findTaskByRequest(host string, headers map[string]string) (
 	defer s.mx.RUnlock()
 
 	var hash string
-	if routerHeaderValue, ok := headers[routerHeaderName]; ok {
+	if routerHeaderValue, ok := headers[idRouterHeaderName]; ok {
 		hash = routerHeaderValue
+	} else if operationID, ok := headers[operationIDRouterHeaderName]; ok {
+		hash = taskHash(operationID, headers[taskNameRouterHeaderName], headers[serviceRouteHeaderName])
+	} else if operationAlias, ok := headers[operationAliasRouterHeaderName]; ok {
+		operationID, ok := s.operationAliasToID[operationAlias]
+		if !ok {
+			return nil, fmt.Errorf("operation by alias %q from header was not found", operationAlias)
+		}
+		hash = taskHash(operationID, headers[taskNameRouterHeaderName], headers[serviceRouteHeaderName])
 	} else if host != "" {
 		subdomain := strings.Split(host, ".")[0]
 		if operationAlias, taskName, service, ok := tryParseAliasSubdomain(subdomain); ok {
 			operationID, ok := s.operationAliasToID[operationAlias]
 			if !ok {
 				return nil, fmt.Errorf("operation by alias %q from subdomain was not found", operationAlias)
 			}
-			hash = (&Task{operationID: operationID, taskName: taskName, service: service}).Hash()
+			hash = taskHash(operationID, taskName, service)
 		} else {
 			hash = subdomain
 		}
 	} else {
-		return nil, fmt.Errorf("authority (host) or %s headers are missing in request", routerHeaderName)
+		return nil, fmt.Errorf("authority (host) or %s headers are missing in request", idRouterHeaderName)
 	}
 
 	if task, ok := s.hashToTasks[hash]; !ok {
@@ -215,3 +223,7 @@ var (
 		},
 	}
 )
+
+func taskHash(operationID, taskName, service string) string {
+	return (&Task{operationID: operationID, taskName: taskName, service: service}).Hash()
+}

server/pkg/auth_test.go

@@ -86,7 +86,93 @@ func TestFindTaskByRequest(t *testing.T) {
 			expectedID: task1.operationID,
 		},
 
-		// Source 2: Alias-based subdomain (format: alias-taskname-service)
+		// Source 2: (operation-id, task-name, service) headers
+		{
+			name: "operation-id headers - valid task",
+			host: "ignored.example.com",
+			headers: map[string]string{
+				"x-yt-taskproxy-operation-id": task1.operationID,
+				"x-yt-taskproxy-task-name":    task1.taskName,
+				"x-yt-taskproxy-service":      task1.service,
+			},
+			expectedID: task1.operationID,
+		},
+		{
+			name: "operation-id headers - unknown task",
+			host: "ignored.example.com",
+			headers: map[string]string{
+				"x-yt-taskproxy-operation-id": "op-unknown",
+				"x-yt-taskproxy-task-name":    "worker",
+				"x-yt-taskproxy-service":      "api",
+			},
+			errorMsg: "no entry for hash",
+		},
+		{
+			name: "operation-id headers - takes precedence over host",
+			host: task3Hash + ".example.com",
+			headers: map[string]string{
+				"x-yt-taskproxy-operation-id": task1.operationID,
+				"x-yt-taskproxy-task-name":    task1.taskName,
+				"x-yt-taskproxy-service":      task1.service,
+			},
+			expectedID: task1.operationID,
+		},
+		{
+			name: "operation-id headers - id header takes precedence over operation-id headers",
+			host: "ignored.example.com",
+			headers: map[string]string{
+				"x-yt-taskproxy-id":           task1Hash,
+				"x-yt-taskproxy-operation-id": task3.operationID,
+				"x-yt-taskproxy-task-name":    task3.taskName,
+				"x-yt-taskproxy-service":      task3.service,
+			},
+			expectedID: task1.operationID,
+		},
+
+		// Source 3: (operation-alias, task-name, service) headers
+		{
+			name: "operation-alias headers - valid alias",
+			host: "ignored.example.com",
+			headers: map[string]string{
+				"x-yt-taskproxy-operation-alias": task2.operationAlias,
+				"x-yt-taskproxy-task-name":       task2.taskName,
+				"x-yt-taskproxy-service":         task2.service,
+			},
+			expectedID: task2.operationID,
+		},
+		{
+			name: "operation-alias headers - unknown alias",
+			host: "ignored.example.com",
+			headers: map[string]string{
+				"x-yt-taskproxy-operation-alias": "unknownalias",
+				"x-yt-taskproxy-task-name":       "worker",
+				"x-yt-taskproxy-service":         "api",
+			},
+			errorMsg: "operation by alias \"unknownalias\" from header was not found",
+		},
+		{
+			name: "operation-alias headers - takes precedence over host",
+			host: task3Hash + ".example.com",
+			headers: map[string]string{
+				"x-yt-taskproxy-operation-alias": task2.operationAlias,
+				"x-yt-taskproxy-task-name":       task2.taskName,
+				"x-yt-taskproxy-service":         task2.service,
+			},
+			expectedID: task2.operationID,
+		},
+		{
+			name: "operation-alias headers - operation-id headers take precedence",
+			host: "ignored.example.com",
+			headers: map[string]string{
+				"x-yt-taskproxy-operation-id":    task1.operationID,
+				"x-yt-taskproxy-operation-alias": task2.operationAlias,
+				"x-yt-taskproxy-task-name":       task1.taskName,
+				"x-yt-taskproxy-service":         task1.service,
+			},
+			expectedID: task1.operationID,
+		},
+
+		// Source 4: Alias-based subdomain (format: alias-taskname-service)
 		{
 			name: "alias subdomain - valid alias",
 			host: "myalias-master-ui.example.com",
@@ -111,7 +197,7 @@ func TestFindTaskByRequest(t *testing.T) {
 			expectedID: task2.operationID,
 		},
 
-		// Source 3: Direct hash from subdomain (fallback)
+		// Source 5: Direct hash from subdomain (fallback)
 		{
 			name:       "direct hash subdomain - valid hash",
 			host:       task1Hash + ".example.com",

server/pkg/xds.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"log"
 	"net"
+	"sort"
 	"time"
 
 	"google.golang.org/grpc"
@@ -37,7 +38,13 @@ import (
 
 const (
 	extAuthClusterName = "extAuthz"
-	routerHeaderName   = "x-yt-taskproxy-id"
+
+	idRouterHeaderName = "x-yt-taskproxy-id" // hash
+
+	operationIDRouterHeaderName    = "x-yt-taskproxy-operation-id"
+	operationAliasRouterHeaderName = "x-yt-taskproxy-operation-alias"
+	taskNameRouterHeaderName       = "x-yt-taskproxy-task-name"
+	serviceRouteHeaderName         = "x-yt-taskproxy-service"
 )
 
 func ServeGRPC(s serverv3.Server, authServer *authServer) error {
@@ -99,25 +106,38 @@ func makeSnapshot(hashToTask map[string]Task, version string, baseDomain string,
 				Action: action,
 			}},
 		})
-		// ... or by custom header
+		// ... or by custom header(-s)
 		defaultVhostRoutes = append(defaultVhostRoutes, &routev3.Route{
 			Match: &routev3.RouteMatch{
 				PathSpecifier: &routev3.RouteMatch_Prefix{Prefix: "/"},
-				Headers: []*routev3.HeaderMatcher{
-					{
-						Name: routerHeaderName,
-						HeaderMatchSpecifier: &routev3.HeaderMatcher_StringMatch{
-							StringMatch: &matcherv3.StringMatcher{
-								MatchPattern: &matcherv3.StringMatcher_Exact{
-									Exact: hash,
-								},
-							},
-						},
-					},
-				},
+				Headers:       makeHeaderMatchers(map[string]string{idRouterHeaderName: hash}),
 			},
 			Action: action,
 		})
+		defaultVhostRoutes = append(defaultVhostRoutes, &routev3.Route{
+			Match: &routev3.RouteMatch{
+				PathSpecifier: &routev3.RouteMatch_Prefix{Prefix: "/"},
+				Headers: makeHeaderMatchers(map[string]string{
+					operationIDRouterHeaderName: task.operationID,
+					taskNameRouterHeaderName:    task.taskName,
+					serviceRouteHeaderName:      task.service,
+				}),
+			},
+			Action: action,
+		})
+		if task.operationAlias != "" {
+			defaultVhostRoutes = append(defaultVhostRoutes, &routev3.Route{
+				Match: &routev3.RouteMatch{
+					PathSpecifier: &routev3.RouteMatch_Prefix{Prefix: "/"},
+					Headers: makeHeaderMatchers(map[string]string{
+						operationAliasRouterHeaderName: task.operationAlias,
+						taskNameRouterHeaderName:       task.taskName,
+						serviceRouteHeaderName:         task.service,
+					}),
+				},
+				Action: action,
+			})
+		}
 	}
 
 	defaultVhostRoutes = append(defaultVhostRoutes, &routev3.Route{
@@ -320,3 +340,28 @@ func mustAny(m proto.Message) *anypb.Any {
 	}
 	return a
 }
+
+func makeHeaderMatchers(headers map[string]string) []*routev3.HeaderMatcher {
+	// Sort keys for deterministic order
+	keys := make([]string, 0, len(headers))
+	for name := range headers {
+		keys = append(keys, name)
+	}
+	sort.Strings(keys)
+
+	matchers := make([]*routev3.HeaderMatcher, 0, len(headers))
+	for _, name := range keys {
+		value := headers[name]
+		matchers = append(matchers, &routev3.HeaderMatcher{
+			Name: name,
+			HeaderMatchSpecifier: &routev3.HeaderMatcher_StringMatch{
+				StringMatch: &matcherv3.StringMatcher{
+					MatchPattern: &matcherv3.StringMatcher_Exact{
+						Exact: value,
+					},
+				},
+			},
+		})
+	}
+	return matchers
+}