ConsumerSession.java

package com.trendmicro.mist.session;

import java.io.IOException;
import java.util.Enumeration;
import java.util.Iterator;

import javax.jms.BytesMessage;
import javax.jms.Destination;
import javax.jms.JMSException;
import javax.jms.Message;
import javax.jms.MessageListener;
import javax.jms.Queue;
import javax.jms.Topic;

import com.google.protobuf.ByteString;
import com.trendmicro.mist.Client;
import com.trendmicro.mist.Daemon;
import com.trendmicro.mist.ExchangeMetric;
import com.trendmicro.mist.MistException;
import com.trendmicro.mist.proto.GateTalk;
import com.trendmicro.mist.proto.MistMessage;
import com.trendmicro.mist.proto.MistMessage.KeyValuePair;
import com.trendmicro.mist.proto.MistMessage.MessageBlock;
import com.trendmicro.mist.util.Exchange;
import com.trendmicro.mist.util.MessageFilter;
import com.trendmicro.mist.util.Packet;
import com.trendmicro.spn.common.util.Utils;

public class ConsumerSession extends Session implements MessageListener {
    public static class MessagePrepared {
        public MistMessage.MessageBlock msgBlock;
        public Exchange from;
        public MessageBlock.Builder builder;
        
        private byte[] convertJMSMessage(BytesMessage msg) throws JMSException, IOException {
            ByteString.Output byteOut = ByteString.newOutput();
            byte[] buf = new byte[1024];
            int len = -1;
            while((len = msg.readBytes(buf)) > 0)
                byteOut.write(buf, 0, len);
            return byteOut.toByteString().toByteArray();
        }
        
        public MessagePrepared(BytesMessage msg) throws MistException, JMSException, IOException {
            Destination d = msg.getJMSDestination();
            if(d instanceof Queue)
                from = new Exchange("queue:" + ((Queue) d).getQueueName());
            else if(d instanceof Topic)
                from = new Exchange("topic:" + ((Topic) d).getTopicName());
            
            builder = MistMessage.MessageBlock.newBuilder();
            builder.setId(from.toString());
            
            byte[] payload = convertJMSMessage(msg);
            builder.setMessage(ByteString.copyFrom(payload));

            Enumeration<?> propNames = msg.getPropertyNames();
            while(propNames.hasMoreElements()) {
                String key = (String) propNames.nextElement();
                String value = msg.getStringProperty(key);
                if(key.equals("MIST_TTL"))
                    builder.setTtl(Long.valueOf(value));
                else
                    builder.addProperties(KeyValuePair.newBuilder().setKey(key).setValue(value).build());
            }

            Iterator<MessageFilter> iter = Daemon.messageFilters.descendingIterator();
            while(iter.hasNext()){
                iter.next().postReceive(this);
            }
            msgBlock = builder.build();
        }
    }
    
    private Packet pack = new Packet();
    private String prefixThreadName = "";
    private boolean attached = false;
    private boolean unack = false;
    
    public ConsumerSession(int sessId, GateTalk.Session sessConfig) throws MistException {
        super(sessId, sessConfig);
        prefixThreadName = "Session-" + sessId + "-";
    }
    
    @Override
    public void run() {
        if(!acceptConnection())
            return;
        open(false);
        isReady = true;
        attached = true;
        unack = false;
        
        for(Client c : allClients.values()) {
            try {
                c.getConsumer().setMessageListener(this);
            }
            catch(Exception e) {
                logger.error(e.getMessage(), e);
            }
        }
    }
    
    @Override
    public void addClientIfAttached(Client c) {
        try {
            c.getConsumer().setMessageListener(this);
        }
        catch(Exception e) {
            logger.error(e.getMessage(), e);
        }
    }
    
    @Override
    protected void detach() {
        attached = false;
        detachNow = true;
        for(int i = 0; i < 10; i++) {
            if(!unack)
                break;
            Utils.justSleep(500);
        }
        unack = false;
        try {
            socketInput.close();
        }
        catch(Exception e) {
        }
    }
    
    private void onMessageProcess(Message msg) throws IOException {
        if(detachNow)
            return;
        MessagePrepared mp = null;
        try {
            mp = new MessagePrepared((BytesMessage) msg);
        }
        catch(Exception e) {
            logger.error(e.getMessage(), e);
            try {
                msg.acknowledge();
            }
            catch(JMSException e1) {
            }
            return;
        }
        
        pack.setPayload(mp.msgBlock.toByteArray());
        unack = true;
        pack.write(socketOutput);
        
        ExchangeMetric metric = ExchangeMetric.getExchangeMetric(mp.from);
        metric.increaseMessageIn(mp.msgBlock.getMessage().size());
        
        if(pack.read(socketInput) <= 0)
            return;
        
        try {
            if(GateTalk.Response.newBuilder().mergeFrom(pack.getPayload()).build().getSuccess())
                msg.acknowledge();
            unack = false;
        }
        catch(Exception e) {
            logger.error(e.getMessage(), e);
        }
    }
    
    @Override
    public synchronized void onMessage(Message msg) {
        if(detachNow){
            return;
        }
        String imqThreadName = Thread.currentThread().getName();
        Thread.currentThread().setName(prefixThreadName + imqThreadName);
        try {
            onMessageProcess(msg);
        }
        catch(Exception e) {
            logger.error("Encounter error during delivering message, detaching the session", e);
            try {
                detach(GateTalk.Request.Role.SOURCE);
            }
            catch(MistException e1) {
                logger.error(e1.getMessage(), e1);
            }
        }
        Thread.currentThread().setName(imqThreadName);
    }
    
    @Override
    public boolean isAttached() {
        return attached;
    }
}