Improved reverse subscription operation

ritchiecarroll · ritchiecarroll · commit ff0dd815f9b3 · 2023-03-24T17:35:18.000-05:00
diff --git a/src/lib/transport/DataSubscriber.cpp b/src/lib/transport/DataSubscriber.cpp
@@ -41,6 +41,30 @@ using namespace sttp::transport::tssc;
 constexpr float32_t MissingCacheWarningInterval = 20.0;
 const DataSubscriberPtr DataSubscriber::NullPtr = nullptr;
 
+class ReverseConnection;
+typedef SharedPtr<ReverseConnection> ReverseConnectionPtr;
+
+class ReverseConnection : public boost::enable_shared_from_this<ReverseConnection>
+{
+private:
+    TcpSocket m_socket;
+
+public:
+    explicit ReverseConnection(IOContext& service) :
+        m_socket(service)
+    {
+    }
+
+    TcpSocket& GetSocket()
+    {
+        return m_socket;
+    }
+
+    static ReverseConnectionPtr New(IOContext& service) {
+        return NewSharedPtr<ReverseConnection, IOContext&>(service);
+    }
+};
+
 DataSubscriber::DataSubscriber() :
     m_subscriberID(Empty::Guid),
     m_compressPayloadData(true),
@@ -127,7 +151,8 @@ void DataSubscriber::RunCommandChannelResponseThread()
         ReadPayloadHeader(error, bytesTransferred);
     });
 
-    m_commandChannelService.run();
+    if (!m_listening)
+        m_commandChannelService.run();
 }
 
 // Callback for async read of the payload header.
@@ -1096,7 +1121,7 @@ void DataSubscriber::Connect(const string& hostname, const uint16_t port, const
     EstablishConnection(hostEndpoint, false);
 }
 
-void DataSubscriber::Listen(const sttp::TcpEndPoint& endpoint)
+void DataSubscriber::Listen(const sttp::TcpEndPoint& endPoint)
 {
     // This function fails by exception, consumers should try/catch calls to Start
     if (IsListening())
@@ -1115,14 +1140,14 @@ void DataSubscriber::Listen(const sttp::TcpEndPoint& endpoint)
 #endif
 
     // TODO: Add TLS implementation options
-    m_clientAcceptor = TcpAcceptor(m_commandChannelService, endpoint, false); //-V601
+    m_clientAcceptor = TcpAcceptor(m_commandChannelService, endPoint, false); //-V601
 
     // Run command channel accept thread
     m_commandChannelAcceptThread = Thread([this]
-        {
-            StartAccept();
-            m_commandChannelService.run();
-        });
+    {
+        StartAccept();
+        m_commandChannelService.run();
+    });
 
     m_listening = true;
 }
@@ -1139,49 +1164,53 @@ void DataSubscriber::Listen(const std::string& networkInterfaceIP, const uint16_
 
 void DataSubscriber::StartAccept()
 {
-    TcpSocket commandChannelSocket(m_commandChannelService);
+    ReverseConnectionPtr reverseConnection = ReverseConnection::New(m_commandChannelService);
 
-    m_clientAcceptor.async_accept(commandChannelSocket, [this, &commandChannelSocket]<typename T0>(T0 && error)
+    m_clientAcceptor.async_accept(reverseConnection->GetSocket(), [this, reverseConnection]<typename T0>(T0 && error)
     {
-        AcceptConnection(commandChannelSocket, error);
+        AcceptConnection(&reverseConnection->GetSocket(), error);
     });
 }
 
-void DataSubscriber::AcceptConnection(TcpSocket& commandChannelSocket, const ErrorCode& error)
+void DataSubscriber::AcceptConnection(TcpSocket* commandChannelSocket, const ErrorCode& error)
 {
     if (!IsListening())
         return;
 
+    // TODO: For secured connections, validate certificate and IP information here to assign subscriberID
+    const TcpEndPoint endPoint = commandChannelSocket->remote_endpoint();
+    string address = "<unknown>";
+    string errorMessage = "closed.";
+
+    auto closeSocket = [&]
+    {
+        try
+        {
+            address = ResolveDNSName(m_commandChannelService, endPoint);
+            commandChannelSocket->close();
+        }
+        catch (SystemError& ex)
+        {
+            errorMessage = "close error: " + string(ex.what());
+        }
+        catch (...)
+        {
+            errorMessage = "close error: " + boost::current_exception_diagnostic_information(true);
+        }
+    };
+
     if (error)
     {
-        DispatchErrorMessage("Error while attempting to accept DataPublisher connection for reverse connection: " + string(SystemError(error).what()));
+        closeSocket();
+        DispatchErrorMessage("Error while attempting to accept DataPublisher connection for reverse connection: " + string(SystemError(error).what()) + " - connection " + errorMessage);
     }
     else
     {
-        // TODO: For secured connections, validate certificate and IP information here to assign subscriberID
-        const TcpEndPoint endPoint = commandChannelSocket.remote_endpoint();
-
         // Will only accept one active connection at a time, this may be indicative
         // of a rouge connection attempt - consumer should log warning
         if (IsConnected())
         {
-            string address = "<unknown>";
-            string errorMessage = "closed.";
-
-            try
-            {
-                address = ResolveDNSName(m_commandChannelService, endPoint);
-                commandChannelSocket.close();
-            }
-            catch (SystemError& ex)
-            {
-                errorMessage = "close error: " + string(ex.what());
-            }
-            catch (...)
-            {
-                errorMessage = "close error: " + boost::current_exception_diagnostic_information(true);
-            }
-
+            closeSocket();
             DispatchErrorMessage("WARNING: Duplicate connection attempt detected from: \"" + address + "\". Existing data publisher connection already established, data subscriber will only accept one connection at a time - connection " + errorMessage);
         }
         else
@@ -1194,7 +1223,7 @@ void DataSubscriber::AcceptConnection(TcpSocket& commandChannelSocket, const Err
             SetupConnection();
 
             // Hold on to primary socket
-            m_commandChannelSocket = std::move(commandChannelSocket);
+            m_commandChannelSocket = std::move(*commandChannelSocket);
 
             // Create new command channel
             EstablishConnection(endPoint, true);
diff --git a/src/lib/transport/DataSubscriber.h b/src/lib/transport/DataSubscriber.h
@@ -153,7 +153,7 @@ namespace sttp::transport
         void ReadPacket(const ErrorCode& error, size_t bytesTransferred);
         void WriteHandler(const ErrorCode& error, size_t bytesTransferred);
         void StartAccept();
-        void AcceptConnection(TcpSocket& commandChannelSocket, const ErrorCode& error);
+        void AcceptConnection(TcpSocket* commandChannelSocket, const ErrorCode& error);
 
         // Server response handlers
         void ProcessServerResponse(uint8_t* buffer, uint32_t offset, uint32_t length);
@@ -274,7 +274,7 @@ namespace sttp::transport
         void Connect(const std::string& hostname, uint16_t port);
 
         // Establish a reverse listening connection for subscriber using specified connection info
-        void Listen(const sttp::TcpEndPoint& endpoint);
+        void Listen(const sttp::TcpEndPoint& endPoint);
         void Listen(uint16_t port, bool ipV6 = false);                       // Bind to default NIC
         void Listen(const std::string& networkInterfaceIP, uint16_t port);   // Bind to specified NIC IP, format determines IP version
 
diff --git a/src/lib/transport/SubscriberInstance.cpp b/src/lib/transport/SubscriberInstance.cpp
@@ -64,6 +64,7 @@ SubscriberInstance::SubscriberInstance() :
     m_subscriber->RegisterSubscriptionUpdatedCallback(&HandleSubscriptionUpdated);
     m_subscriber->RegisterNewMeasurementsCallback(&HandleNewMeasurements);
     m_subscriber->RegisterConfigurationChangedCallback(&HandleConfigurationChanged);
+    m_subscriber->RegisterConnectionEstablishedCallback(&HandleConnectionEstablished);
     m_subscriber->RegisterConnectionTerminatedCallback(&HandleConnectionTerminated);
 }
 
@@ -186,6 +187,35 @@ void SubscriberInstance::SetMetadataFilters(const std::string& metadataFilters)
     m_metadataFilters = metadataFilters;
 }
 
+void SubscriberInstance::HandleConnect()
+{
+    // If automatically parsing metadata, request metadata upon successful connection,
+    // after metadata is received the SubscriberInstance will then initiate subscribe;
+    // otherwise, subscribe is initiated immediately (when auto subscribe requested)
+    if (m_autoParseMetadata)
+        SendMetadataRefreshCommand();
+    else
+        m_subscriber->Subscribe();
+}
+
+void SubscriberInstance::SetupSubscriptionInfo()
+{
+    m_subscriptionInfo = CreateSubscriptionInfo();
+
+    if (!m_startTime.empty() && !m_stopTime.empty())
+    {
+        m_subscriber->RegisterProcessingCompleteCallback(&HandleProcessingComplete);
+        m_subscriptionInfo.StartTime = m_startTime;
+        m_subscriptionInfo.StopTime = m_stopTime;
+    }
+
+    if (m_udpPort > 0)
+    {
+        m_subscriptionInfo.UdpDataChannel = true;
+        m_subscriptionInfo.DataChannelLocalPort = m_udpPort;
+    }
+}
+
 bool SubscriberInstance::Connect()
 {
     if (IsConnected())
@@ -206,20 +236,7 @@ bool SubscriberInstance::Connect()
     if (m_port == 0)
         throw SubscriberException("No port specified for connection: call Initialize before Connect");
 
-    m_subscriptionInfo = CreateSubscriptionInfo();
-
-    if (!m_startTime.empty() && !m_stopTime.empty())
-    {
-        m_subscriber->RegisterProcessingCompleteCallback(&HandleProcessingComplete);
-        m_subscriptionInfo.StartTime = m_startTime;
-        m_subscriptionInfo.StopTime = m_stopTime;
-    }
-
-    if (m_udpPort > 0)
-    {
-        m_subscriptionInfo.UdpDataChannel = true;
-        m_subscriptionInfo.DataChannelLocalPort = m_udpPort;
-    }
+    SetupSubscriptionInfo();
 
     // Connect and subscribe to publisher
     const int result = connector.Connect(*m_subscriber, m_subscriptionInfo);
@@ -239,14 +256,7 @@ bool SubscriberInstance::Connect()
         }
 
         ConnectionEstablished();
-
-        // If automatically parsing metadata, request metadata upon successful connection,
-        // after metadata is handled the SubscriberInstance will then initiate subscribe;
-        // otherwise, initiate subscribe immediately
-        if (m_autoParseMetadata)
-            SendMetadataRefreshCommand();
-        else
-            m_subscriber->Subscribe();
+        HandleConnect();
 
         return true;
     }
@@ -280,6 +290,12 @@ bool SubscriberInstance::Listen(const sttp::TcpEndPoint& endPoint)
 
     try
     {
+        m_hostname = endPoint.address().to_string();
+        m_port = endPoint.port();;
+
+        SetupSubscriptionInfo();
+        m_subscriber->SetSubscriptionInfo(m_subscriptionInfo);
+
         m_subscriber->Listen(endPoint);
     }
     catch (SubscriberException& ex)
@@ -301,7 +317,7 @@ bool SubscriberInstance::Listen(const sttp::TcpEndPoint& endPoint)
     return m_subscriber->IsListening();
 }
 
-bool SubscriberInstance::ListenStart(const uint16_t port, const bool ipV6)
+bool SubscriberInstance::Listen(const uint16_t port, const bool ipV6)
 {
     return Listen(TcpEndPoint(ipV6 ? tcp::v6() : tcp::v4(), port));
 }
@@ -1260,10 +1276,12 @@ void SubscriberInstance::HistoricalReadComplete()
 
 void SubscriberInstance::ConnectionEstablished()
 {
+    StatusMessage("Connection from \"" + GetConnectionID() + "\" established.");
 }
 
 void SubscriberInstance::ConnectionTerminated()
 {
+    StatusMessage("Connection from \"" + GetConnectionID() + "\" terminated.");
 }
 
 void SubscriberInstance::GetAssemblyInfo(std::string& source, std::string& version, std::string& updatedOn) const
@@ -1291,14 +1309,7 @@ void SubscriberInstance::HandleResubscribe(DataSubscriber* source)
     if (source->IsConnected())
     {
         instance->ConnectionEstablished();
-
-        // If automatically parsing metadata, request metadata upon successful connection,
-        // after metadata is handled the SubscriberInstance will then initiate subscribe;
-        // otherwise, initiate subscribe immediately
-        if (instance->m_autoParseMetadata)
-            instance->SendMetadataRefreshCommand();
-        else
-            source->Subscribe();
+        instance->HandleConnect();
     }
     else
     {
@@ -1410,6 +1421,7 @@ void SubscriberInstance::HandleConnectionEstablished(const DataSubscriber* sourc
         return;
 
     instance->ConnectionEstablished();
+    instance->HandleConnect();
 }
 
 void SubscriberInstance::HandleConnectionTerminated(const DataSubscriber* source)
diff --git a/src/lib/transport/SubscriberInstance.h b/src/lib/transport/SubscriberInstance.h
@@ -56,6 +56,8 @@ namespace sttp::transport
         Mutex m_configurationUpdateLock;
         void* m_userData;
 
+        void HandleConnect();
+        void SetupSubscriptionInfo();
         void SendMetadataRefreshCommand();        
 
         static void ConstructConfigurationFrames(const sttp::StringMap<DeviceMetadataPtr>& devices, const std::unordered_map<Guid, MeasurementMetadataPtr>& measurements, sttp::StringMap<ConfigurationFramePtr>& configurationFrames);
@@ -201,7 +203,7 @@ namespace sttp::transport
         // Returns true if subscriber was successfully started
         // Throws SubscriberException for implementation logic errors
         virtual bool Listen(const sttp::TcpEndPoint& endPoint);
-        virtual bool ListenStart(uint16_t port, bool ipV6 = false);                       // Bind to default NIC
+        virtual bool Listen(uint16_t port, bool ipV6 = false);                       // Bind to default NIC
         virtual bool Listen(const std::string& networkInterfaceIP, uint16_t port);   // Bind to specified NIC IP, format determines IP version
 
         // Disconnects from the STTP publisher, normal or reverse connection mode
