StreamingBase64Decoder

/*
 * Copyright 2020 LINE Corporation
 *
 * LINE Corporation licenses this file to you 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:
 *
 *   https://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 com.linecorp.armeria.common.grpc.protocol;

import javax.annotation.Nullable;

import io.netty.buffer.ByteBuf;
import io.netty.buffer.ByteBufAllocator;
import io.netty.util.ByteProcessor;

/**
 * A stateful Base64decoder. Unlike {@link Base64#getDecoder()}, this decoder does not end when it meets
 * padding('='), but continues to decode until the end of {@link ByteBuf}. If the 
 */
final class StreamingBase64Decoder implements ByteProcessor {

    // Forked from https://github.com/netty/netty/blob/netty-4.1.51.Final/codec
    // /src/main/java/io/netty/handler/codec/base64/Base64.java

    /** The equals sign (=) as a byte. */
    private static final byte EQUALS_SIGN = (byte) '=';

    private static final byte WHITE_SPACE_ENC = -5; // Indicates white space in encoding

    private static final byte EQUALS_SIGN_ENC = -1; // Indicates equals sign in encoding

    private static final byte[] decodabet = {
            -9, -9, -9, -9, -9, -9,
            -9, -9, -9, // Decimal  0 -  8
            -5, -5, // Whitespace: Tab and Linefeed
            -9, -9, // Decimal 11 - 12
            -5, // Whitespace: Carriage Return
            -9, -9, -9, -9, -9, -9, -9, -9, -9, -9, -9, -9, -9, // Decimal 14 - 26
            -9, -9, -9, -9, -9, // Decimal 27 - 31
            -5, // Whitespace: Space
            -9, -9, -9, -9, -9, -9, -9, -9, -9, -9, // Decimal 33 - 42
            62, // Plus sign at decimal 43
            -9, -9, -9, // Decimal 44 - 46
            63, // Slash at decimal 47
            52, 53, 54, 55, 56, 57, 58, 59, 60, 61, // Numbers zero through nine
            -9, -9, -9, // Decimal 58 - 60
            -1, // Equals sign at decimal 61
            -9, -9, -9, // Decimal 62 - 64
            0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, // Letters 'A' through 'N'
            14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, // Letters 'O' through 'Z'
            -9, -9, -9, -9, -9, -9, // Decimal 91 - 96
            26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, // Letters 'a' through 'm'
            39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, // Letters 'n' through 'z'
            -9, -9, -9, -9, -9 // Decimal 123 - 127
            /* -9,-9,-9,-9,-9,-9,-9,-9,-9,-9,-9,-9,-9,     // Decimal 128 - 140
               -9,-9,-9,-9,-9,-9,-9,-9,-9,-9,-9,-9,-9,     // Decimal 141 - 153
               -9,-9,-9,-9,-9,-9,-9,-9,-9,-9,-9,-9,-9,     // Decimal 154 - 166
               -9,-9,-9,-9,-9,-9,-9,-9,-9,-9,-9,-9,-9,     // Decimal 167 - 179
               -9,-9,-9,-9,-9,-9,-9,-9,-9,-9,-9,-9,-9,     // Decimal 180 - 192
               -9,-9,-9,-9,-9,-9,-9,-9,-9,-9,-9,-9,-9,     // Decimal 193 - 205
               -9,-9,-9,-9,-9,-9,-9,-9,-9,-9,-9,-9,-9,     // Decimal 206 - 218
               -9,-9,-9,-9,-9,-9,-9,-9,-9,-9,-9,-9,-9,     // Decimal 219 - 231
               -9,-9,-9,-9,-9,-9,-9,-9,-9,-9,-9,-9,-9,     // Decimal 232 - 244
               -9,-9,-9,-9,-9,-9,-9,-9,-9,-9,-9            // Decimal 245 - 255 */
    };

    private final ByteBufAllocator allocator;
    private final byte[] b4 = new byte[4];
    private int b4Posn;
    private int outBuffPosn;
    @Nullable
    private ByteBuf dest;

    StreamingBase64Decoder(ByteBufAllocator allocator) {
        this.allocator = allocator;
    }

    ByteBuf decode(ByteBuf src) {
        final ByteBuf dest = allocator.buffer(decodedBufferSize(src.readableBytes()));
        this.dest = dest;
        try {
            src.forEachByte(this);
            final ByteBuf result = dest.slice(0, outBuffPosn);

            // Reset except b4 and b4Posn which will be used for the next chunk.
            outBuffPosn = 0;
            this.dest = null;
            return result;
        } catch (Throwable t) {
            dest.release();
            throw t;
        } finally {
            src.release();
        }
    }

    private int decodedBufferSize(int len) {
        return (len + b4Posn) / 4 * 3;
    }

    @Override
    public boolean process(byte value) throws Exception {
        if (value > 0) {
            final byte sbiDecode = decodabet[value];
            if (sbiDecode >= WHITE_SPACE_ENC) { // White space, Equals sign or better
                if (sbiDecode >= EQUALS_SIGN_ENC) { // Equals sign or better
                    b4[b4Posn++] = value;
                    if (b4Posn > 3) { // Quartet built
                        outBuffPosn += decode4to3(b4, dest, outBuffPosn, decodabet);
                        b4Posn = 0;
                    }
                }
                return true;
            }
        }
        throw new IllegalArgumentException(
                "invalid Base64 input character: " + (short) (value & 0xFF) + " (decimal)");
    }

    private static int decode4to3(byte[] src, ByteBuf dest, int destOffset, byte[] decodabet) {
        final byte src0 = src[0];
        final byte src1 = src[1];
        final byte src2 = src[2];
        final int decodedValue;
        if (src2 == EQUALS_SIGN) {
            // Example: Dk==
            try {
                decodedValue = (decodabet[src0] & 0xff) << 2 | (decodabet[src1] & 0xff) >>> 4;
            } catch (IndexOutOfBoundsException ignored) {
                throw new IllegalArgumentException("not encoded in Base64");
            }
            dest.setByte(destOffset, decodedValue);
            return 1;
        }

        final byte src3 = src[3];
        if (src3 == EQUALS_SIGN) {
            // Example: DkL=
            final byte b1 = decodabet[src1];
            // Packing bytes into a short to reduce bound and reference count checking.
            try {
                // The decodabet bytes are meant to straddle byte boundaries and so we must carefully mask out
                // the bits we care about.
                decodedValue = ((decodabet[src0] & 0x3f) << 2 | (b1 & 0xf0) >> 4) << 8 |
                               (b1 & 0xf) << 4 | (decodabet[src2] & 0xfc) >>> 2;
            } catch (IndexOutOfBoundsException ignored) {
                throw new IllegalArgumentException("not encoded in Base64");
            }
            dest.setShort(destOffset, decodedValue);
            return 2;
        }

        // Example: DkLE
        try {
            decodedValue = (decodabet[src0] & 0x3f) << 18 |
                           (decodabet[src1] & 0xff) << 12 |
                           (decodabet[src2] & 0xff) << 6 |
                           decodabet[src3] & 0xff;
        } catch (IndexOutOfBoundsException ignored) {
            throw new IllegalArgumentException("not encoded in Base64");
        }
        dest.setMedium(destOffset, decodedValue);
        return 3;
    }
}