DecoFashion/src/main/java/com/razz/dfashion/client/ClientSharedCosmeticUploader.java

package com.razz.dfashion.client;

import com.razz.dfashion.DecoFashion;
import com.razz.dfashion.bbmodel.BbModelParser;
import com.razz.dfashion.bbmodel.Bbmodel;
import com.razz.dfashion.cosmetic.share.BbmodelCodec;
import com.razz.dfashion.cosmetic.share.SharedCosmeticCache;
import com.razz.dfashion.cosmetic.share.packet.UploadCosmeticChunk;
import com.razz.dfashion.skin.SafePngReader;

import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import net.minecraft.client.Minecraft;
import net.minecraft.client.multiplayer.ClientPacketListener;
import net.minecraft.network.protocol.common.ServerboundCustomPayloadPacket;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.Reader;
import java.nio.file.Files;
import java.nio.file.Path;
import java.security.NoSuchAlgorithmException;
import java.util.zip.Deflater;

/**
 * Author-side upload pipeline for shared cosmetics. Runs only on the uploading client.
 *
 * <p>Security-critical: this is the <b>only</b> place in the entire system where
 * {@link BbModelParser} and {@link SafePngReader} see untrusted bytes. Everything
 * downstream (server, receivers) operates on the validated binary produced here and on
 * already-deflated RGBA. The output of this file is what every other client will
 * eventually decode — if this file doesn't reject something, nothing else will either.
 *
 * <p>Anti-DoS: before sending, the server-authoritative {@link
 * ClientSharedCosmeticCache#libraryContains} check short-circuits uploads that the
 * server already has in this player's library. Re-uploads of identical content are
 * therefore zero-bandwidth after the first success.
 */
public final class ClientSharedCosmeticUploader {

    private ClientSharedCosmeticUploader() {}

    public static String upload(Path bbmodelFile, Path textureFile, String displayName) {
        return upload(bbmodelFile, textureFile, displayName, 0, 1);
    }

    public static String upload(Path bbmodelFile, Path textureFile, String displayName,
                                int frametime, int frames) {
        String result = doUpload(bbmodelFile, textureFile, displayName, frametime, frames);
        DecoFashion.LOGGER.info("Shared cosmetic upload: {}", result);
        return result;
    }

    private static String doUpload(Path bbmodelFile, Path textureFile, String displayName,
                                   int frametime, int frames) {
        // Validate the flipbook spec early so we don't waste a parse on an obviously bad input.
        boolean staticFb = (frametime == 0 && frames == 1);
        boolean animatedFb = (frametime > 0
                && frametime <= com.razz.dfashion.cosmetic.share.SharedCosmeticCache.MAX_FLIPBOOK_FRAMETIME
                && frames > 1
                && frames <= com.razz.dfashion.cosmetic.share.SharedCosmeticCache.MAX_FLIPBOOK_FRAMES);
        if (!staticFb && !animatedFb) {
            return "Bad flipbook: frametime=" + frametime + " frames=" + frames;
        }
        if (!Files.isRegularFile(bbmodelFile)) return "Not a file: " + bbmodelFile;
        if (!Files.isRegularFile(textureFile)) return "Not a file: " + textureFile;

        long bbSize;
        long texSize;
        try {
            bbSize = Files.size(bbmodelFile);
            texSize = Files.size(textureFile);
        } catch (IOException ex) {
            return "Size check failed: " + ex.getMessage();
        }
        if (bbSize > 16L * 1024 * 1024) return "bbmodel too large: " + bbSize;
        if (texSize > SafeImageLoader.MAX_PNG_FILE_BYTES) return "PNG too large: " + texSize;

        // 1. Parse bbmodel via BbModelParser (validates everything).
        Bbmodel bbmodel;
        try (Reader reader = Files.newBufferedReader(bbmodelFile)) {
            bbmodel = BbModelParser.parse(reader);
        } catch (Exception ex) {
            return "bbmodel rejected: " + ex.getMessage();
        }

        // 2. Decode PNG via SafePngReader (memory-safe, RGBA only).
        SafePngReader.Image decoded;
        try {
            byte[] pngBytes = Files.readAllBytes(textureFile);
            decoded = SafePngReader.decode(pngBytes);
        } catch (IOException ex) {
            return "PNG rejected: " + ex.getMessage();
        }

        // 3. Encode bbmodel canonically via BbmodelCodec. This is the form that will cross
        //    the wire, be stored, hashed, and eventually decoded by every receiver.
        byte[] canonicalBbmodelBin;
        ByteBuf out = Unpooled.buffer(1024);
        try {
            BbmodelCodec.CODEC.encode(out, bbmodel);
            if (out.readableBytes() > SharedCosmeticCache.MAX_BBMODEL_BIN_BYTES) {
                return "canonical bbmodel too large: " + out.readableBytes();
            }
            canonicalBbmodelBin = new byte[out.readableBytes()];
            out.readBytes(canonicalBbmodelBin);
        } finally {
            out.release();
        }

        // 4. For flipbooks, the PNG height must divide cleanly by frame count. Reject early
        //    so the server doesn't have to.
        if (frames > 1 && decoded.height % frames != 0) {
            return "Flipbook frames=" + frames + " doesn't divide texture height " + decoded.height;
        }

        // 5. Compute content hash — must match what the server will compute.
        String hash;
        try {
            hash = SharedCosmeticCache.hashContent(canonicalBbmodelBin,
                    decoded.width, decoded.height,
                    frametime, frames,
                    decoded.rgba);
        } catch (NoSuchAlgorithmException ex) {
            return "SHA-256 unavailable";
        }

        // 5. Anti-DoS: if the server already has this hash in our library, skip the wire.
        if (ClientSharedCosmeticCache.libraryContains(hash)) {
            return "Already in library (" + hash.substring(0, 8) + "…); skipped";
        }

        // 6. Deflate raw RGBA.
        byte[] deflated = deflate(decoded.rgba);
        if (deflated.length > SharedCosmeticCache.MAX_DEFLATED_BYTES) {
            return "RGBA deflated too large: " + deflated.length;
        }

        ClientPacketListener conn = Minecraft.getInstance().getConnection();
        if (conn == null) return "Not connected";

        // 7. Assemble wire payload = bbmodel binary || deflated RGBA. Chunk and send.
        int bbmodelLen = canonicalBbmodelBin.length;
        int payloadLen = bbmodelLen + deflated.length;
        int chunkSize = SharedCosmeticCache.MAX_CHUNK_BYTES;
        int total = Math.max(1, (payloadLen + chunkSize - 1) / chunkSize);

        for (int i = 0; i < total; i++) {
            int off = i * chunkSize;
            int len = Math.min(chunkSize, payloadLen - off);
            byte[] slice = sliceConcat(canonicalBbmodelBin, deflated, off, len);
            conn.send(new ServerboundCustomPayloadPacket(
                    new UploadCosmeticChunk(
                            i, total, bbmodelLen,
                            decoded.width, decoded.height,
                            frametime, frames,
                            slice
                    )));
        }

        return "Uploading " + decoded.width + "x" + decoded.height + " ("
                + payloadLen + " bytes, " + total + " chunks, hash "
                + hash.substring(0, 8) + "…) [" + displayName + "]";
    }

    /** Slice bytes from {@code a ‖ b} into a new array of length {@code len} starting at {@code off}. */
    private static byte[] sliceConcat(byte[] a, byte[] b, int off, int len) {
        byte[] out = new byte[len];
        int dst = 0;
        if (off < a.length) {
            int take = Math.min(len, a.length - off);
            System.arraycopy(a, off, out, dst, take);
            dst += take;
            len -= take;
            off += take;
        }
        if (len > 0) {
            int bOff = off - a.length;
            System.arraycopy(b, bOff, out, dst, len);
        }
        return out;
    }

    private static byte[] deflate(byte[] raw) {
        Deflater def = new Deflater(Deflater.BEST_COMPRESSION);
        try {
            def.setInput(raw);
            def.finish();
            ByteArrayOutputStream out = new ByteArrayOutputStream(raw.length / 2 + 64);
            byte[] buf = new byte[64 * 1024];
            while (!def.finished()) {
                int n = def.deflate(buf);
                if (n == 0) break;
                out.write(buf, 0, n);
            }
            return out.toByteArray();
        } finally {
            def.end();
        }
    }
}