feat[agent-manager](tenants): persist collector tenant from connectio…

agent-manager/agent/collector_imp.go

@@ -148,11 +148,20 @@ func InitCollectorService() {
 }
 
 func (s *CollectorService) RegisterCollector(ctx context.Context, req *RegisterRequest) (*AuthResponse, error) {
+	tenantID, ok := tenantFromContext(ctx)
+	if !ok {
+		// The interceptor is the only writer of this key on the connection-key
+		// path; missing it means the route was reached through some other auth
+		// that shouldn't be allowed to enrol collectors.
+		return nil, status.Error(codes.PermissionDenied, "missing tenant on register")
+	}
+
 	collector := &models.Collector{
 		Ip:       req.GetIp(),
 		Hostname: req.GetHostname(),
 		Version:  req.GetVersion(),
 		Module:   models.CollectorModule(req.GetCollector().String()),
+		TenantID: tenantID,
 	}
 
 	oldCollector := &models.Collector{}
@@ -178,7 +187,7 @@ func (s *CollectorService) RegisterCollector(ctx context.Context, req *RegisterR
 	}
 
 	s.CacheCollectorKeyMutex.Lock()
-	entry := utils.ConnectorAuth{Key: key, TenantID: tenantOrDefault(collector.TenantID)}
+	entry := utils.ConnectorAuth{Key: key, TenantID: collector.TenantID}
 	s.CacheCollectorKey[collector.ID] = entry
 	AuthCache.PublishCollector(collector.ID, entry)
 	s.CacheCollectorKeyMutex.Unlock()

agent-manager/agent/interceptor.go

@@ -20,29 +20,45 @@ func UnaryInterceptor(ctx context.Context, req interface{}, info *grpc.UnaryServ
 		return nil, status.Error(codes.Unauthenticated, "metadata is not provided")
 	}
 
-	err := authHeaders(md, info.FullMethod)
+	tenantID, err := authHeaders(md, info.FullMethod)
 	if err != nil {
 		return nil, err
 	}
 
+	if tenantID != "" {
+		ctx = withTenant(ctx, tenantID)
+	}
 	return handler(ctx, req)
 }
 
+type tenantServerStream struct {
+	grpc.ServerStream
+	ctx context.Context
+}
+
+func (s *tenantServerStream) Context() context.Context { return s.ctx }
+
 func StreamInterceptor(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
 	md, ok := metadata.FromIncomingContext(ss.Context())
 	if !ok {
 		return status.Error(codes.Unauthenticated, "metadata is not provided")
 	}
 
-	err := authHeaders(md, info.FullMethod)
+	tenantID, err := authHeaders(md, info.FullMethod)
 	if err != nil {
 		return err
 	}
 
+	if tenantID != "" {
+		ss = &tenantServerStream{ServerStream: ss, ctx: withTenant(ss.Context(), tenantID)}
+	}
 	return handler(srv, ss)
 }
 
-func authHeaders(md metadata.MD, fullMethod string) error {
+// authHeaders returns the tenantID resolved from the auth material when the
+// connection-key path matched; otherwise "". A non-nil error means the call
+// is unauthenticated and must be rejected.
+func authHeaders(md metadata.MD, fullMethod string) (string, error) {
 	var authType string
 	var routes []string
 	authKey := md.Get("key")
@@ -61,43 +77,45 @@ func authHeaders(md metadata.MD, fullMethod string) error {
 		authType = "internal-key"
 		routes = config.InternalKeyRoutes()
 	} else {
-		return status.Error(codes.Unauthenticated, "auth is not provided")
+		return "", status.Error(codes.Unauthenticated, "auth is not provided")
 	}
 
 	if !isInRoute(fullMethod, routes) {
-		return status.Error(codes.PermissionDenied, fmt.Sprintf("route is not registered for authentication with %s auth type", authType))
+		return "", status.Error(codes.PermissionDenied, fmt.Sprintf("route is not registered for authentication with %s auth type", authType))
 	}
 
 	switch authType {
 	case "key":
 		key := authKey[0]
 		id, err := strconv.ParseUint(authId[0], 10, 32)
 		if err != nil {
-			return status.Error(codes.PermissionDenied, "id is not valid")
+			return "", status.Error(codes.PermissionDenied, "id is not valid")
 		}
 		typ := strings.ToLower(connectorType[0])
 		switch typ {
 		case "agent":
 			if !AgentServ.ValidateAgentKey(key, uint(id)) {
-				return status.Error(codes.PermissionDenied, "invalid key")
+				return "", status.Error(codes.PermissionDenied, "invalid key")
 			}
 		case "collector":
 			if !CollectorServ.ValidateCollectorKey(key, uint(id)) {
-				return status.Error(codes.PermissionDenied, "invalid key")
+				return "", status.Error(codes.PermissionDenied, "invalid key")
 			}
 		default:
-			return status.Error(codes.PermissionDenied, "invalid type")
+			return "", status.Error(codes.PermissionDenied, "invalid type")
 		}
 	case "connection-key":
-		if !AgentServ.ValidateConnectionKey(authConnectionKey[0]) {
-			return status.Error(codes.PermissionDenied, "invalid connection key")
+		tenantID, ok := AgentServ.TenantForConnectionKey(authConnectionKey[0])
+		if !ok {
+			return "", status.Error(codes.PermissionDenied, "invalid connection key")
 		}
+		return tenantID, nil
 	case "internal-key":
 		if !isInternalKeyValid(authInternalKey[0]) {
-			return status.Error(codes.PermissionDenied, "internal key does not match")
+			return "", status.Error(codes.PermissionDenied, "internal key does not match")
 		}
 	}
-	return nil
+	return "", nil
 }
 
 func isInternalKeyValid(token string) bool {

agent-manager/agent/tenant_ctx.go

@@ -0,0 +1,20 @@
+package agent
+
+import "context"
+
+// Handlers on the key/internal-key paths must fall back to the request
+// payload (or tenantOrDefault) if tenantFromContext returns false.
+
+type tenantCtxKey struct{}
+
+func withTenant(ctx context.Context, tenantID string) context.Context {
+	return context.WithValue(ctx, tenantCtxKey{}, tenantID)
+}
+
+func tenantFromContext(ctx context.Context) (string, bool) {
+	t, ok := ctx.Value(tenantCtxKey{}).(string)
+	if !ok || t == "" {
+		return "", false
+	}
+	return t, true
+}