purpur-ignis/src/main/java/games/dmg/spigotignis/FileStateRepository.java

package games.dmg.spigotignis;

import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.nio.channels.FileChannel;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.StandardCopyOption;
import java.nio.file.StandardOpenOption;
import java.util.HashMap;
import java.util.Map;
import java.util.UUID;

/** Versioned data, never Java object deserialization. Call from the persistence worker only. */
public final class FileStateRepository implements StateRepository {
    private static final int MAGIC = 0x49474e49;
    private static final int VERSION = 2;
    private static final int MAX_RECORDS = 100_000;
    private final Path path;

    public FileStateRepository(Path path) { this.path = path.toAbsolutePath(); }

    @Override
    public IgnisState load() throws IOException {
        if (Files.notExists(path)) return IgnisState.empty();
        try (var input = new DataInputStream(new BufferedInputStream(Files.newInputStream(path)))) {
            if (input.readInt() != MAGIC) throw new IOException("Unsupported Ignis state header.");
            int version = input.readInt();
            if (version < 1 || version > VERSION) throw new IOException("Unsupported Ignis state version.");
            var limits = new ProgressLimits(input.readLong(), input.readLong(), input.readLong());
            Map<UUID, PlayerState> players = new HashMap<>();
            int count = count(input, MAX_RECORDS);
            for (int i = 0; i < count; i++) {
                UUID id = uuid(input);
                Map<Track, TrackState> tracks = new HashMap<>();
                int trackCount = count(input, Track.values().length);
                for (int j = 0; j < trackCount; j++) {
                    Track track = Track.valueOf(input.readUTF());
                    var state = new TrackState(input.readLong(), flag(input),
                            flag(input), input.readLong());
                    if (tracks.putIfAbsent(track, state) != null) throw new IOException("Duplicate track.");
                }
                Map<UUID, Track> pending = new HashMap<>();
                int rewards = count(input, MAX_RECORDS);
                for (int j = 0; j < rewards; j++) {
                    if (pending.putIfAbsent(uuid(input), Track.valueOf(input.readUTF())) != null) {
                        throw new IOException("Duplicate pending reward.");
                    }
                }
                Map<UUID, Track> issued = new HashMap<>(pending);
                if (version >= 2) {
                    issued.clear();
                    int receipts = count(input, MAX_RECORDS);
                    for (int j = 0; j < receipts; j++) {
                        if (issued.putIfAbsent(uuid(input), Track.valueOf(input.readUTF())) != null) {
                            throw new IOException("Duplicate issued receipt.");
                        }
                    }
                    if (!issued.entrySet().containsAll(pending.entrySet())) {
                        throw new IOException("Pending reward has no matching issuance receipt.");
                    }
                }
                if (players.putIfAbsent(id, new PlayerState(tracks, pending, issued)) != null) {
                    throw new IOException("Duplicate player.");
                }
            }
            if (input.read() != -1) throw new IOException("Trailing Ignis state data.");
            return new IgnisState(limits, players);
        } catch (IllegalArgumentException exception) {
            throw new IOException("Invalid Ignis state; original file retained.", exception);
        }
    }

    @Override
    public void save(IgnisState state) throws IOException {
        if (state.players().size() > MAX_RECORDS
                || state.players().values().stream().anyMatch(p -> p.pendingRewards().size() > MAX_RECORDS
                        || p.issuedRewards().size() > MAX_RECORDS)) {
            throw new IOException("Ignis state exceeds supported record count.");
        }
        Files.createDirectories(path.getParent());
        Path temporary = Files.createTempFile(path.getParent(), ".ignis-", ".tmp");
        try {
            try (var output = new DataOutputStream(new BufferedOutputStream(Files.newOutputStream(temporary)))) {
                output.writeInt(MAGIC);
                output.writeInt(VERSION);
                output.writeLong(state.limits().invisibilityMillis());
                output.writeLong(state.limits().pearls());
                output.writeLong(state.limits().darknessMillis());
                output.writeInt(state.players().size());
                for (var entry : state.players().entrySet()) {
                    uuid(output, entry.getKey());
                    output.writeInt(entry.getValue().tracks().size());
                    for (var track : entry.getValue().tracks().entrySet()) {
                        output.writeUTF(track.getKey().name());
                        output.writeLong(track.getValue().progress());
                        output.writeBoolean(track.getValue().unlocked());
                        output.writeBoolean(track.getValue().activated());
                        output.writeLong(track.getValue().generation());
                    }
                    output.writeInt(entry.getValue().pendingRewards().size());
                    for (var reward : entry.getValue().pendingRewards().entrySet()) {
                        uuid(output, reward.getKey());
                        output.writeUTF(reward.getValue().name());
                    }
                    output.writeInt(entry.getValue().issuedRewards().size());
                    for (var receipt : entry.getValue().issuedRewards().entrySet()) {
                        uuid(output, receipt.getKey());
                        output.writeUTF(receipt.getValue().name());
                    }
                }
            }
            try (var file = FileChannel.open(temporary, StandardOpenOption.WRITE)) { file.force(true); }
            // Deliberately fail rather than fall back to a partial non-atomic replacement.
            Files.move(temporary, path, StandardCopyOption.ATOMIC_MOVE, StandardCopyOption.REPLACE_EXISTING);
        } finally {
            Files.deleteIfExists(temporary);
        }
    }

    private static boolean flag(DataInputStream input) throws IOException {
        int value = input.readUnsignedByte();
        if (value > 1) throw new IOException("Invalid state flag.");
        return value == 1;
    }

    private static int count(DataInputStream input, int maximum) throws IOException {
        int result = input.readInt();
        if (result < 0 || result > maximum) throw new IOException("Invalid record count.");
        return result;
    }

    private static UUID uuid(DataInputStream input) throws IOException {
        return new UUID(input.readLong(), input.readLong());
    }

    private static void uuid(DataOutputStream output, UUID id) throws IOException {
        output.writeLong(id.getMostSignificantBits());
        output.writeLong(id.getLeastSignificantBits());
    }
}