/*_############################################################################
  _##
  _##  SNMP4J-Agent 3 - AgentStateDeserializationFilter.java
  _##
  _##  Copyright (C) 2005-2026  Frank Fock (SNMP4J.org)
  _##
  _##  Licensed under the Apache License, Version 2.0 (the "License");
  _##  you may not use this file except in compliance with the License.
  _##  You may obtain a copy of the License at
  _##
  _##      http://www.apache.org/licenses/LICENSE-2.0
  _##
  _##  Unless required by applicable law or agreed to in writing, software
  _##  distributed under the License is distributed on an "AS IS" BASIS,
  _##  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  _##  See the License for the specific language governing permissions and
  _##  limitations under the License.
  _##
  _##########################################################################*/

package org.snmp4j.agent.io;

import org.snmp4j.log.LogAdapter;
import org.snmp4j.log.LogFactory;
import org.snmp4j.security.SecretOctetString;
import org.snmp4j.smi.*;

import java.io.ObjectInputFilter;
import java.util.Collection;
import java.util.Collections;
import java.util.HashSet;
import java.util.Set;

/**
 * The {@code AgentStateDeserializationFilter} is a strict {@link ObjectInputFilter} that protects the
 * deserialization of persistent agent state (as read by {@link DefaultMOInput}) against Java
 * deserialization attacks. It allows <em>only</em> the exact set of classes that {@link DefaultMOOutput}
 * serializes (and the transitive field types of those classes) to be deserialized and rejects every
 * other class. In addition, it bounds the depth of the object graph as well as array length, the number
 * of object references and the number of stream bytes per object graph as a denial-of-service backstop.
 * <p>
 * The allowed classes are matched by exact identity, i.e., registering a non-final class does <em>not</em>
 * implicitly allow its subclasses. Custom {@link Variable} or {@link Context} subclasses that are not part
 * of the default allowlist have to be registered explicitly through
 * {@link #AgentStateDeserializationFilter(Collection)}.
 * <p>
 * This filter only uses Java&nbsp;9 serialization filter API.
 *
 * @author Frank Fock
 * @version 3.12.0
 * @since 3.12.0
 */
public final class AgentStateDeserializationFilter implements ObjectInputFilter {

    private static final LogAdapter logger = LogFactory.getLogger(AgentStateDeserializationFilter.class);

    /**
     * Maximum allowed depth of the deserialized object graph. The deepest legitimate graph written by
     * {@link DefaultMOOutput} is about five levels deep
     * ({@code IndexedVariables -> Variable[] -> IpAddress -> InetAddress -> byte[]}); the limit adds
     * headroom while still rejecting deeply nested gadget chains.
     */
    public static final long MAX_DEPTH = 10;

    /**
     * Maximum allowed array length. Set generously: legitimate arrays (table column values, OID
     * sub-identifiers, IP address bytes) are tiny, but a stored binary {@code OctetString} blob may be
     * sizable. The limit only guards against multi-gigabyte allocation attacks.
     */
    public static final long MAX_ARRAY = 1_000_000L;

    /**
     * Maximum allowed number of object references in a single object graph (handle table size).
     */
    public static final long MAX_REFS = 5_000_000L;

    /**
     * Maximum allowed number of stream bytes consumed for a single object graph (i.e. per
     * {@code readObject} call), as a denial-of-service backstop.
     */
    public static final long MAX_BYTES = 200_000_000L;

    /**
     * The set of classes that are always allowed by this filter. This is the exact transitive closure of
     * the types serialized by {@link DefaultMOOutput} including all serializable super classes (which
     * each emit their own class descriptor) and the referenced JDK types.
     */
    private static final Set<Class<?>> DEFAULT_ALLOWED = Set.of(
            // -- Agent IO model types (top level objects written by DefaultMOOutput)
            Context.class,
            IndexedVariables.class,
            MOInfo.class,
            Sequence.class,
            // -- SMI base/abstract/intermediate classes (each emits its own class descriptor)
            Variable.class,            // interface; also allows Variable[] (array component type)
            AbstractVariable.class,
            UnsignedInteger32.class,
            OctetString.class,
            SecretOctetString.class,
            SMIAddress.class,
            IpAddress.class,
            TransportIpAddress.class,
            // -- SMI concrete Variable implementations
            Integer32.class,
            Counter32.class,
            Gauge32.class,
            TimeTicks.class,
            Counter64.class,
            Opaque.class,
            BitString.class,
            OID.class,
            Null.class,
            TcpAddress.class,
            UdpAddress.class,
            TlsAddress.class,
            DtlsAddress.class,
            // -- JDK types referenced by the model/SMI fields
            String.class,
            java.net.InetAddress.class,
            java.net.Inet4Address.class,
            java.net.Inet6Address.class);

    private final Set<Class<?>> allowed;

    /**
     * Creates a strict filter that allows only the default set of agent state classes.
     */
    public AgentStateDeserializationFilter() {
        this.allowed = DEFAULT_ALLOWED;
    }

    /**
     * Creates a strict filter that allows the default set of agent state classes plus the supplied
     * additional classes. The additional classes are matched by exact identity (subclasses are not
     * implicitly allowed).
     *
     * @param additionalAllowed
     *    additional classes to allow for deserialization, or {@code null}/empty for the default set only.
     */
    public AgentStateDeserializationFilter(Collection<Class<?>> additionalAllowed) {
        if ((additionalAllowed == null) || additionalAllowed.isEmpty()) {
            this.allowed = DEFAULT_ALLOWED;
        }
        else {
            Set<Class<?>> merged = new HashSet<>(DEFAULT_ALLOWED);
            merged.addAll(additionalAllowed);
            this.allowed = Collections.unmodifiableSet(merged);
        }
    }

    @Override
    public Status checkInput(FilterInfo filterInfo) {
        // Resource limit guards. These are also evaluated for callbacks without a class
        // (filterInfo.serialClass() == null), e.g. array length checks.
        if ((filterInfo.depth() > MAX_DEPTH) ||
                (filterInfo.references() > MAX_REFS) ||
                (filterInfo.streamBytes() > MAX_BYTES) ||
                ((filterInfo.arrayLength() >= 0) && (filterInfo.arrayLength() > MAX_ARRAY))) {
            if (logger.isDebugEnabled()) {
                logger.debug("Rejecting deserialization due to exceeded limit: depth=" + filterInfo.depth() +
                        ", references=" + filterInfo.references() + ", streamBytes=" + filterInfo.streamBytes() +
                        ", arrayLength=" + filterInfo.arrayLength());
            }
            return Status.REJECTED;
        }
        Class<?> clazz = filterInfo.serialClass();
        if (clazz == null) {
            // No class to check (e.g. a limit-only callback): leave the decision undecided.
            return Status.UNDECIDED;
        }
        if (clazz.isArray()) {
            // The only arrays written by DefaultMOOutput are primitive arrays (byte[], int[]) and the
            // Variable[] of IndexedVariables. Reject arrays of any other (even otherwise allowed) type.
            Class<?> component = clazz;
            while (component.isArray()) {
                component = component.getComponentType();
            }
            if (component.isPrimitive() || (component == Variable.class)) {
                return Status.ALLOWED;
            }
            logger.warn("Rejecting deserialization of disallowed array class: " + clazz.getName());
            return Status.REJECTED;
        }
        if (allowed.contains(clazz)) {
            return Status.ALLOWED;
        }
        logger.warn("Rejecting deserialization of disallowed class: " + clazz.getName());
        return Status.REJECTED;
    }
}
