[Bug] Script routing bypassed with two service references #16545

package org.apache.dubbo.samples.script;

public interface OrderProbeService {
    String instance();
}

package org.apache.dubbo.samples.script;

public interface UserProbeService {
    String ping();
}

package org.apache.dubbo.samples.script;

import com.alibaba.nacos.api.NacosFactory;
import com.alibaba.nacos.api.PropertyKeyConst;
import com.alibaba.nacos.api.config.ConfigService;
import org.apache.curator.test.TestingServer;
import org.apache.dubbo.config.ApplicationConfig;
import org.apache.dubbo.config.ConfigCenterConfig;
import org.apache.dubbo.config.ProtocolConfig;
import org.apache.dubbo.config.ReferenceConfig;
import org.apache.dubbo.config.RegistryConfig;
import org.apache.dubbo.config.ServiceConfig;
import org.apache.dubbo.config.bootstrap.DubboBootstrap;
import org.apache.dubbo.rpc.cluster.router.condition.config.AppStateRouterFactory;

import java.io.BufferedReader;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.Properties;
import java.util.concurrent.CountDownLatch;
import java.util.stream.Collectors;

/** Real RPC reproduction of a shared AppStateRouter.nextRouter across two service chains. */
public final class RouterChainReproMain {
    private static final String REGISTRY_URL = "zookeeper://127.0.0.1:23281";
    private static final String ORDER_APP = "route-chain-repro-order";
    private static final String USER_APP = "route-chain-repro-user";
    private static final String RULE_DATA_ID = ORDER_APP + ".script-router";

    private RouterChainReproMain() {}

    public static void main(String[] args) throws Exception {
        if (args.length == 0) {
            throw new IllegalArgumentException("zookeeper | order-provider <id> <port> | user-provider | publish | consumer <single|multi> <count>");
        }
        switch (args[0]) {
            case "zookeeper":
                try (TestingServer server = new TestingServer(23281)) {
                    System.out.println("ZooKeeper ready: " + server.getConnectString());
                    new CountDownLatch(1).await();
                }
                return;
            case "order-provider":
                requireArgs(args, 3);
                runOrderProvider(args[1], Integer.parseInt(args[2]));
                return;
            case "user-provider":
                runUserProvider();
                return;
            case "publish":
                publishRule();
                return;
            case "consumer":
                requireArgs(args, 3);
                if (!"single".equals(args[1]) && !"multi".equals(args[1])
                        && !"fixed".equals(args[1]) && !"startup".equals(args[1])
                        && !"fixed-startup".equals(args[1])) {
                    throw new IllegalArgumentException("Unknown consumer mode: " + args[1]);
                }
                runConsumer(args[1], Integer.parseInt(args[2]));
                return;
            default:
                throw new IllegalArgumentException("Unknown role: " + args[0]);
        }
    }

    private static void requireArgs(String[] args, int count) {
        if (args.length != count) {
            throw new IllegalArgumentException("Expected " + count + " arguments, got " + args.length);
        }
    }

    private static void runOrderProvider(String id, int port) {
        ServiceConfig<OrderProbeService> service = new ServiceConfig<>();
        service.setInterface(OrderProbeService.class);
        service.setRef(() -> id);
        DubboBootstrap.getInstance()
                .application(new ApplicationConfig(ORDER_APP))
                .registry(new RegistryConfig(REGISTRY_URL))
                .protocol(new ProtocolConfig("dubbo", port))
                .service(service)
                .start()
                .await();
    }

    private static void runUserProvider() {
        ServiceConfig<UserProbeService> service = new ServiceConfig<>();
        service.setInterface(UserProbeService.class);
        service.setRef(() -> "user-provider");
        DubboBootstrap.getInstance()
                .application(new ApplicationConfig(USER_APP))
                .registry(new RegistryConfig(REGISTRY_URL))
                .protocol(new ProtocolConfig("dubbo", 23983))
                .service(service)
                .start()
                .await();
    }

    private static void publishRule() throws Exception {
        InputStream input = RouterChainReproMain.class.getClassLoader().getResourceAsStream("route-chain-repro.yaml");
        if (input == null) {
            throw new IllegalStateException("route-chain-repro.yaml is missing");
        }
        String rule;
        try (BufferedReader reader = new BufferedReader(new InputStreamReader(input, StandardCharsets.UTF_8))) {
            rule = reader.lines().collect(Collectors.joining("\n"));
        }
        ConfigService config = NacosFactory.createConfigService(nacosProperties());
        try {
            if (!config.publishConfig(RULE_DATA_ID, "dubbo", rule)) {
                throw new IllegalStateException("Nacos rejected " + RULE_DATA_ID);
            }
        } finally {
            config.shutDown();
        }
        System.out.println("Published " + RULE_DATA_ID + " to Nacos group dubbo");
    }

    private static Properties nacosProperties() {
        Properties properties = new Properties();
        properties.setProperty(PropertyKeyConst.SERVER_ADDR, nacosAddress());
        String username = System.getenv("NACOS_USERNAME");
        String password = System.getenv("NACOS_PASSWORD");
        if (username != null && !username.isEmpty()) {
            properties.setProperty(PropertyKeyConst.USERNAME, username);
        }
        if (password != null && !password.isEmpty()) {
            properties.setProperty(PropertyKeyConst.PASSWORD, password);
        }
        return properties;
    }

    private static String nacosAddress() {
        String address = System.getenv("NACOS_SERVER_ADDR");
        return address == null || address.isEmpty() ? "127.0.0.1:18848" : address;
    }

    private static void runConsumer(String mode, int requests) {
        boolean multi = !"single".equals(mode);
        boolean startup = "startup".equals(mode) || "fixed-startup".equals(mode);
        boolean fixed = "fixed".equals(mode) || "fixed-startup".equals(mode);
        if (requests < 2) {
            throw new IllegalArgumentException("At least two requests are needed to observe round-robin leakage");
        }
        System.out.println("Java version=" + System.getProperty("java.version"));
        System.out.println("AppStateRouterFactory source="
                + AppStateRouterFactory.class.getProtectionDomain().getCodeSource().getLocation());
        ConfigCenterConfig configCenter = new ConfigCenterConfig();
        configCenter.setAddress("nacos://" + nacosAddress());
        String username = System.getenv("NACOS_USERNAME");
        String password = System.getenv("NACOS_PASSWORD");
        if (username != null && !username.isEmpty()) {
            configCenter.setUsername(username);
        }
        if (password != null && !password.isEmpty()) {
            configCenter.setPassword(password);
        }

        ReferenceConfig<OrderProbeService> orderRef = new ReferenceConfig<>();
        orderRef.setInterface(OrderProbeService.class);
        orderRef.setLoadbalance("roundrobin");
        orderRef.setRetries(0);

        ReferenceConfig<UserProbeService> userRef = null;
        if (multi) {
            userRef = new ReferenceConfig<>();
            userRef.setInterface(UserProbeService.class);
            userRef.setRetries(0);
        }

        DubboBootstrap bootstrap = DubboBootstrap.getInstance()
                .application(new ApplicationConfig("route-chain-repro-consumer"))
                .configCenter(configCenter)
                .registry(new RegistryConfig(REGISTRY_URL))
                .reference(orderRef);
        if (startup) {
            bootstrap.reference(userRef);
        }
        bootstrap.start();
        try {
            OrderProbeService order = bootstrap.getCache().get(orderRef);
            if (!startup) {
                for (int i = 0; i < 4; i++) {
                    String result = order.instance();
                    if (!"order-v2".equals(result)) {
                        throw new IllegalStateException("Order rule did not work before second service was referenced: " + result);
                    }
                }
                System.out.println("Before second service: 4/4 Order RPCs reached order-v2");
            }

            if (multi) {
                if (!startup) {
                    bootstrap.reference(userRef);
                }
                UserProbeService user = bootstrap.getCache().get(userRef);
                if (!"user-provider".equals(user.ping())) {
                    throw new IllegalStateException("Second service is not reachable");
                }
                System.out.println("Second service reference initialized and invoked in the same JVM");
            }

            Map<String, Integer> counts = new LinkedHashMap<>();
            for (int i = 0; i < requests; i++) {
                String id = order.instance();
                counts.merge(id, 1, Integer::sum);
            }
            System.out.println("Mode=" + mode + " Order RPC results=" + counts);
            if (!multi && !counts.keySet().equals(java.util.Collections.singleton("order-v2"))) {
                throw new IllegalStateException("Single-service control failed: " + counts);
            }
            if (!fixed && multi && !counts.containsKey("order-v1")) {
                throw new IllegalStateException("The multi-service bug did not reproduce: " + counts);
            }
            if (fixed && !counts.keySet().equals(java.util.Collections.singleton("order-v2"))) {
                throw new IllegalStateException("The router factory patch did not restore Order routing: " + counts);
            }
            System.out.println(fixed ? "PATCHED CONTROL PASSED"
                    : multi ? "BUG REPRODUCED" : "SINGLE-SERVICE CONTROL PASSED");
        } finally {
            bootstrap.stop();
        }
    }
}