storageNode/src/test/java/lion/Service/SubscriptionSnapshotStoreTest.java

package lion.Service;

import com.fasterxml.jackson.databind.ObjectMapper;
import lion.CustomUtil;
import lion.Message.Main.SubscriptionAccountSnapshot;
import lion.Message.Main.SubscriptionBindingSnapshot;
import lion.Message.Main.SubscriptionSnapshotMessage;
import lion.Message.Main.SubscriptionSnapshotPayload;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.io.TempDir;

import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import java.io.ByteArrayOutputStream;
import java.nio.charset.StandardCharsets;
import java.nio.file.Path;
import java.security.MessageDigest;
import java.util.Base64;
import java.util.HexFormat;
import java.util.zip.GZIPOutputStream;

import static org.junit.jupiter.api.Assertions.*;

class SubscriptionSnapshotStoreTest {
    private static final String SECRET = "snapshot-test-secret";
    private final ObjectMapper mapper = CustomUtil.objectMapper;

    @Test
    void appliesSnapshotAndServesByHashedPublicKey(@TempDir Path directory) throws Exception {
        SubscriptionSnapshotStore store = new SubscriptionSnapshotStore(directory, SECRET, 3600, 1024 * 1024);
        SubscriptionSnapshotMessage message = message("public-key-1", "v2-content", "clash-content", System.currentTimeMillis());

        assertEquals(SubscriptionSnapshotStore.APPLY_SUCCESS, store.apply(message).code());
        SubscriptionSnapshotStore.Lookup v2 = store.lookup("v2", "public-key-1");
        SubscriptionSnapshotStore.Lookup clash = store.lookup("cat", "public-key-1");
        assertArrayEquals("v2-content".getBytes(StandardCharsets.UTF_8), v2.content());
        assertArrayEquals("clash-content".getBytes(StandardCharsets.UTF_8), clash.content());
        assertEquals(SubscriptionSnapshotStore.APPLY_OLD, store.apply(message).code());
    }

    @Test
    void rejectsTamperedPayloadAndKeepsPreviousSnapshot(@TempDir Path directory) throws Exception {
        SubscriptionSnapshotStore store = new SubscriptionSnapshotStore(directory, SECRET, 3600, 1024 * 1024);
        SubscriptionSnapshotMessage message = message("public-key-1", "v2-content", "clash-content", System.currentTimeMillis());
        assertEquals(SubscriptionSnapshotStore.APPLY_SUCCESS, store.apply(message).code());
        message.setPayloadBase64(Base64.getEncoder().encodeToString("tampered".getBytes(StandardCharsets.UTF_8)));
        assertEquals(SubscriptionSnapshotStore.APPLY_INVALID, store.apply(message).code());
        assertArrayEquals("v2-content".getBytes(StandardCharsets.UTF_8), store.lookup("v2", "public-key-1").content());
    }

    @Test
    void loadsLastGoodSnapshotAfterRestart(@TempDir Path directory) throws Exception {
        SubscriptionSnapshotMessage message = message("public-key-1", "v2-content", "clash-content", System.currentTimeMillis());
        SubscriptionSnapshotStore first = new SubscriptionSnapshotStore(directory, SECRET, 3600, 1024 * 1024);
        assertEquals(SubscriptionSnapshotStore.APPLY_SUCCESS, first.apply(message).code());
        SubscriptionSnapshotStore second = new SubscriptionSnapshotStore(directory, SECRET, 3600, 1024 * 1024);
        second.load();
        assertArrayEquals("clash-content".getBytes(StandardCharsets.UTF_8), second.lookup("cat", "public-key-1").content());
    }

    /**
     * 回归：节点曾按「内容里的 generatedAt」判过期，而 revision 是内容寻址的——
     * 内容长期不变时主站每分钟重发同一 revision 会被判为 APPLY_OLD 直接丢弃，
     * 新鲜度永远停在首次接收那天，于是主站在线且持续同步，备机仍在第 7 天开始 503。
     * 现在过期只取决于「最近一次成功接收主站快照的时刻」。
     */
    @Test
    void repeatedSameRevisionRefreshesFreshnessAndNeverExpires(@TempDir Path directory) throws Exception {
        // 有效期 1 秒，便于在用例内观察「不续期会过期、续期后恢复」。
        SubscriptionSnapshotStore store = new SubscriptionSnapshotStore(directory, SECRET, 1, 1024 * 1024);
        SubscriptionSnapshotMessage message = message("public-key-1", "v2-content", "clash-content", System.currentTimeMillis());

        assertEquals(SubscriptionSnapshotStore.APPLY_SUCCESS, store.apply(message).code());
        assertEquals("ready", store.status().state());

        // 超过有效期且期间没有任何推送 → 过期（证明判定确实生效，不是恒 ready）
        Thread.sleep(1_200);
        assertEquals("expired", store.status().state(), "超过有效期且无推送应判过期");

        // 同 revision 重发：仍是 APPLY_OLD（内容幂等），但必须刷新新鲜度
        assertEquals(SubscriptionSnapshotStore.APPLY_OLD, store.apply(message).code());
        assertEquals("ready", store.status().state(), "重新收到同 revision 后应恢复可用");
        assertNotNull(store.lookup("v2", "public-key-1"), "持续同步期间必须能取到订阅");
    }

    /** 超过有效期后仍未收到任何快照，才判过期并停止分发。 */
    @Test
    void expiresOnlyAfterNoSnapshotArrivesWithinMaxAge(@TempDir Path directory) throws Exception {
        SubscriptionSnapshotStore store = new SubscriptionSnapshotStore(directory, SECRET, 3600, 1024 * 1024);
        // 内容生成于 2 小时前，且重启后再没收到主站推送 → 超期。
        SubscriptionSnapshotMessage message = message("public-key-1", "v2-content", "clash-content",
                System.currentTimeMillis() - 7_200_000L);
        assertEquals(SubscriptionSnapshotStore.APPLY_SUCCESS, store.apply(message).code());

        SubscriptionSnapshotStore reloaded = new SubscriptionSnapshotStore(directory, SECRET, 3600, 1024 * 1024);
        reloaded.load();
        assertEquals("expired", reloaded.status().state(), "重启后未再收到快照，超期应判过期");
        assertNull(reloaded.lookup("v2", "public-key-1"), "过期快照不得继续分发");
    }

    private SubscriptionSnapshotMessage message(String publicKey, String v2, String clash, long generatedAt) throws Exception {
        byte[] v2Bytes = v2.getBytes(StandardCharsets.UTF_8);
        byte[] clashBytes = clash.getBytes(StandardCharsets.UTF_8);
        SubscriptionAccountSnapshot account = new SubscriptionAccountSnapshot();
        account.setAccountId(1);
        account.setEnabled(true);
        account.setV2ContentBase64(Base64.getEncoder().encodeToString(v2Bytes));
        account.setV2Sha256(sha256(v2Bytes));
        account.setClashContentBase64(Base64.getEncoder().encodeToString(clashBytes));
        account.setClashSha256(sha256(clashBytes));
        SubscriptionBindingSnapshot binding = new SubscriptionBindingSnapshot();
        binding.setPublicKeySha256(sha256(publicKey.getBytes(StandardCharsets.UTF_8)));
        binding.setAccountId(1);
        SubscriptionSnapshotPayload payload = new SubscriptionSnapshotPayload();
        payload.setSchemaVersion(1);
        payload.setAccounts(java.util.List.of(account));
        payload.setBindings(java.util.List.of(binding));
        byte[] json = mapper.writeValueAsBytes(payload);
        byte[] compressed = gzip(json);
        SubscriptionSnapshotMessage message = new SubscriptionSnapshotMessage();
        message.setSchemaVersion(1);
        message.setRevision(sha256(json));
        message.setGeneratedAt(generatedAt);
        message.setPayloadBase64(Base64.getEncoder().encodeToString(compressed));
        message.setPayloadSha256(sha256(compressed));
        String input = "1\n" + message.getRevision() + "\n" + generatedAt + "\n" + message.getPayloadSha256();
        Mac mac = Mac.getInstance("HmacSHA256");
        mac.init(new SecretKeySpec(SECRET.getBytes(StandardCharsets.UTF_8), "HmacSHA256"));
        message.setSignature(hex(mac.doFinal(input.getBytes(StandardCharsets.UTF_8))));
        return message;
    }

    private static byte[] gzip(byte[] bytes) throws Exception {
        ByteArrayOutputStream output = new ByteArrayOutputStream();
        try (GZIPOutputStream gzip = new GZIPOutputStream(output)) { gzip.write(bytes); }
        return output.toByteArray();
    }

    private static String sha256(byte[] bytes) throws Exception {
        return hex(MessageDigest.getInstance("SHA-256").digest(bytes));
    }

    private static String hex(byte[] bytes) { return HexFormat.of().formatHex(bytes); }
}