<commit_before>package de.bwaldvogel.mongo.wire;

import io.netty.buffer.ByteBuf;
import io.netty.channel.Channel;
import io.netty.channel.ChannelHandlerContext;
import io.netty.handler.codec.LengthFieldBasedFrameDecoder;

import java.io.IOException;
import java.nio.ByteOrder;
import java.util.ArrayList;
import java.util.List;

import org.bson.BSONObject;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import de.bwaldvogel.mongo.wire.message.ClientRequest;
import de.bwaldvogel.mongo.wire.message.MessageHeader;
import de.bwaldvogel.mongo.wire.message.MongoDelete;
import de.bwaldvogel.mongo.wire.message.MongoInsert;
import de.bwaldvogel.mongo.wire.message.MongoQuery;
import de.bwaldvogel.mongo.wire.message.MongoUpdate;

/**
 * Based on information from <a
 * href="http://www.mongodb.org/display/DOCS/Mongo+Wire+Protocol"
 * >http://www.mongodb.org/display/DOCS/Mongo+Wire+Protocol</a>
 */
public class MongoWireProtocolHandler extends LengthFieldBasedFrameDecoder {

    public static final int MAX_MESSAGE_SIZE_BYTES = 48 * 1000 * 1000;

    private static final Logger log = LoggerFactory.getLogger(MongoWireProtocolHandler.class);

    private static final int maxFrameLength = Integer.MAX_VALUE;
    private static final int lengthFieldOffset = 0;
    private static final int lengthFieldLength = 4;
    private static final int lengthAdjustment = -lengthFieldLength;
    private static final int initialBytesToStrip = 0;

    private final BsonDecoder bsonDecoder;

    public MongoWireProtocolHandler() {
        super(maxFrameLength, lengthFieldOffset, lengthFieldLength, lengthAdjustment, initialBytesToStrip);
        bsonDecoder = new BsonDecoder();
    }

    @Override
    protected ClientRequest decode(ChannelHandlerContext ctx, ByteBuf buf) throws Exception {

        ByteBuf in = buf.order(ByteOrder.LITTLE_ENDIAN);

        if (in.readableBytes() < 4) {
            return null;
        }

        in.markReaderIndex();
        int totalLength = in.readInt();

        if (totalLength > MAX_MESSAGE_SIZE_BYTES) {
            throw new IOException("message too large: " + totalLength + " bytes");
        }

        if (in.readableBytes() < totalLength - lengthFieldLength) {
            in.resetReaderIndex();
            return null; // retry
        }
        in = in.readSlice(totalLength - lengthFieldLength);
        int readable = in.readableBytes();
        if (readable != totalLength - lengthFieldLength) {
            throw new IllegalStateException();
        }

        final int requestID = in.readInt();
        final int responseTo = in.readInt();
        final MessageHeader header = new MessageHeader(requestID, responseTo);

        int opCodeId = in.readInt();
        final OpCode opCode = OpCode.getById(opCodeId);
        if (opCode == null) {
            throw new IOException("opCode " + opCodeId + " not supported");
        }

        final Channel channel = ctx.channel();
        final ClientRequest ret;

        switch (opCode) {
        case OP_QUERY:
            ret = handleQuery(channel, header, in);
            break;
        case OP_INSERT:
            ret = handleInsert(channel, header, in);
            break;
        case OP_DELETE:
            ret = handleDelete(channel, header, in);
            break;
        case OP_UPDATE:
            ret = handleUpdate(channel, header, in);
            break;
        default:
            throw new UnsupportedOperationException("unsupported opcode: " + opCode);
        }

        if (in.isReadable()) {
            throw new IOException();
        }

        return ret;
    }

    private ClientRequest handleDelete(Channel channel, MessageHeader header, ByteBuf buffer) throws IOException {

        buffer.skipBytes(4); // reserved

        final String fullCollectionName = bsonDecoder.decodeCString(buffer);

        final int flags = buffer.readInt();
        boolean singleRemove = false;
        if (flags == 0) {
            // ignore
        } else if (flags == 1) {
            singleRemove = true;
        } else {
            throw new UnsupportedOperationException("flags=" + flags + " not yet supported");
        }

        BSONObject selector = bsonDecoder.decodeBson(buffer);
        log.debug("delete {} from {}", selector, fullCollectionName);
        return new MongoDelete(channel, header, fullCollectionName, selector, singleRemove);
    }

    private ClientRequest handleUpdate(Channel channel, MessageHeader header, ByteBuf buffer) throws IOException {

        buffer.skipBytes(4); // reserved

        final String fullCollectionName = bsonDecoder.decodeCString(buffer);

        final int flags = buffer.readInt();
        boolean upsert = UpdateFlag.UPSERT.isSet(flags);
        boolean multi = UpdateFlag.MULTI_UPDATE.isSet(flags);

        BSONObject selector = bsonDecoder.decodeBson(buffer);
        BSONObject update = bsonDecoder.decodeBson(buffer);
        log.debug("update {} in {}", selector, fullCollectionName);
        return new MongoUpdate(channel, header, fullCollectionName, selector, update, upsert, multi);
    }

    private ClientRequest handleInsert(Channel channel, MessageHeader header, ByteBuf buffer) throws IOException {

        final int flags = buffer.readInt();
        if (flags != 0)
            throw new UnsupportedOperationException("flags=" + flags + " not yet supported");

        final String fullCollectionName = bsonDecoder.decodeCString(buffer);

        List<BSONObject> documents = new ArrayList<BSONObject>();
        while (buffer.isReadable()) {
            BSONObject document = bsonDecoder.decodeBson(buffer);
            if (document == null) {
                return null;
            }
            documents.add(document);
        }
        log.debug("insert {} in {}", documents, fullCollectionName);
        return new MongoInsert(channel, header, fullCollectionName, documents);
    }

    private ClientRequest handleQuery(Channel channel, MessageHeader header, ByteBuf buffer) throws IOException {

        int flags = buffer.readInt();

        final String fullCollectionName = bsonDecoder.decodeCString(buffer);
        final int numberToSkip = buffer.readInt();
        final int numberToReturn = buffer.readInt();

        BSONObject query = bsonDecoder.decodeBson(buffer);
        BSONObject returnFieldSelector = null;
        if (buffer.isReadable()) {
            returnFieldSelector = bsonDecoder.decodeBson(buffer);
        }

        MongoQuery mongoQuery = new MongoQuery(channel, header, fullCollectionName, numberToSkip, numberToReturn,
                query, returnFieldSelector);

        if (QueryFlag.SLAVE_OK.isSet(flags)) {
            flags = QueryFlag.SLAVE_OK.removeFrom(flags);
        }

        if (QueryFlag.NO_CURSOR_TIMEOUT.isSet(flags)) {
            flags = QueryFlag.NO_CURSOR_TIMEOUT.removeFrom(flags);
        }

        if (flags != 0)
            throw new UnsupportedOperationException("flags=" + flags + " not yet supported");

        log.debug("query {} from {}", query, fullCollectionName);

        return mongoQuery;
    }

}<commit_msg>code formatting<commit_after>package de.bwaldvogel.mongo.wire;

import io.netty.buffer.ByteBuf;
import io.netty.channel.Channel;
import io.netty.channel.ChannelHandlerContext;
import io.netty.handler.codec.LengthFieldBasedFrameDecoder;

import java.io.IOException;
import java.nio.ByteOrder;
import java.util.ArrayList;
import java.util.List;

import org.bson.BSONObject;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import de.bwaldvogel.mongo.wire.message.ClientRequest;
import de.bwaldvogel.mongo.wire.message.MessageHeader;
import de.bwaldvogel.mongo.wire.message.MongoDelete;
import de.bwaldvogel.mongo.wire.message.MongoInsert;
import de.bwaldvogel.mongo.wire.message.MongoQuery;
import de.bwaldvogel.mongo.wire.message.MongoUpdate;

/**
 * Based on information from <a
 * href="http://www.mongodb.org/display/DOCS/Mongo+Wire+Protocol"
 * >http://www.mongodb.org/display/DOCS/Mongo+Wire+Protocol</a>
 */
public class MongoWireProtocolHandler extends LengthFieldBasedFrameDecoder {

    public static final int MAX_MESSAGE_SIZE_BYTES = 48 * 1000 * 1000;

    private static final Logger log = LoggerFactory.getLogger(MongoWireProtocolHandler.class);

    private static final int maxFrameLength = Integer.MAX_VALUE;
    private static final int lengthFieldOffset = 0;
    private static final int lengthFieldLength = 4;
    private static final int lengthAdjustment = -lengthFieldLength;
    private static final int initialBytesToStrip = 0;

    private final BsonDecoder bsonDecoder;

    public MongoWireProtocolHandler() {
        super(maxFrameLength, lengthFieldOffset, lengthFieldLength, lengthAdjustment, initialBytesToStrip);
        bsonDecoder = new BsonDecoder();
    }

    @Override
    protected ClientRequest decode(ChannelHandlerContext ctx, ByteBuf buf) throws Exception {

        ByteBuf in = buf.order(ByteOrder.LITTLE_ENDIAN);

        if (in.readableBytes() < 4) {
            return null;
        }

        in.markReaderIndex();
        int totalLength = in.readInt();

        if (totalLength > MAX_MESSAGE_SIZE_BYTES) {
            throw new IOException("message too large: " + totalLength + " bytes");
        }

        if (in.readableBytes() < totalLength - lengthFieldLength) {
            in.resetReaderIndex();
            return null; // retry
        }
        in = in.readSlice(totalLength - lengthFieldLength);
        int readable = in.readableBytes();
        if (readable != totalLength - lengthFieldLength) {
            throw new IllegalStateException();
        }

        final int requestID = in.readInt();
        final int responseTo = in.readInt();
        final MessageHeader header = new MessageHeader(requestID, responseTo);

        int opCodeId = in.readInt();
        final OpCode opCode = OpCode.getById(opCodeId);
        if (opCode == null) {
            throw new IOException("opCode " + opCodeId + " not supported");
        }

        final Channel channel = ctx.channel();
        final ClientRequest ret;

        switch (opCode) {
        case OP_QUERY:
            ret = handleQuery(channel, header, in);
            break;
        case OP_INSERT:
            ret = handleInsert(channel, header, in);
            break;
        case OP_DELETE:
            ret = handleDelete(channel, header, in);
            break;
        case OP_UPDATE:
            ret = handleUpdate(channel, header, in);
            break;
        default:
            throw new UnsupportedOperationException("unsupported opcode: " + opCode);
        }

        if (in.isReadable()) {
            throw new IOException();
        }

        return ret;
    }

    private ClientRequest handleDelete(Channel channel, MessageHeader header, ByteBuf buffer) throws IOException {

        buffer.skipBytes(4); // reserved

        final String fullCollectionName = bsonDecoder.decodeCString(buffer);

        final int flags = buffer.readInt();
        boolean singleRemove = false;
        if (flags == 0) {
            // ignore
        } else if (flags == 1) {
            singleRemove = true;
        } else {
            throw new UnsupportedOperationException("flags=" + flags + " not yet supported");
        }

        BSONObject selector = bsonDecoder.decodeBson(buffer);
        log.debug("delete {} from {}", selector, fullCollectionName);
        return new MongoDelete(channel, header, fullCollectionName, selector, singleRemove);
    }

    private ClientRequest handleUpdate(Channel channel, MessageHeader header, ByteBuf buffer) throws IOException {

        buffer.skipBytes(4); // reserved

        final String fullCollectionName = bsonDecoder.decodeCString(buffer);

        final int flags = buffer.readInt();
        boolean upsert = UpdateFlag.UPSERT.isSet(flags);
        boolean multi = UpdateFlag.MULTI_UPDATE.isSet(flags);

        BSONObject selector = bsonDecoder.decodeBson(buffer);
        BSONObject update = bsonDecoder.decodeBson(buffer);
        log.debug("update {} in {}", selector, fullCollectionName);
        return new MongoUpdate(channel, header, fullCollectionName, selector, update, upsert, multi);
    }

    private ClientRequest handleInsert(Channel channel, MessageHeader header, ByteBuf buffer) throws IOException {

        final int flags = buffer.readInt();
        if (flags != 0) {
            throw new UnsupportedOperationException("flags=" + flags + " not yet supported");
        }

        final String fullCollectionName = bsonDecoder.decodeCString(buffer);

        List<BSONObject> documents = new ArrayList<BSONObject>();
        while (buffer.isReadable()) {
            BSONObject document = bsonDecoder.decodeBson(buffer);
            if (document == null) {
                return null;
            }
            documents.add(document);
        }
        log.debug("insert {} in {}", documents, fullCollectionName);
        return new MongoInsert(channel, header, fullCollectionName, documents);
    }

    private ClientRequest handleQuery(Channel channel, MessageHeader header, ByteBuf buffer) throws IOException {

        int flags = buffer.readInt();

        final String fullCollectionName = bsonDecoder.decodeCString(buffer);
        final int numberToSkip = buffer.readInt();
        final int numberToReturn = buffer.readInt();

        BSONObject query = bsonDecoder.decodeBson(buffer);
        BSONObject returnFieldSelector = null;
        if (buffer.isReadable()) {
            returnFieldSelector = bsonDecoder.decodeBson(buffer);
        }

        MongoQuery mongoQuery = new MongoQuery(channel, header, fullCollectionName, numberToSkip, numberToReturn,
                query, returnFieldSelector);

        if (QueryFlag.SLAVE_OK.isSet(flags)) {
            flags = QueryFlag.SLAVE_OK.removeFrom(flags);
        }

        if (QueryFlag.NO_CURSOR_TIMEOUT.isSet(flags)) {
            flags = QueryFlag.NO_CURSOR_TIMEOUT.removeFrom(flags);
        }

        if (flags != 0) {
            throw new UnsupportedOperationException("flags=" + flags + " not yet supported");
        }

        log.debug("query {} from {}", query, fullCollectionName);

        return mongoQuery;
    }

}
{'fraction_non_alphanumeric': 0.09460791214979586, 'fraction_numerical': 0.0030972828382373643, 'mean_word_length': 2.772437599575146, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 4, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6891177', 'n_tokens_mistral': 3895, 'n_tokens_neox': 3732, 'n_words': 921}