/*
 * Copyright (C) 2023 The Android Open Source Project
 *
 * 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 android.net.apf;

import static android.net.apf.ApfCounterTracker.Counter.APF_PROGRAM_ID;
import static android.net.apf.ApfCounterTracker.Counter.FILTER_AGE_SECONDS;

import android.annotation.NonNull;
import android.util.ArrayMap;
import android.util.Log;
import android.util.Pair;

import com.android.internal.annotations.VisibleForTesting;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;

/**
 * Counter class for {@code ApfFilter}.
 *
 * @hide
 */
public class ApfCounterTracker {
    /**
     * APF packet counters.
     *
     * Packet counters are 32bit big-endian values, and allocated near the end of the APF data
     * buffer, using negative byte offsets, where -4 is equivalent to maximumApfProgramSize - 4,
     * the last writable 32bit word.
     */
    public enum Counter {
        RESERVED_OOB,  // Points to offset 0 from the end of the buffer (out-of-bounds)
        ENDIANNESS,              // APFv6 interpreter stores 0x12345678 here
        TOTAL_PACKETS,           // hardcoded in APFv6 interpreter
        PASSED_ALLOCATE_FAILURE, // hardcoded in APFv6 interpreter
        PASSED_TRANSMIT_FAILURE, // hardcoded in APFv6 interpreter
        CORRUPT_DNS_PACKET,      // hardcoded in APFv6 interpreter
        EXCEPTIONS,              // hardcoded in APFv6.1 interpreter
        FILTER_AGE_SECONDS,
        FILTER_AGE_16384THS,
        APF_VERSION,
        APF_PROGRAM_ID,
        // The counter sequence should keep the same as ApfSessionInfoMetrics.java
        PASSED_ARP_BROADCAST_REPLY,  // see also MIN_PASS_COUNTER below.
        PASSED_ARP_REQUEST,
        PASSED_ARP_UNICAST_REPLY,
        PASSED_DHCP,
        PASSED_DUE_TO_REPLY_OVER_MTU,
        PASSED_ETHER_OUR_SRC_MAC,
        PASSED_IPV4,
        PASSED_IPV4_FROM_DHCPV4_SERVER,
        PASSED_IPV4_UNICAST,
        PASSED_IPV6_HOPOPTS,
        PASSED_IPV6_ICMP,
        PASSED_IPV6_NON_ICMP,
        PASSED_IPV6_UNICAST_NON_ICMP,
        PASSED_LOW_POWER_STANDBY_MAGIC_PACKET,
        PASSED_NON_IP_UNICAST,
        PASSED_MDNS,
        PASSED_RA,  // see also MAX_PASS_COUNTER below
        DROPPED_ETH_BROADCAST,  // see also MIN_DROP_COUNTER below
        DROPPED_ETHER_OUR_SRC_MAC,
        DROPPED_RA,
        DROPPED_IPV4_L2_BROADCAST,
        DROPPED_IPV4_BROADCAST_ADDR,
        DROPPED_IPV4_BROADCAST_NET,
        DROPPED_IPV4_ICMP_INVALID,
        DROPPED_IPV4_MULTICAST,
        DROPPED_IPV4_NON_DHCP4,
        DROPPED_IPV4_PING_REQUEST_REPLIED,
        DROPPED_IPV6_ICMP6_ECHO_REQUEST_INVALID,
        DROPPED_IPV6_ICMP6_ECHO_REQUEST_REPLIED,
        DROPPED_IPV6_ROUTER_SOLICITATION,
        DROPPED_IPV6_MLD_INVALID,
        DROPPED_IPV6_MLD_REPORT,
        DROPPED_IPV6_MLD_V1_GENERAL_QUERY_REPLIED,
        DROPPED_IPV6_MLD_V2_GENERAL_QUERY_REPLIED,
        DROPPED_IPV6_MULTICAST_NA,
        DROPPED_IPV6_NON_ICMP_MULTICAST,
        DROPPED_IPV6_NS_INVALID,
        DROPPED_IPV6_NS_OTHER_HOST,
        DROPPED_IPV6_NS_REPLIED_NON_DAD,
        DROPPED_802_3_FRAME,
        DROPPED_ETHERTYPE_NOT_ALLOWED,
        DROPPED_IPV4_KEEPALIVE_ACK,
        DROPPED_IPV4_NATT_KEEPALIVE,
        DROPPED_MDNS,
        DROPPED_MDNS_REPLIED,
        DROPPED_NON_UNICAST_TDLS,
        DROPPED_IPV4_TCP_PORT7_UNICAST,
        DROPPED_ARP_NON_IPV4,
        DROPPED_ARP_OTHER_HOST,
        DROPPED_ARP_REPLY_SPA_NO_HOST,
        DROPPED_ARP_REQUEST_REPLIED,
        DROPPED_ARP_UNKNOWN,
        DROPPED_ARP_V6_ONLY,
        DROPPED_IGMP_V2_GENERAL_QUERY_REPLIED,
        DROPPED_IGMP_V3_GENERAL_QUERY_REPLIED,
        DROPPED_IGMP_INVALID,
        DROPPED_IGMP_REPORT,
        DROPPED_GARP_REPLY;  // see also MAX_DROP_COUNTER below

        // Cached count of the number of valid Counter enum values, excluding
        // RESERVED_OOB. This avoids redundant and costly values() queries.
        public static final int NUM_VALID_COUNTERS = values().length - 1;

        /**
         * Returns the negative byte offset from the end of the APF data segment for
         * a given counter.
         */
        public int offset() {
            return -this.value() * 4;  // Currently, all counters are 32bit long.
        }

        /**
         * Returns the counter sequence number from the end of the APF data segment for
         * a given counter.
         */
        public int value() {
            return this.ordinal();
        }

        /**
         * Returns the total size of the data segment in bytes.
         */
        public static int totalSize() {
            return NUM_VALID_COUNTERS * 4;
        }

        /**
         * Returns the counter enum based on the offset.
         */
        @VisibleForTesting(visibility = VisibleForTesting.Visibility.PACKAGE)
        public static Counter getCounterEnumFromOffset(int offset) {
            for (Counter cnt : Counter.values()) {
                if (cnt.offset() == offset) {
                    return cnt;
                }
            }
            return RESERVED_OOB;
        }

        private void checkCounterRange(Counter lowerBound, Counter upperBound) {
            if (value() < lowerBound.value() || value() > upperBound.value()) {
                throw new IllegalArgumentException(
                        String.format("Counter %s, is not in range [%s, %s]", this,
                                lowerBound, upperBound));
            }
        }

        /**
         * Return the label such that if we jump to it, the counter will be increased by 1 and
         * the packet will be passed.
         */
        public short getJumpPassLabel() {
            checkCounterRange(MIN_PASS_COUNTER, MAX_PASS_COUNTER);
            return (short) (2 * this.value());
        }

        /**
         * Return the label such that if we jump to it, the counter will be increased by 1 and
         * the packet will be dropped.
         */
        public short getJumpDropLabel() {
            checkCounterRange(MIN_DROP_COUNTER, MAX_DROP_COUNTER);
            return (short) (2 * this.value() + 1);
        }
    }

    public static final Counter MIN_DROP_COUNTER = Counter.DROPPED_ETH_BROADCAST;
    public static final Counter MAX_DROP_COUNTER = Counter.DROPPED_GARP_REPLY;
    public static final Counter MIN_PASS_COUNTER = Counter.PASSED_ARP_BROADCAST_REPLY;
    public static final Counter MAX_PASS_COUNTER = Counter.PASSED_RA;

    private static final String TAG = ApfCounterTracker.class.getSimpleName();

    private final List<Counter> mCounterList;
    // Store the counters' value
    private final Map<Counter, Long> mCounters = new ArrayMap<>();

    public ApfCounterTracker() {
        Counter[] counters = Counter.values();
        mCounterList = Arrays.asList(counters).subList(1, counters.length);
    }

    /**
     * Get the value of a counter from APF data.
     */
    public static long getCounterValue(byte[] data, Counter counter)
            throws ArrayIndexOutOfBoundsException {
        int offset = data.length + Counter.ENDIANNESS.offset();
        int endianness = 0;
        for (int i = 0; i < 4; i++) {
            endianness = endianness << 8 | (data[offset + i] & 0xff);
        }
        // Follow the same wrap-around addressing scheme of the interpreter.
        offset = data.length + counter.offset();

        boolean isBe = true;
        switch (endianness) {
            case 0:
            case 0x12345678:
                isBe = true;
                break;
            case 0x78563412:
                isBe = false;
                break;
            default:
                Log.wtf(TAG, "Unknown endianness: 0x" + Integer.toHexString(endianness));
        }

        // Decode 32bit big-endian integer into a long so we can count up beyond 2^31.
        long value = 0;
        for (int i = 0; i < 4; i++) {
            value = value << 8 | (data[offset + (isBe ? i : 3 - i)] & 0xff);
        }
        return value;
    }

    /**
     * Update counters from APF data.
     */
    public void updateCountersFromData(byte[] data) {
        if (data == null) return;
        for (Counter counter : mCounterList) {
            long value;
            try {
                value = getCounterValue(data, counter);
            } catch (ArrayIndexOutOfBoundsException e) {
                value = 0;
            }
            long oldValue = mCounters.getOrDefault(counter, 0L);
            // All counters are incremental
            if (value > oldValue) {
                mCounters.put(counter, value);
            }
        }
    }

    /**
     * Get counters map.
     */
    public Map<Counter, Long> getCounters() {
        return mCounters;
    }

    /**
     * Clear all counters.
     */
    public void clearCounters() {
        mCounters.clear();
    }

    /**
     * Return readable counter for testing purposes.
     */
    public List<Pair<Counter, String>> dumpCountersFromData(
            @NonNull byte[] data,
            int filterAgeSeconds,
            int numProgramUpdates,
            int apfVersionSupported) throws ArrayIndexOutOfBoundsException {
        List<Pair<Counter, String>> counterList = new ArrayList<>();
        Counter[] counters = Counter.values();
        long counterFilterAgeSeconds =
                getCounterValue(data, FILTER_AGE_SECONDS);
        long counterApfProgramId =
                getCounterValue(data, APF_PROGRAM_ID);

        for (Counter c : Arrays.asList(counters).subList(1, counters.length)) {
            long value = getCounterValue(data, c);

            String note = "";
            boolean checkValueIncreases = true;
            switch (c) {
                case FILTER_AGE_SECONDS:
                    checkValueIncreases = false;
                    if (value != counterFilterAgeSeconds) {
                        note = " [ERROR: impossible]";
                    } else if (counterApfProgramId < numProgramUpdates) {
                        note = " [IGNORE: obsolete program]";
                    } else if (value > filterAgeSeconds) {
                        long offset = value - filterAgeSeconds;
                        note = " [ERROR: in the future by " + offset + "s]";
                    }
                    break;
                case FILTER_AGE_16384THS:
                    if (apfVersionSupported > BaseApfGenerator.APF_VERSION_4) {
                        checkValueIncreases = false;
                        if (value % 16384 == 0) {
                            // valid, but unlikely
                            note = " [INFO: zero fractional portion]";
                        }
                        if (value / 16384 != counterFilterAgeSeconds) {
                            // should not be able to happen
                            note = " [ERROR: mismatch with FILTER_AGE_SECONDS]";
                        }
                    } else if (value != 0) {
                        note = " [UNEXPECTED: APF<=4, yet non-zero]";
                    }
                    break;
                case APF_PROGRAM_ID:
                    if (value != counterApfProgramId) {
                        note = " [ERROR: impossible]";
                    } else if (value < numProgramUpdates) {
                        note = " [WARNING: OBSOLETE PROGRAM]";
                    } else if (value > numProgramUpdates) {
                        note = " [ERROR: INVALID FUTURE ID]";
                    }
                    break;
                default:
                    break;
            }

            // Only print non-zero counters (or those with a note)
            if (value != 0 || !note.equals("")) {
                counterList.add(new Pair<>(c, value + note));
            }

            if (checkValueIncreases) {
                // If the counter's value decreases, it may have been cleaned up or there
                // may be a bug.
                long oldValue = getCounters().getOrDefault(c, 0L);
                if (value < oldValue) {
                    Log.e(TAG, String.format(
                            "Apf Counter: %s unexpectedly decreased. oldValue: %d. "
                            + "newValue: %d", c.toString(), oldValue, value));
                }
            }
        }

        return counterList;
    }
}