1 files changed, 29 insertions, 7 deletions
diff --git a/net/curvecp/client_packetizer.cc b/net/curvecp/client_packetizer.cc
index 72bf4a0..77ddfda 100644
--- a/net/curvecp/client_packetizer.cc
+++ b/net/curvecp/client_packetizer.cc
@@ -33,7 +33,7 @@ ClientPacketizer::ClientPacketizer()
     : Packetizer(),
       next_state_(NONE),
       listener_(NULL),
-      user_callback_(NULL),
+      old_user_callback_(NULL),
       current_address_(NULL),
       hello_attempts_(0),
       initiate_sent_(false),
@@ -52,7 +52,23 @@ ClientPacketizer::~ClientPacketizer() {
 int ClientPacketizer::Connect(const AddressList& server,
                               Packetizer::Listener* listener,
                               OldCompletionCallback* callback) {
-  DCHECK(!user_callback_);
+  DCHECK(!old_user_callback_);
+  DCHECK(!socket_.get());
+  DCHECK(!listener_);
+
+  listener_ = listener;
+
+  addresses_ = server;
+
+  old_user_callback_ = callback;
+  next_state_ = LOOKUP_COOKIE;
+
+  return DoLoop(OK);
+}
+int ClientPacketizer::Connect(const AddressList& server,
+                              Packetizer::Listener* listener,
+                              const net::CompletionCallback& callback) {
+  DCHECK(user_callback_.is_null());
   DCHECK(!socket_.get());
   DCHECK(!listener_);
 
@@ -279,11 +295,17 @@ int ClientPacketizer::DoConnected(int rv) {
 
 void ClientPacketizer::DoCallback(int result) {
   DCHECK_NE(result, ERR_IO_PENDING);
-  DCHECK(user_callback_);
-
-  OldCompletionCallback* callback = user_callback_;
-  user_callback_ = NULL;
-  callback->Run(result);
+  DCHECK(old_user_callback_ || !user_callback_.is_null());
+
+  if (old_user_callback_) {
+    OldCompletionCallback* callback = old_user_callback_;
+    old_user_callback_ = NULL;
+    callback->Run(result);
+  } else {
+    CompletionCallback callback = user_callback_;
+    user_callback_.Reset();
+    callback.Run(result);
+  }
 }
 
 int ClientPacketizer::ConnectNextAddress() {