org.neo4j.bolt.v1.transport.BoltV1Dechunker Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of neo4j-bolt Show documentation
The core of Neo4j Bolt Protocol, this contains the state machine for Bolt sessions.
There is a newer version: 5.26.1
/*
 * Copyright (c) 2002-2016 "Neo Technology,"
 * Network Engine for Objects in Lund AB [http://neotechnology.com]
 *
 * This file is part of Neo4j.
 *
 * Neo4j is free software: you can redistribute it and/or modify
 * it under the terms of the GNU 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 General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see .
 */
package org.neo4j.bolt.v1.transport;

import io.netty.buffer.ByteBuf;

import java.io.IOException;

import org.neo4j.bolt.v1.messaging.MessageHandler;
import org.neo4j.bolt.v1.messaging.Neo4jPack;
import org.neo4j.bolt.v1.messaging.PackStreamMessageFormatV1;

public class BoltV1Dechunker
{
    private final ChunkedInput input;
    private final PackStreamMessageFormatV1.Reader unpacker;
    private final MessageHandler onMessage;
    private final Runnable onMessageStarted;

    public enum State
    {
        AWAITING_CHUNK,
        IN_CHUNK,
        IN_HEADER,
        CLOSED
    }

    private State state = State.AWAITING_CHUNK;
    private int chunkSize = 0;

    public BoltV1Dechunker( MessageHandler messageHandler, Runnable onMessageStarted )
    {
        this.onMessage = messageHandler;
        this.onMessageStarted = onMessageStarted;
        this.input = new ChunkedInput();
        this.unpacker = new PackStreamMessageFormatV1.Reader( new Neo4jPack.Unpacker( input ) );
    }

    /** Check if we are currently "in the middle of" a message, eg. we've gotten parts of it, but are waiting for more. */
    public boolean isInMiddleOfAMessage()
    {
        return chunkSize != 0;
    }

    public void handle( ByteBuf data ) throws IOException
    {
        while ( data.readableBytes() > 0 )
        {
            switch ( state )
            {
            case AWAITING_CHUNK:
            {
                if ( data.readableBytes() >= 2 )
                {
                    // Whole header available, read that
                    chunkSize = data.readUnsignedShort();
                    handleHeader();
                }
                else
                {
                    // Only one byte available, read that and wait for the second byte
                    chunkSize = data.readUnsignedByte() << 8;
                    state = State.IN_HEADER;
                }
                break;
            }
            case IN_HEADER:
            {
                // First header byte read, now we read the next one
                chunkSize = chunkSize | data.readUnsignedByte();
                handleHeader();
                break;
            }
            case IN_CHUNK:
            {
                if ( chunkSize < data.readableBytes() )
                {
                    // Current packet is larger than current chunk, slice of the chunk
                    input.append( data.readSlice( chunkSize ) );
                    state = State.AWAITING_CHUNK;
                }
                else if ( chunkSize == data.readableBytes() )
                {
                    // Current packet perfectly maps to current chunk
                    input.append( data );
                    state = State.AWAITING_CHUNK;
                    return;
                }
                else
                {
                    // Current packet is smaller than the chunk we're reading, split the current chunk itself up
                    chunkSize -= data.readableBytes();
                    input.append( data );
                    return;
                }
                break;
            }
            case CLOSED:
            {
                // No-op
                return;
            }
            }
        }
    }

    public synchronized void close()
    {
        state = State.CLOSED;
        input.close();
    }

    private void handleHeader() throws IOException
    {
        if(chunkSize == 0)
        {
            // Message boundary
            try
            {
                onMessageStarted.run();
                unpacker.read( onMessage );
            }
            finally
            {
                input.clear();
            }
            state = State.AWAITING_CHUNK;
        }
        else
        {
            state = State.IN_CHUNK;
        }
    }

}