package com.cmvr.quic.gateway.server;

import io.netty.buffer.ByteBuf;
import io.netty.handler.codec.quic.QuicTokenHandler;

import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import java.net.InetSocketAddress;
import java.security.GeneralSecurityException;
import java.security.MessageDigest;
import java.security.SecureRandom;

/** 使用进程级随机密钥签发短期QUIC Retry Token。 */
final class HmacQuicTokenHandler implements QuicTokenHandler {

    private static final int VERSION = 1;
    private static final int MAC_BYTES = 16;
    private static final int MAX_CONNECTION_ID_BYTES = 20;
    private static final int MAX_TOKEN_BYTES = 1 + 8 + 1 + 16 + MAC_BYTES + MAX_CONNECTION_ID_BYTES;
    private static final long TOKEN_TTL_MS = 30000L;

    private final SecretKeySpec key;

    HmacQuicTokenHandler() {
        byte[] secret = new byte[32];
        new SecureRandom().nextBytes(secret);
        key = new SecretKeySpec(secret, "HmacSHA256");
    }

    @Override
    public boolean writeToken(ByteBuf out, ByteBuf dcid, InetSocketAddress address) {
        byte[] ip = address.getAddress().getAddress();
        int timestampIndex = out.writerIndex() + 1;
        out.writeByte(VERSION);
        out.writeLong(System.currentTimeMillis());
        out.writeByte(ip.length);
        out.writeBytes(ip);
        int macIndex = out.writerIndex();
        out.writeZero(MAC_BYTES);
        int dcidIndex = out.writerIndex();
        out.writeBytes(dcid, dcid.readerIndex(), dcid.readableBytes());
        byte[] signature = sign(out, timestampIndex, 8 + 1 + ip.length, dcidIndex, dcid.readableBytes());
        out.setBytes(macIndex, signature, 0, MAC_BYTES);
        return true;
    }

    @Override
    public int validateToken(ByteBuf token, InetSocketAddress address) {
        int start = token.readerIndex();
        if (token.readableBytes() < 1 + 8 + 1 + 4 + MAC_BYTES + 1
                || token.getUnsignedByte(start) != VERSION) {
            return -1;
        }
        long issuedAt = token.getLong(start + 1);
        long age = System.currentTimeMillis() - issuedAt;
        int ipLength = token.getUnsignedByte(start + 9);
        byte[] actualIp = address.getAddress().getAddress();
        int macIndex = start + 10 + ipLength;
        int dcidIndex = macIndex + MAC_BYTES;
        int dcidLength = token.writerIndex() - dcidIndex;
        if (age < 0 || age > TOKEN_TTL_MS || ipLength != actualIp.length
                || dcidLength <= 0 || dcidLength > MAX_CONNECTION_ID_BYTES
                || dcidIndex > token.writerIndex()) {
            return -1;
        }
        for (int i = 0; i < ipLength; i++) {
            if (token.getByte(start + 10 + i) != actualIp[i]) return -1;
        }
        byte[] expected = sign(token, start + 1, 8 + 1 + ipLength, dcidIndex, dcidLength);
        byte[] supplied = new byte[MAC_BYTES];
        token.getBytes(macIndex, supplied);
        byte[] truncated = new byte[MAC_BYTES];
        System.arraycopy(expected, 0, truncated, 0, MAC_BYTES);
        return MessageDigest.isEqual(truncated, supplied) ? dcidIndex - start : -1;
    }

    @Override
    public int maxTokenLength() {
        return MAX_TOKEN_BYTES;
    }

    private byte[] sign(ByteBuf source, int firstIndex, int firstLength,
                        int secondIndex, int secondLength) {
        try {
            Mac mac = Mac.getInstance("HmacSHA256");
            mac.init(key);
            byte[] first = new byte[firstLength];
            byte[] second = new byte[secondLength];
            source.getBytes(firstIndex, first);
            source.getBytes(secondIndex, second);
            mac.update(first);
            mac.update(second);
            return mac.doFinal();
        } catch (GeneralSecurityException error) {
            throw new IllegalStateException("无法计算QUIC Retry Token签名", error);
        }
    }
}
