com.spotify.netty4.handler.codec.zmtp.ZMTPCodec.java Source code

Introduction

Here is the source code for com.spotify.netty4.handler.codec.zmtp.ZMTPCodec.java

Source

/*
 * Copyright (c) 2012-2015 Spotify AB
 *
 * 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 com.spotify.netty4.handler.codec.zmtp;

import java.nio.ByteBuffer;
import java.nio.channels.ClosedChannelException;
import java.util.List;

import io.netty.buffer.ByteBuf;
import io.netty.channel.Channel;
import io.netty.channel.ChannelHandler;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.CombinedChannelDuplexHandler;
import io.netty.handler.codec.ReplayingDecoder;

import static com.spotify.netty4.handler.codec.zmtp.ZMTPUtils.checkNotNull;

/**
 * A ZMTP codec for Netty.
 *
 * Note: A single codec instance is not {@link Sharable} among multiple {@link Channel} instances.
 */
public class ZMTPCodec extends ReplayingDecoder<Void> {

    private final ZMTPSession session;
    private final ZMTPHandshaker handshaker;

    private final ZMTPConfig config;

    public ZMTPCodec(final ZMTPSession session) {
        this.config = session.config();
        this.session = checkNotNull(session, "session");
        this.handshaker = config.protocol().handshaker(config);
    }

    /**
     * Get the {@link ZMTPSession} for this codec.
     */
    public ZMTPSession session() {
        return session;
    }

    @Override
    public void channelActive(final ChannelHandlerContext ctx) throws Exception {
        super.channelActive(ctx);
        ctx.writeAndFlush(handshaker.greeting());
    }

    @Override
    public void channelInactive(final ChannelHandlerContext ctx) throws Exception {
        super.channelInactive(ctx);
        if (!session.handshakeFuture().isDone()) {
            session.handshakeFailure(new ClosedChannelException());
            ctx.fireUserEventTriggered(new ZMTPHandshakeFailure(session));
        }
    }

    @Override
    protected void decode(final ChannelHandlerContext ctx, final ByteBuf in, final List<Object> out)
            throws Exception {

        // Discard input if handshake failed. It is expected that the user will close the channel.
        if (session.handshakeFuture().isDone()) {
            assert !session.handshakeFuture().isSuccess();
            in.skipBytes(in.readableBytes());
        }

        // Shake hands
        final ZMTPHandshake handshake;
        try {
            handshake = handshaker.handshake(in, ctx);
            if (handshake == null) {
                // Handshake is not yet done. Await more input.
                return;
            }
        } catch (Exception e) {
            session.handshakeFailure(e);
            ctx.fireUserEventTriggered(new ZMTPHandshakeFailure(session));
            throw e;
        }

        // Handshake is done.
        session.handshakeSuccess(handshake);

        // Replace this handler with the framing encoder and decoder
        if (actualReadableBytes() > 0) {
            out.add(in.readBytes(actualReadableBytes()));
        }
        final ZMTPDecoder decoder = config.decoder().decoder(session);
        final ZMTPEncoder encoder = config.encoder().encoder(session);
        final ZMTPWireFormat wireFormat = ZMTPWireFormats.wireFormat(session.negotiatedVersion());
        final ChannelHandler handler = new CombinedChannelDuplexHandler<ZMTPFramingDecoder, ZMTPFramingEncoder>(
                new ZMTPFramingDecoder(wireFormat, decoder), new ZMTPFramingEncoder(wireFormat, encoder));
        ctx.pipeline().replace(this, ctx.name(), handler);

        // Tell the user that the handshake is complete
        ctx.fireUserEventTriggered(new ZMTPHandshakeSuccess(session, handshake));
    }

    public static Builder builder() {
        return new Builder();
    }

    public static ZMTPCodec from(final ZMTPConfig config) {
        return new ZMTPCodec(ZMTPSession.from(config));
    }

    public static ZMTPCodec of(final ZMTPSocketType socketType) {
        return builder().socketType(socketType).build();
    }

    public static ZMTPCodec from(final ZMTPSession session) {
        return new ZMTPCodec(session);
    }

    public static class Builder {

        private final ZMTPConfig.Builder config = ZMTPConfig.builder();

        private Builder() {
        }

        public Builder protocol(final ZMTPProtocol protocol) {
            config.protocol(protocol);
            return this;
        }

        public Builder interop(final boolean interop) {
            config.interop(interop);
            return this;
        }

        public Builder socketType(final ZMTPSocketType socketType) {
            config.socketType(socketType);
            return this;
        }

        public Builder localIdentity(final CharSequence localIdentity) {
            config.localIdentity(localIdentity);
            return this;
        }

        public Builder localIdentity(final byte[] localIdentity) {
            config.localIdentity(localIdentity);
            return this;
        }

        public Builder localIdentity(final ByteBuffer localIdentity) {
            config.localIdentity(localIdentity);
            return this;
        }

        public Builder encoder(final ZMTPEncoder.Factory encoder) {
            config.encoder(encoder);
            return this;
        }

        public Builder encoder(final Class<? extends ZMTPEncoder> encoder) {
            config.encoder(encoder);
            return this;
        }

        public Builder decoder(final ZMTPDecoder.Factory decoder) {
            config.decoder(decoder);
            return this;
        }

        public Builder decoder(final Class<? extends ZMTPDecoder> decoder) {
            config.decoder(decoder);
            return this;
        }

        public Builder identityGenerator(final ZMTPIdentityGenerator identityGenerator) {
            config.identityGenerator(identityGenerator);
            return this;
        }

        public ZMTPCodec build() {
            return ZMTPCodec.from(config.build());
        }
    }
}