axiom-api/src/main/java/org/apache/axiom/mime/MultipartBody.java - ws-axiom - Git at Google

 /*
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements. See the NOTICE file
  * distributed with this work for additional information
  * regarding copyright ownership. The ASF licenses this file
  * to you 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 org.apache.axiom.mime;

 import java.io.IOException;
 import java.io.InputStream;
 import java.text.ParseException;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;

 import org.apache.axiom.blob.MemoryBlob;
 import org.apache.axiom.blob.WritableBlobFactory;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.apache.james.mime4j.MimeException;
 import org.apache.james.mime4j.stream.EntityState;
 import org.apache.james.mime4j.stream.Field;
 import org.apache.james.mime4j.stream.MimeConfig;
 import org.apache.james.mime4j.stream.MimeTokenStream;
 import org.apache.james.mime4j.stream.RecursionMode;

 /**
  * A MIME multipart message read from a stream. This class exposes an API that represents the
  * message as a sequence of {@link Part} instances. It implements the {@link Iterable} interface to
  * enable access to all parts. In addition, it supports lookup by content ID. Data is read from the
  * stream on demand. This means that the stream must be kept open until all parts have been
  * processed or {@link #detach()} has been called. It also means that any invocation of a method on
  * an instance of this class or an individual {@link Part} instance may trigger a
  * {@link MIMEException} if there is an I/O error on the stream or a MIME parsing error.
  * <p>
  * Instances of this class are created using a fluent builder; see {@link #builder()}.
  */
 public final class MultipartBody implements Iterable<Part> {
     public interface PartCreationListener {
         void partCreated(Part part);
     }

     public final static class Builder {
         private InputStream inputStream;
         private ContentType contentType;
         private WritableBlobFactory<?> attachmentBlobFactory;
         private PartBlobFactory partBlobFactory;
         private PartCreationListener partCreationListener;

         Builder() {}

         public Builder setInputStream(InputStream inputStream) {
             this.inputStream = inputStream;
             return this;
         }

         public Builder setContentType(ContentType contentType) {
             this.contentType = contentType;
             return this;
         }

         public Builder setContentType(String contentType) {
             try {
                 this.contentType = new ContentType(contentType);
             } catch (ParseException ex) {
                 throw new MIMEException(ex);
             }
             return this;
         }

         public Builder setAttachmentBlobFactory(WritableBlobFactory<?> attachmentBlobFactory) {
             this.attachmentBlobFactory = attachmentBlobFactory;
             return this;
         }

         public Builder setPartBlobFactory(PartBlobFactory partBlobFactory) {
             this.partBlobFactory = partBlobFactory;
             return this;
         }

         public Builder setPartCreationListener(PartCreationListener partCreationListener) {
             this.partCreationListener = partCreationListener;
             return this;
         }

         public MultipartBody build() {
             if (inputStream == null) {
                 throw new IllegalArgumentException("inputStream is mandatory");
             }
             if (contentType == null) {
                 throw new IllegalArgumentException("contentType is mandatory");
             }
             return new MultipartBody(
                     inputStream,
                     contentType,
                     attachmentBlobFactory == null ? MemoryBlob.FACTORY : attachmentBlobFactory,
                     partBlobFactory == null ? PartBlobFactory.DEFAULT : partBlobFactory,
                     partCreationListener);
         }
     }

     private static final Log log = LogFactory.getLog(MultipartBody.class);

     private static final MimeConfig config = MimeConfig.custom().setStrictParsing(true).build();

     /** <code>ContentType</code> of the MIME message */
     private final ContentType contentType;
     private final String rootPartContentID;
     private final MimeTokenStream parser;

     /**
      * Stores the already parsed MIME parts by Content IDs.
      */
     private final Map<String,PartImpl> partMap = new HashMap<String,PartImpl>();

     /**
      * The MIME part currently being processed.
      */
     private PartImpl currentPart;

     private PartImpl firstPart;
     private PartImpl rootPart;

     private int partCount;

     private final WritableBlobFactory<?> attachmentBlobFactory;
     private final PartBlobFactory partBlobFactory;
     private final PartCreationListener partCreationListener;

     MultipartBody(InputStream inStream, ContentType contentType,
             WritableBlobFactory<?> attachmentBlobFactory,
             PartBlobFactory partBlobFactory,
             PartCreationListener partCreationListener) {
         this.attachmentBlobFactory = attachmentBlobFactory;
         this.partBlobFactory = partBlobFactory;
         this.partCreationListener = partCreationListener;
         this.contentType = contentType;

         String start = contentType.getParameter("start");
         rootPartContentID = start == null ? null : normalizeContentID(start);

         parser = new MimeTokenStream(config);
         parser.setRecursionMode(RecursionMode.M_NO_RECURSE);
         parser.parseHeadless(inStream, contentType.toString());

         // Move the parser to the beginning of the first part
         while (parser.getState() != EntityState.T_START_BODYPART) {
             try {
                 parser.next();
             } catch (IOException ex) {
                 throw new MIMEException(ex);
             } catch (MimeException ex) {
                 throw new MIMEException(ex);
             }
         }
     }

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

     private static String normalizeContentID(String contentID) {
         contentID = contentID.trim();
         if (contentID.length() >= 2 && contentID.charAt(0) == '<'
                 && contentID.charAt(contentID.length()-1) == '>') {
             contentID = contentID.substring(1, contentID.length()-1);
         }
         // There is some evidence that some broken MIME implementations add
         // a "cid:" prefix to the Content-ID; remove it if necessary.
         if (contentID.length() > 4 && contentID.startsWith("cid:")) {
             contentID = contentID.substring(4);
         }
         return contentID;
     }

     PartBlobFactory getPartBlobFactory() {
         return partBlobFactory;
     }

     public ContentType getContentType() {
         return contentType;
     }

     /**
      * Get the MIME part with the given content ID.
      *
      * @param contentID
      *            the content ID of the part to retrieve
      * @return the MIME part, or {@code null} if the message doesn't have a part with the given
      *         content ID
      */
     public Part getPart(String contentID) {
         do {
             PartImpl part = partMap.get(contentID);
             if (part != null) {
                 return part;
             }
         } while (getNextPart() != null);
         return null;
     }

     /**
      * Get the number of parts in this multipart.
      *
      * @return the number of parts
      */
     public int getPartCount() {
         detach();
         return partCount;
     }

     PartImpl getFirstPart() {
         if (firstPart == null) {
             getNextPart();
         }
         return firstPart;
     }

     public Part getRootPart() {
         do {
             if (rootPart != null) {
                 return rootPart;
             }
         } while (getNextPart() != null);
         throw new MIMEException(
                 "Mandatory root MIME part is missing");
     }

     PartImpl getNextPart() {
         if (currentPart != null) {
             currentPart.fetch();
         }
         if (parser.getState() == EntityState.T_END_MULTIPART) {
             currentPart = null;
         } else {
             String partContentID = null;
             boolean isRootPart;

             try {
                 checkParserState(parser.next(), EntityState.T_START_HEADER);

                 List<Header> headers = new ArrayList<Header>();
                 while (parser.next() == EntityState.T_FIELD) {
                     Field field = parser.getField();
                     String name = field.getName();
                     String value = field.getBody();

                     if (log.isDebugEnabled()){
                         log.debug("addHeader: (" + name + ") value=(" + value +")");
                     }
                     headers.add(new Header(name, value));
                     if (partContentID == null && name.equalsIgnoreCase("Content-ID")) {
                         partContentID = normalizeContentID(value);
                     }
                 }

                 checkParserState(parser.next(), EntityState.T_BODY);

                 if (rootPartContentID == null) {
                     isRootPart = firstPart == null;
                 } else {
                     isRootPart = rootPartContentID.equals(partContentID);
                 }

                 PartImpl part = new PartImpl(this, isRootPart ? MemoryBlob.FACTORY : attachmentBlobFactory, partContentID, headers, parser);
                 if (currentPart == null) {
                     firstPart = part;
                 } else {
                     currentPart.setNextPart(part);
                 }
                 currentPart = part;
             } catch (IOException ex) {
                 throw new MIMEException(ex);
             } catch (MimeException ex) {
                 throw new MIMEException(ex);
             }

             partCount++;
             if (partContentID != null) {
                 if (partMap.containsKey(partContentID)) {
                     throw new MIMEException(
                             "Two MIME parts with the same Content-ID not allowed.");
                 }
                 partMap.put(partContentID, currentPart);
             }
             if (isRootPart) {
                 rootPart = currentPart;
             }
             if (partCreationListener != null) {
                 partCreationListener.partCreated(currentPart);
             }
         }
         return currentPart;
     }

     private static void checkParserState(EntityState state, EntityState expected) throws IllegalStateException {
         if (expected != state) {
             throw new IllegalStateException("Internal error: expected parser to be in state "
                     + expected + ", but got " + state);
         }
     }

     @Override
     public Iterator<Part> iterator() {
         return new PartIterator(this);
     }

     public void detach() {
         while (getNextPart() != null) {
             // Just loop
         }
     }
 }
	/*
	* Licensed to the Apache Software Foundation (ASF) under one
	* or more contributor license agreements. See the NOTICE file
	* distributed with this work for additional information
	* regarding copyright ownership. The ASF licenses this file
	* to you 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 org.apache.axiom.mime;

	import java.io.IOException;
	import java.io.InputStream;
	import java.text.ParseException;
	import java.util.ArrayList;
	import java.util.HashMap;
	import java.util.Iterator;
	import java.util.List;
	import java.util.Map;

	import org.apache.axiom.blob.MemoryBlob;
	import org.apache.axiom.blob.WritableBlobFactory;
	import org.apache.commons.logging.Log;
	import org.apache.commons.logging.LogFactory;
	import org.apache.james.mime4j.MimeException;
	import org.apache.james.mime4j.stream.EntityState;
	import org.apache.james.mime4j.stream.Field;
	import org.apache.james.mime4j.stream.MimeConfig;
	import org.apache.james.mime4j.stream.MimeTokenStream;
	import org.apache.james.mime4j.stream.RecursionMode;

	/**
	* A MIME multipart message read from a stream. This class exposes an API that represents the
	* message as a sequence of {@link Part} instances. It implements the {@link Iterable} interface to
	* enable access to all parts. In addition, it supports lookup by content ID. Data is read from the
	* stream on demand. This means that the stream must be kept open until all parts have been
	* processed or {@link #detach()} has been called. It also means that any invocation of a method on
	* an instance of this class or an individual {@link Part} instance may trigger a
	* {@link MIMEException} if there is an I/O error on the stream or a MIME parsing error.
	* <p>
	* Instances of this class are created using a fluent builder; see {@link #builder()}.
	*/
	public final class MultipartBody implements Iterable<Part> {
	public interface PartCreationListener {
	void partCreated(Part part);
	}

	public final static class Builder {
	private InputStream inputStream;
	private ContentType contentType;
	private WritableBlobFactory<?> attachmentBlobFactory;
	private PartBlobFactory partBlobFactory;
	private PartCreationListener partCreationListener;

	Builder() {}

	public Builder setInputStream(InputStream inputStream) {
	this.inputStream = inputStream;
	return this;
	}

	public Builder setContentType(ContentType contentType) {
	this.contentType = contentType;
	return this;
	}

	public Builder setContentType(String contentType) {
	try {
	this.contentType = new ContentType(contentType);
	} catch (ParseException ex) {
	throw new MIMEException(ex);
	}
	return this;
	}

	public Builder setAttachmentBlobFactory(WritableBlobFactory<?> attachmentBlobFactory) {
	this.attachmentBlobFactory = attachmentBlobFactory;
	return this;
	}

	public Builder setPartBlobFactory(PartBlobFactory partBlobFactory) {
	this.partBlobFactory = partBlobFactory;
	return this;
	}

	public Builder setPartCreationListener(PartCreationListener partCreationListener) {
	this.partCreationListener = partCreationListener;
	return this;
	}

	public MultipartBody build() {
	if (inputStream == null) {
	throw new IllegalArgumentException("inputStream is mandatory");
	}
	if (contentType == null) {
	throw new IllegalArgumentException("contentType is mandatory");
	}
	return new MultipartBody(
	inputStream,
	contentType,
	attachmentBlobFactory == null ? MemoryBlob.FACTORY : attachmentBlobFactory,
	partBlobFactory == null ? PartBlobFactory.DEFAULT : partBlobFactory,
	partCreationListener);
	}
	}

	private static final Log log = LogFactory.getLog(MultipartBody.class);

	private static final MimeConfig config = MimeConfig.custom().setStrictParsing(true).build();

	/** <code>ContentType</code> of the MIME message */
	private final ContentType contentType;
	private final String rootPartContentID;
	private final MimeTokenStream parser;

	/**
	* Stores the already parsed MIME parts by Content IDs.
	*/
	private final Map<String,PartImpl> partMap = new HashMap<String,PartImpl>();

	/**
	* The MIME part currently being processed.
	*/
	private PartImpl currentPart;

	private PartImpl firstPart;
	private PartImpl rootPart;

	private int partCount;

	private final WritableBlobFactory<?> attachmentBlobFactory;
	private final PartBlobFactory partBlobFactory;
	private final PartCreationListener partCreationListener;

	MultipartBody(InputStream inStream, ContentType contentType,
	WritableBlobFactory<?> attachmentBlobFactory,
	PartBlobFactory partBlobFactory,
	PartCreationListener partCreationListener) {
	this.attachmentBlobFactory = attachmentBlobFactory;
	this.partBlobFactory = partBlobFactory;
	this.partCreationListener = partCreationListener;
	this.contentType = contentType;

	String start = contentType.getParameter("start");
	rootPartContentID = start == null ? null : normalizeContentID(start);

	parser = new MimeTokenStream(config);
	parser.setRecursionMode(RecursionMode.M_NO_RECURSE);
	parser.parseHeadless(inStream, contentType.toString());

	// Move the parser to the beginning of the first part
	while (parser.getState() != EntityState.T_START_BODYPART) {
	try {
	parser.next();
	} catch (IOException ex) {
	throw new MIMEException(ex);
	} catch (MimeException ex) {
	throw new MIMEException(ex);
	}
	}
	}

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

	private static String normalizeContentID(String contentID) {
	contentID = contentID.trim();
	if (contentID.length() >= 2 && contentID.charAt(0) == '<'
	&& contentID.charAt(contentID.length()-1) == '>') {
	contentID = contentID.substring(1, contentID.length()-1);
	}
	// There is some evidence that some broken MIME implementations add
	// a "cid:" prefix to the Content-ID; remove it if necessary.
	if (contentID.length() > 4 && contentID.startsWith("cid:")) {
	contentID = contentID.substring(4);
	}
	return contentID;
	}

	PartBlobFactory getPartBlobFactory() {
	return partBlobFactory;
	}

	public ContentType getContentType() {
	return contentType;
	}

	/**
	* Get the MIME part with the given content ID.
	*
	* @param contentID
	* the content ID of the part to retrieve
	* @return the MIME part, or {@code null} if the message doesn't have a part with the given
	* content ID
	*/
	public Part getPart(String contentID) {
	do {
	PartImpl part = partMap.get(contentID);
	if (part != null) {
	return part;
	}
	} while (getNextPart() != null);
	return null;
	}

	/**
	* Get the number of parts in this multipart.
	*
	* @return the number of parts
	*/
	public int getPartCount() {
	detach();
	return partCount;
	}

	PartImpl getFirstPart() {
	if (firstPart == null) {
	getNextPart();
	}
	return firstPart;
	}

	public Part getRootPart() {
	do {
	if (rootPart != null) {
	return rootPart;
	}
	} while (getNextPart() != null);
	throw new MIMEException(
	"Mandatory root MIME part is missing");
	}

	PartImpl getNextPart() {
	if (currentPart != null) {
	currentPart.fetch();
	}
	if (parser.getState() == EntityState.T_END_MULTIPART) {
	currentPart = null;
	} else {
	String partContentID = null;
	boolean isRootPart;

	try {
	checkParserState(parser.next(), EntityState.T_START_HEADER);

	List<Header> headers = new ArrayList<Header>();
	while (parser.next() == EntityState.T_FIELD) {
	Field field = parser.getField();
	String name = field.getName();
	String value = field.getBody();

	if (log.isDebugEnabled()){
	log.debug("addHeader: (" + name + ") value=(" + value +")");
	}
	headers.add(new Header(name, value));
	if (partContentID == null && name.equalsIgnoreCase("Content-ID")) {
	partContentID = normalizeContentID(value);
	}
	}

	checkParserState(parser.next(), EntityState.T_BODY);

	if (rootPartContentID == null) {
	isRootPart = firstPart == null;
	} else {
	isRootPart = rootPartContentID.equals(partContentID);
	}

	PartImpl part = new PartImpl(this, isRootPart ? MemoryBlob.FACTORY : attachmentBlobFactory, partContentID, headers, parser);
	if (currentPart == null) {
	firstPart = part;
	} else {
	currentPart.setNextPart(part);
	}
	currentPart = part;
	} catch (IOException ex) {
	throw new MIMEException(ex);
	} catch (MimeException ex) {
	throw new MIMEException(ex);
	}

	partCount++;
	if (partContentID != null) {
	if (partMap.containsKey(partContentID)) {
	throw new MIMEException(
	"Two MIME parts with the same Content-ID not allowed.");
	}
	partMap.put(partContentID, currentPart);
	}
	if (isRootPart) {
	rootPart = currentPart;
	}
	if (partCreationListener != null) {
	partCreationListener.partCreated(currentPart);
	}
	}
	return currentPart;
	}

	private static void checkParserState(EntityState state, EntityState expected) throws IllegalStateException {
	if (expected != state) {
	throw new IllegalStateException("Internal error: expected parser to be in state "
	+ expected + ", but got " + state);
	}
	}

	@Override
	public Iterator<Part> iterator() {
	return new PartIterator(this);
	}

	public void detach() {
	while (getNextPart() != null) {
	// Just loop
	}
	}
	}