/*
 * Copyright (C) 2026  Tim Zimmermann <tim@aglais.org>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

package org.aglais.zms.server.codec.quic;

import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import io.netty.handler.codec.quic.Quic;
import io.netty.handler.codec.quic.QuicTokenHandler;
import io.netty.util.concurrent.FastThreadLocal;

import javax.crypto.KeyGenerator;
import javax.crypto.Mac;
import java.net.InetSocketAddress;
import java.security.Key;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;

/**
 * Secure {@link QuicTokenHandler} implementation based on a MAC.
 *
 * @author Jonas Konrad
 */
public class SecureQuicTokenHandler implements QuicTokenHandler {
    private static final int MAC_LENGTH = 256 / 8;
    private static final int MAX_CONNECTION_ID_LENGTH = Quic.MAX_CONN_ID_LEN;
    /**
     * Timestamp will be modulo'd by this window size, and included in the MAC. In verification, we
     * check the last two windows, so a given token expires in roughly this time frame.
     */
    private static final long TIMESTAMP_WINDOW_SIZE = 5 * 60 * 1000;

    private final Key key;
    /**
     * Making this non-static is not ideal, but it allows initializing the mac with the key only
     * once.
     */
    private final FastThreadLocal<Mac> macCache = new FastThreadLocal<>() {
        @Override
        protected Mac initialValue() throws Exception {
            Mac mac = Mac.getInstance("HmacSHA256");
            mac.init(key);
            return mac;
        }
    };

    /**
     * Create a new {@link SecureQuicTokenHandler} instance
     */
    public SecureQuicTokenHandler() {
        try {
            key = KeyGenerator.getInstance("HmacSHA256").generateKey();
        } catch (NoSuchAlgorithmException e) {
            throw new AssertionError(e);
        }
    }

    @Override
    public boolean writeToken(ByteBuf out, ByteBuf dcid, InetSocketAddress address) {
        byte[] hash = hash(address, dcid, currentWindowId());
        out.writeBytes(hash);
        out.writeBytes(dcid, dcid.readerIndex(), dcid.readableBytes());
        return true;
    }

    @Override
    public int validateToken(ByteBuf token, InetSocketAddress address) {
        byte[] actual = new byte[MAC_LENGTH];
        token.getBytes(token.readerIndex(), actual);
        ByteBuf dcid = token.slice(token.readerIndex() + MAC_LENGTH, token.readableBytes() - MAC_LENGTH);

        long windowId = currentWindowId();
        byte[] expectedHashNow = hash(address, dcid, windowId);
        byte[] expectedHashPrev = hash(address, dcid, windowId - 1);
        // constant-time comparison
        boolean equalNow = MessageDigest.isEqual(expectedHashNow, actual);
        boolean equalPrev = MessageDigest.isEqual(expectedHashPrev, actual);
        if (equalNow | equalPrev) {
            return MAC_LENGTH;
        } else {
            return -1;
        }
    }

    private byte[] hash(InetSocketAddress address, ByteBuf dcid, long windowId) {
        ByteBuf textToVerify = buildTextToVerify(address, dcid, windowId);
        byte[] hash;
        try {
            Mac mac = macCache.get();
            mac.update(textToVerify.array(), textToVerify.arrayOffset() + textToVerify.readerIndex(), textToVerify.readableBytes());
            hash = mac.doFinal();
        } finally {
            textToVerify.release();
        }
        assert hash.length == MAC_LENGTH;
        return hash;
    }

    private ByteBuf buildTextToVerify(InetSocketAddress address, ByteBuf dcid, long windowId) {
        if (dcid.readableBytes() > MAX_CONNECTION_ID_LENGTH) {
            throw new IllegalArgumentException("Connection ID may not exceed " + MAX_CONNECTION_ID_LENGTH
                    + " bytes length");
        }
        ByteBuf textToVerify = Unpooled.buffer();
        byte[] addressBytes = address.getAddress().getAddress();
        textToVerify.writeByte(addressBytes.length);
        textToVerify.writeBytes(addressBytes);
        textToVerify.writeShort(address.getPort());
        textToVerify.writeByte(dcid.readableBytes());
        textToVerify.writeBytes(dcid, dcid.readerIndex(), dcid.readableBytes());
        textToVerify.writeLong(windowId);
        return textToVerify;
    }

    @Override
    public int maxTokenLength() {
        return MAC_LENGTH + MAX_CONNECTION_ID_LENGTH;
    }

    // overridden in test
    long currentWindowId() {
        return System.currentTimeMillis() / TIMESTAMP_WINDOW_SIZE;
    }
}