Emulate extraction of header+body for a specific part (IMAP).

3 files changed, 86 insertions, 19 deletions

diff --git a/src/vmime/net/imap/IMAPMessage.cpp b/src/vmime/net/imap/IMAPMessage.cpp
index f32434b4..123a94a9 100644
--- a/src/vmime/net/imap/IMAPMessage.cpp
+++ b/src/vmime/net/imap/IMAPMessage.cpp
@@ -61,11 +61,11 @@ class IMAPMessage_literalHandler : public IMAPParser::literalHandler
 public:
 
 	IMAPMessage_literalHandler(utility::outputStream& os, utility::progressListener* progress)
-		: m_os(os), m_progress(progress)
 	{
+		m_target = shared_ptr <target>(new targetStream(progress, os));
 	}
 
-	target* targetFor(const IMAPParser::component& comp, const int /* data */)
+	shared_ptr <target> targetFor(const IMAPParser::component& comp, const int /* data */)
 	{
 		if (typeid(comp) == typeid(IMAPParser::msg_att_item))
 		{
@@ -75,17 +75,21 @@ public:
 			if (type == IMAPParser::msg_att_item::BODY_SECTION ||
 			    type == IMAPParser::msg_att_item::RFC822_TEXT)
 			{
-				return new targetStream(m_progress, m_os);
+				return m_target;
 			}
 		}
 
-		return (NULL);
+		return shared_ptr <target>();
+	}
+
+	shared_ptr <target> getTarget()
+	{
+		return m_target;
 	}
 
 private:
 
-	utility::outputStream& m_os;
-	utility::progressListener* m_progress;
+	shared_ptr <target> m_target;
 };
 
 #endif // VMIME_BUILDING_DOC
@@ -268,7 +272,7 @@ void IMAPMessage::fetchPartHeaderForStructure(shared_ptr <messageStructure> str)
 }
 
 
-void IMAPMessage::extractImpl
+size_t IMAPMessage::extractImpl
 	(shared_ptr <const messagePart> p,
 	 utility::outputStream& os,
 	 utility::progressListener* progress,
@@ -279,6 +283,9 @@ void IMAPMessage::extractImpl
 
 	IMAPMessage_literalHandler literalHandler(os, progress);
 
+	if (length == 0)
+		return 0;
+
 	// Construct section identifier
 	std::ostringstream section;
 	section.imbue(std::locale::classic());
@@ -343,16 +350,62 @@ void IMAPMessage::extractImpl
 
 		// header + body
 		if ((extractFlags & EXTRACT_HEADER) && (extractFlags & EXTRACT_BODY))
-			throw exceptions::operation_not_supported();
+		{
+			// First, extract header
+			std::ostringstream header;
+			utility::outputStreamAdapter headerStream(header);
+
+			const size_t headerLength = extractImpl
+				(p, headerStream, /* progress */ NULL,
+				 /* start */ 0, /* length */ -1, extractFlags & ~EXTRACT_BODY);
+
+			size_t s = start;
+			size_t l = length;
+
+			if (s < headerLength)
+			{
+				if (l == static_cast <size_t>(-1))
+				{
+					os.write(header.str().data() + s, headerLength - s);
+				}
+				else
+				{
+					size_t headerCopyLength = l;
+
+					if (start + headerCopyLength > headerLength)
+						headerCopyLength = headerLength - start;
+
+					os.write(header.str().data() + s, headerCopyLength);
+
+					l -= headerCopyLength;
+				}
+
+				s = 0;
+			}
+			else
+			{
+				s -= headerLength;
+			}
+
+			// Then, extract body
+			return extractImpl(p, os, progress, s, l, extractFlags & ~EXTRACT_HEADER);
+		}
 		// header only
 		else if (extractFlags & EXTRACT_HEADER)
+		{
 			bodyDesc << ".MIME";   // "MIME" not "HEADER" for parts
+		}
 	}
 
 	bodyDesc << "]";
 
 	if (start != 0 || length != static_cast <size_t>(-1))
-		bodyDesc << "<" << start << "." << length << ">";
+	{
+		if (length == static_cast <size_t>(-1))
+			bodyDesc << "<" << start << "." << static_cast <unsigned int>(-1) << ">";
+		else
+			bodyDesc << "<" << start << "." << length << ">";
+	}
 
 	std::vector <std::string> fetchParams;
 	fetchParams.push_back(bodyDesc.str());
@@ -379,6 +432,8 @@ void IMAPMessage::extractImpl
 	{
 		// TODO: update the flags (eg. flag "\Seen" may have been set)
 	}
+
+	return literalHandler.getTarget()->getBytesWritten();
 }
 
 
diff --git a/src/vmime/net/imap/IMAPMessage.hpp b/src/vmime/net/imap/IMAPMessage.hpp
index 2dea1725..93545f22 100644
--- a/src/vmime/net/imap/IMAPMessage.hpp
+++ b/src/vmime/net/imap/IMAPMessage.hpp
@@ -154,7 +154,7 @@ private:
 		EXTRACT_PEEK = 0x10
 	};
 
-	void extractImpl
+	size_t extractImpl
 		(shared_ptr <const messagePart> p,
 		 utility::outputStream& os,
 		 utility::progressListener* progress,
diff --git a/src/vmime/net/imap/IMAPParser.hpp b/src/vmime/net/imap/IMAPParser.hpp
index ef239ec2..9d74b608 100644
--- a/src/vmime/net/imap/IMAPParser.hpp
+++ b/src/vmime/net/imap/IMAPParser.hpp
@@ -368,6 +368,8 @@ public:
 
 			virtual void putData(const string& chunk) = 0;
 
+			virtual size_t getBytesWritten() const = 0;
+
 		private:
 
 			utility::progressListener* m_progress;
@@ -380,7 +382,7 @@ public:
 		public:
 
 			targetString(utility::progressListener* progress, vmime::string& str)
-				: target(progress), m_string(str) { }
+				: target(progress), m_string(str), m_bytesWritten(0) { }
 
 			const vmime::string& string() const { return (m_string); }
 			vmime::string& string() { return (m_string); }
@@ -389,11 +391,18 @@ public:
 			void putData(const vmime::string& chunk)
 			{
 				m_string += chunk;
+				m_bytesWritten += chunk.length();
+			}
+
+			size_t getBytesWritten() const
+			{
+				return m_bytesWritten;
 			}
 
 		private:
 
 			vmime::string& m_string;
+			size_t m_bytesWritten;
 		};
 
 
@@ -403,7 +412,7 @@ public:
 		public:
 
 			targetStream(utility::progressListener* progress, utility::outputStream& stream)
-				: target(progress), m_stream(stream) { }
+				: target(progress), m_stream(stream), m_bytesWritten(0) { }
 
 			const utility::outputStream& stream() const { return (m_stream); }
 			utility::outputStream& stream() { return (m_stream); }
@@ -412,11 +421,18 @@ public:
 			void putData(const string& chunk)
 			{
 				m_stream.write(chunk.data(), chunk.length());
+				m_bytesWritten += chunk.length();
+			}
+
+			size_t getBytesWritten() const
+			{
+				return m_bytesWritten;
 			}
 
 		private:
 
 			utility::outputStream& m_stream;
+			size_t m_bytesWritten;
 		};
 
 
@@ -428,7 +444,7 @@ public:
 		//    . == NULL to put the literal into the response
 		//    . != NULL to redirect the literal to the specified target
 
-		virtual target* targetFor(const component& comp, const int data) = 0;
+		virtual shared_ptr <target> targetFor(const component& comp, const int data) = 0;
 	};
 
 
@@ -1142,7 +1158,7 @@ public:
 
 					if (parser.m_literalHandler != NULL)
 					{
-						literalHandler::target* target =
+						shared_ptr <literalHandler::target> target =
 							parser.m_literalHandler->targetFor(*m_component, m_data);
 
 						if (target != NULL)
@@ -1164,8 +1180,6 @@ public:
 								progress->progress(length, length);
 								progress->stop(length);
 							}
-
-							delete (target);
 						}
 						else
 						{
@@ -1196,7 +1210,7 @@ public:
 
 					if (parser.m_literalHandler != NULL)
 					{
-						literalHandler::target* target =
+						shared_ptr <literalHandler::target> target =
 							parser.m_literalHandler->targetFor(*m_component, m_data);
 
 						if (target != NULL)
@@ -1206,8 +1220,6 @@ public:
 							parser.m_progress = target->progressListener();
 							parser.readLiteral(*target, length);
 							parser.m_progress = NULL;
-
-							delete (target);
 						}
 						else
 						{