package com.hubspot.smtp; import static com.hubspot.smtp.ExtensibleNettyServer.NETTY_CHANNEL; import java.io.IOException; import java.io.OutputStream; import java.util.Collection; import java.util.Collections; import java.util.List; import java.util.regex.Matcher; import java.util.regex.Pattern; import org.apache.james.protocols.api.ProtocolSession; import org.apache.james.protocols.api.ProtocolSession.State; import org.apache.james.protocols.api.Request; import org.apache.james.protocols.api.Response; import org.apache.james.protocols.api.future.FutureResponseImpl; import org.apache.james.protocols.api.handler.CommandHandler; import org.apache.james.protocols.api.handler.ExtensibleHandler; import org.apache.james.protocols.api.handler.LineHandler; import org.apache.james.protocols.api.handler.WiringException; import org.apache.james.protocols.netty.HandlerConstants; import org.apache.james.protocols.smtp.MailAddress; import org.apache.james.protocols.smtp.MailEnvelopeImpl; import org.apache.james.protocols.smtp.SMTPResponse; import org.apache.james.protocols.smtp.SMTPRetCode; import org.apache.james.protocols.smtp.SMTPSession; import org.apache.james.protocols.smtp.core.esmtp.EhloExtension; import org.apache.james.protocols.smtp.dsn.DSNStatus; import org.apache.james.protocols.smtp.hook.Hook; import org.apache.james.protocols.smtp.hook.MessageHook; import org.jboss.netty.buffer.ChannelBuffer; import org.jboss.netty.channel.Channel; import org.jboss.netty.channel.ChannelHandlerContext; import org.jboss.netty.channel.ChannelPipeline; import org.jboss.netty.channel.MessageEvent; import org.jboss.netty.channel.SimpleChannelUpstreamHandler; import com.google.common.collect.Lists; public class ChunkingExtension implements EhloExtension, CommandHandler<SMTPSession>, ExtensibleHandler, Hook { private static final String MAIL_ENVELOPE = "mail envelope"; private static final String BDAT_HANDLER_NAME = "BDAT handler"; private static final Pattern BDAT_COMMAND_PATTERN = Pattern.compile("(?<size>[0-9]+)(?<last> LAST)?"); private static final Response DELIVERY_SYNTAX = new SMTPResponse(SMTPRetCode.SYNTAX_ERROR_ARGUMENTS, DSNStatus.getStatus(DSNStatus.PERMANENT, DSNStatus.DELIVERY_SYNTAX) + " Invalid syntax").immutable(); private List<MessageHook> messageHandlers; @Override public Response onCommand(SMTPSession session, Request request) { Matcher matcher = BDAT_COMMAND_PATTERN.matcher(request.getArgument()); if (!matcher.matches()) { return DELIVERY_SYNTAX; } Channel channel = (Channel) session.getAttachment(NETTY_CHANNEL, State.Connection); if (channel == null) { throw new RuntimeException("ExtensibleNettyServer must be used to support chunking"); } BdatHandler bdatHandler = new BdatHandler(channel.getPipeline(), session); bdatHandler.startCapturingData(Integer.parseInt(matcher.group("size")), !matcher.group("last").isEmpty()); return bdatHandler.bdatResponseFuture; } @Override public Collection<String> getImplCommands() { return Collections.singletonList("BDAT"); } @Override public List<String> getImplementedEsmtpFeatures(SMTPSession session) { return Lists.newArrayList("CHUNKING"); } @Override public List<Class<?>> getMarkerInterfaces() { return Lists.newArrayList(MessageHook.class, LineHandler.class); } @Override public void wireExtensions(Class<?> interfaceName, List<?> extension) throws WiringException { // Save MessageHooks so we can tell them about mails we receive if (MessageHook.class.equals(interfaceName)) { messageHandlers = Lists.newArrayList(); for (Object ext : extension) { if (ext instanceof MessageHook) { messageHandlers.add((MessageHook) ext); } } } } private class BdatHandler extends SimpleChannelUpstreamHandler { private final ChannelPipeline pipeline; private final SMTPSession session; private final FutureResponseImpl bdatResponseFuture; private int currentChunkSize; private int bytesRead; private boolean isLastChunk; public BdatHandler(ChannelPipeline pipeline, SMTPSession session) { this.pipeline = pipeline; this.session = session; this.bdatResponseFuture = new FutureResponseImpl(); } @Override public void messageReceived(ChannelHandlerContext ctx, MessageEvent e) throws Exception { if (e.getMessage() instanceof ChannelBuffer) { ChannelBuffer buffer = (ChannelBuffer) e.getMessage(); int bytesToRead = Math.min(currentChunkSize - bytesRead, buffer.readableBytes()); buffer.readBytes(getMailEnvelope().getMessageOutputStream(), bytesToRead); bytesRead += bytesToRead; if (bytesRead == currentChunkSize) { stopCapturingData(); } return; } super.messageReceived(ctx, e); } @SuppressWarnings("unchecked") private void startCapturingData(int bytesToCapture, boolean last) { currentChunkSize = bytesToCapture; isLastChunk = last; pipeline.addBefore(HandlerConstants.FRAMER, BDAT_HANDLER_NAME, this); MailEnvelopeImpl env = new MailEnvelopeImpl(); env.setRecipients(Lists.newArrayList((Collection<MailAddress>) session.getAttachment(SMTPSession.RCPT_LIST, State.Transaction))); env.setSender((MailAddress) session.getAttachment(SMTPSession.SENDER,ProtocolSession.State.Transaction)); session.setAttachment(MAIL_ENVELOPE, env,ProtocolSession.State.Transaction); } private void stopCapturingData() { pipeline.remove(this); bdatResponseFuture.setResponse(new SMTPResponse("250", String.format("Message OK, %d octets received", currentChunkSize))); if (isLastChunk) { callMessageHooks(); } } private void callMessageHooks() { MailEnvelopeImpl env = getMailEnvelope(); try { OutputStream messageOutputStream = env.getMessageOutputStream(); messageOutputStream.flush(); messageOutputStream.close(); } catch (IOException e) { throw new RuntimeException(e); } for (MessageHook hook : messageHandlers) { hook.onMessage(session, env); } } private MailEnvelopeImpl getMailEnvelope() { return (MailEnvelopeImpl) session.getAttachment(MAIL_ENVELOPE, State.Transaction); } } }