4 files changed, 72 insertions, 9 deletions
diff --git a/src/objc/smtp/MCOSMTPOperation.mm b/src/objc/smtp/MCOSMTPOperation.mm
index cd52ad13..9d96613f 100644
--- a/src/objc/smtp/MCOSMTPOperation.mm
+++ b/src/objc/smtp/MCOSMTPOperation.mm
@@ -41,6 +41,8 @@ typedef void (^CompletionType)(NSError *error);
     
     NSError * error = [NSError mco_errorWithErrorCode:MCO_NATIVE_INSTANCE->error()];
     _completionBlock(error);
+    [_completionBlock release];
+    _completionBlock = NULL;
 }
 
 - (void) setSession:(MCOSMTPSession *)session
diff --git a/src/objc/smtp/MCOSMTPSendOperation.mm b/src/objc/smtp/MCOSMTPSendOperation.mm
index e41a9f5a..8a7d5b04 100644
--- a/src/objc/smtp/MCOSMTPSendOperation.mm
+++ b/src/objc/smtp/MCOSMTPSendOperation.mm
@@ -80,16 +80,16 @@ private:
     [self start];
 }
 
+// This method needs to be duplicated from MCOSMTPOperation since _completionBlock
+// references the instance of this subclass and not the one from MCOSMTPOperation.
 - (void)operationCompleted {
     if (_completionBlock == NULL)
         return;
     
-    nativeType *op = MCO_NATIVE_INSTANCE;
-    if (op->error() == mailcore::ErrorNone) {
-        _completionBlock(nil);
-    } else {
-        _completionBlock([NSError mco_errorWithErrorCode:op->error()]);
-    }
+    NSError * error = [NSError mco_errorWithErrorCode:MCO_NATIVE_INSTANCE->error()];
+    _completionBlock(error);
+    [_completionBlock release];
+    _completionBlock = NULL;
 }
 
 - (void) bodyProgress:(unsigned int)current maximum:(unsigned int)maximum
diff --git a/src/objc/smtp/MCOSMTPSession.h b/src/objc/smtp/MCOSMTPSession.h
index 08189bed..44d37942 100644
--- a/src/objc/smtp/MCOSMTPSession.h
+++ b/src/objc/smtp/MCOSMTPSession.h
@@ -64,6 +64,15 @@
 */
 @property (nonatomic, assign, getter=isUseHeloIPEnabled) BOOL useHeloIPEnabled;
 
+/**
+ Sets logger callback. The network traffic will be sent to this block.
+ 
+ [session setConnectionLogger:^(void * connectionID, MCOConnectionLogType type, NSData * data) {
+   ...
+ }];
+ */
+@property (nonatomic, copy) MCOConnectionLogger connectionLogger;
+
 /** @name Operations */
 
 /**
diff --git a/src/objc/smtp/MCOSMTPSession.mm b/src/objc/smtp/MCOSMTPSession.mm
index de6b54cd..1538d53a 100644
--- a/src/objc/smtp/MCOSMTPSession.mm
+++ b/src/objc/smtp/MCOSMTPSession.mm
@@ -17,8 +17,34 @@
 #import "MCOAddress.h"
 #import "MCOSMTPOperation+Private.h"
 
+using namespace mailcore;
+
+@interface MCOSMTPSession ()
+
+- (void) _logWithSender:(void *)sender connectionType:(MCOConnectionLogType)logType data:(NSData *)data;
+
+@end
+
+class MCOSMTPConnectionLoggerBridge : public Object, public ConnectionLogger {
+public:
+    MCOSMTPConnectionLoggerBridge(MCOSMTPSession * session)
+    {
+        mSession = session;
+    }
+    
+    virtual void log(void * context, void * sender, ConnectionLogType logType, Data * data)
+    {
+        [mSession _logWithSender:sender connectionType:(MCOConnectionLogType)logType data:MCO_TO_OBJC(data)];
+    }
+    
+private:
+    MCOSMTPSession * mSession;
+};
+
 @implementation MCOSMTPSession {
     mailcore::SMTPAsyncSession * _session;
+    MCOConnectionLogger _connectionLogger;
+    MCOSMTPConnectionLoggerBridge * _loggerBridge;
 }
 
 #define nativeType mailcore::SMTPAsyncSession
@@ -30,13 +56,16 @@
 
 - (id)init {
     self = [super init];
-    if (self) {
-        _session = new mailcore::SMTPAsyncSession();
-    }
+    
+    _session = new mailcore::SMTPAsyncSession();
+    _loggerBridge = new MCOSMTPConnectionLoggerBridge(self);
+    
     return self;
 }
 
 - (void)dealloc {
+    MC_SAFE_RELEASE(_loggerBridge);
+    [_connectionLogger release];
     _session->release();
     [super dealloc];
 }
@@ -51,6 +80,24 @@ MCO_OBJC_SYNTHESIZE_SCALAR(NSTimeInterval, time_t, setTimeout, timeout)
 MCO_OBJC_SYNTHESIZE_BOOL(setCheckCertificateEnabled, isCheckCertificateEnabled)
 MCO_OBJC_SYNTHESIZE_BOOL(setUseHeloIPEnabled, useHeloIPEnabled)
 
+- (void) setConnectionLogger:(MCOConnectionLogger)connectionLogger
+{
+    [_connectionLogger release];
+    _connectionLogger = [connectionLogger copy];
+    
+    if (_connectionLogger != nil) {
+        _session->setConnectionLogger(_loggerBridge);
+    }
+    else {
+        _session->setConnectionLogger(NULL);
+    }
+}
+
+- (MCOConnectionLogger) connectionLogger
+{
+    return _connectionLogger;
+}
+
 #pragma mark - Operations
 
 - (MCOSMTPSendOperation *) sendOperationWithData:(NSData *)messageData
@@ -69,4 +116,9 @@ MCO_OBJC_SYNTHESIZE_BOOL(setUseHeloIPEnabled, useHeloIPEnabled)
     return result;
 }
 
+- (void) _logWithSender:(void *)sender connectionType:(MCOConnectionLogType)logType data:(NSData *)data
+{
+    _connectionLogger(sender, logType, data);
+}
+
 @end