astralis/src/Components/BrotliCompressor.vala

using Invercargill;
using Invercargill.DataStructures;

namespace Astralis {

    public class BrotliCompressor : Compressor {

        private int _quality = 9;

        public BrotliCompressor(int quality = 9, uint64 max_buffer_size = 1024 * 1024 * 16) {
            base(max_buffer_size);
            this._quality = quality.clamp(Brotli.MIN_QUALITY, Brotli.MAX_QUALITY);
        }

        public override string encoding_token { get { return "br"; } }

        /// Picks an EncoderMode based on MIME type for optimal compression
        private static Brotli.EncoderMode pick_encoder_mode(string? content_type) {
            if (content_type == null) {
                return Brotli.EncoderMode.GENERIC;
            }

            string mime = content_type.down();

            // Text mode for text-based content types
            if (mime.has_prefix("text/")) {
                return Brotli.EncoderMode.TEXT;
            }

            // Text mode for application text formats
            if (mime.has_prefix("application/json") ||
                mime.has_prefix("application/xml") ||
                mime.has_prefix("application/javascript") ||
                mime.has_prefix("application/x-javascript") ||
                mime.has_prefix("application/xhtml+xml")) {
                return Brotli.EncoderMode.TEXT;
            }

            // Font mode for web fonts
            if (mime.has_prefix("font/") ||
                mime.has_prefix("application/font-") ||
                mime.has_prefix("application/x-font-") ||
                mime.has_suffix("woff") ||
                mime.has_suffix("woff2") ||
                mime.has_suffix("ttf") ||
                mime.has_suffix("otf") ||
                mime.has_suffix("eot")) {
                return Brotli.EncoderMode.FONT;
            }

            // Default to generic for everything else
            return Brotli.EncoderMode.GENERIC;
        }

        public override ByteBuffer compress_buffer(ByteBuffer data, string? content_type) throws Error {
            // Get maximum compressed size
            size_t max_output_size = Brotli.encoder_max_compressed_size((size_t) data.length);
            if (max_output_size == 0) {
                throw new IOError.FAILED("Brotli Error");
            }

            // Prepare input
            uint8[] input_bytes = data.to_array();

            // Allocate output buffer
            var output_buffer = new uint8[max_output_size];
            size_t encoded_size = max_output_size;

            // Pick encoder mode based on content type
            var mode = pick_encoder_mode(content_type);

            // Compress in one shot (using pointers for the C API)
            Brotli.Bool result = Brotli.encoder_compress(
                _quality,
                Brotli.DEFAULT_WINDOW,
                mode,
                input_bytes.length,
                input_bytes,
                ref encoded_size,
                output_buffer
            );
            
            if (result != Brotli.TRUE) {
                throw new IOError.FAILED("Brotli Error");
            }
            
            // Return compressed data (base class handles size comparison)
            return new ByteBuffer.from_byte_array(output_buffer[0:encoded_size]);
        }

        public override HttpResult compress_chunked(HttpResult inner_result) {
            // Get content type for compression optimization
            string? content_type = null;
            if (inner_result.headers.has("Content-Type")) {
                content_type = inner_result.headers["Content-Type"];
            }
            var mode = pick_encoder_mode(content_type);

            var streaming_result = new StreamingBrotliResult(inner_result, _quality, mode);
            return streaming_result;
        }

        /// Streaming compression result that compresses data on-the-fly using brotli
        private class StreamingBrotliResult : HttpResult {
            private HttpResult inner_result;
            private int quality;
            private Brotli.EncoderMode mode;

            public StreamingBrotliResult(HttpResult result, int quality, Brotli.EncoderMode mode) {
                base(result.status, null); // No content length for streaming
                copy_headers(result, this, {"content-length", "content-encoding"});
                this.inner_result = result;
                this.quality = quality;
                this.mode = mode;
            }

            public async override void send_body(AsyncOutput output) throws Error {
                // Create a brotli output stream wrapper that compresses on-the-fly
                var brotli_output = new BrotliAsyncOutput(output, quality, mode);

                // Send the body through the brotli wrapper
                yield inner_result.send_body(brotli_output);

                // Finish the compression stream
                yield brotli_output.finish();
            }
        }

        /// An AsyncOutput that compresses data on-the-fly using brotli format
        private class BrotliAsyncOutput : Object, AsyncOutput {
            private AsyncOutput downstream;
            private Brotli.EncoderState encoder;
            private bool finished = false;
            private uint8[] output_buffer;

            public bool connected { get { return downstream.connected; } }
            public bool write_would_block { get { return downstream.write_would_block; } }

            public BrotliAsyncOutput(AsyncOutput downstream, int quality, Brotli.EncoderMode mode) throws Error {
                this.downstream = downstream;

                // Initialize brotli encoder
                this.encoder = new Brotli.EncoderState(null, null, null);

                // Set encoder parameters
                this.encoder.set_parameter(Brotli.EncoderParameter.MODE, (uint32) mode);
                this.encoder.set_parameter(Brotli.EncoderParameter.QUALITY, (uint32) quality);
                this.encoder.set_parameter(Brotli.EncoderParameter.LGWIN, (uint32) Brotli.DEFAULT_WINDOW);

                // Allocate output buffer (16KB should be plenty for compressed chunks)
                this.output_buffer = new uint8[16384];
            }

            public async void write(BinaryData data) throws Error {
                if (!connected) {
                    throw new IOError.CLOSED("Cannot write to closed output stream");
                }
                if (finished) {
                    throw new IOError.FAILED("Cannot write to finished brotli stream");
                }
                
                uint8[] input_bytes = data.to_array();
                if (input_bytes.length == 0) {
                    return;
                }
                
                // Set up input
                size_t available_in = input_bytes.length;
                uint8* next_in = input_bytes;
                
                // Compress with PROCESS operation (don't flush yet)
                while (available_in > 0) {
                    size_t available_out = output_buffer.length;
                    uint8* next_out = output_buffer;
                    size_t total_out = 0;
                    
                    Brotli.Bool result = encoder.compress_stream(
                        Brotli.EncoderOperation.PROCESS,
                        ref available_in,
                        ref next_in,
                        ref available_out,
                        ref next_out,
                        out total_out
                    );
                    
                    if (result != Brotli.TRUE) {
                        throw new IOError.FAILED("Brotli compression stream error");
                    }
                    
                    // Calculate how much compressed data was produced
                    size_t compressed_size = output_buffer.length - available_out;
                    if (compressed_size > 0) {
                        yield downstream.write(new ByteBuffer.from_byte_array(output_buffer[0:compressed_size]));
                    }
                }
                
                // Flush any pending output
                yield flush_pending();
            }

            public async void write_stream(InputStream stream) throws Error {
                uint8[] chunk = new uint8[8192];
                while (true) {
                    ssize_t bytes_read = yield stream.read_async(chunk);
                    if (bytes_read <= 0) {
                        break;
                    }
                    yield write(new ByteBuffer.from_byte_array(chunk[0:bytes_read]));
                }
            }
            
            /// Flush any pending output from the encoder
            private async void flush_pending() throws Error {
                // Check if there's more output available
                while (encoder.has_more_output() == Brotli.TRUE) {
                    size_t size = output_buffer.length;
                    unowned uint8* taken = encoder.take_output(ref size);
                    
                    if (size > 0) {
                        // Copy data from taken pointer to array
                        uint8[] data = new uint8[size];
                        Memory.copy(data, taken, size);
                        yield downstream.write(new ByteBuffer.from_byte_array(data));
                    }
                }
            }
            
            /// Finish the compression stream
            public async void finish() throws Error {
                if (finished) {
                    return;
                }
                finished = true;
                
                // Finish the stream
                while (encoder.is_finished() != Brotli.TRUE) {
                    size_t available_in = 0;
                    uint8* next_in = null;
                    size_t available_out = output_buffer.length;
                    uint8* next_out = output_buffer;
                    size_t total_out = 0;
                    
                    Brotli.Bool result = encoder.compress_stream(
                        Brotli.EncoderOperation.FINISH,
                        ref available_in,
                        ref next_in,
                        ref available_out,
                        ref next_out,
                        out total_out
                    );
                    
                    if (result != Brotli.TRUE) {
                        throw new IOError.FAILED("Brotli compression stream error during finish");
                    }
                    
                    // Write any compressed data produced
                    size_t compressed_size = output_buffer.length - available_out;
                    if (compressed_size > 0) {
                        yield downstream.write(new ByteBuffer.from_byte_array(output_buffer[0:compressed_size]));
                    }
                    
                    // Also take any remaining output
                    yield flush_pending();
                }
                
                // Encoder will be cleaned up by the free_function when the object is finalized
            }
        }

    }

}