2 files changed, 33 insertions, 3 deletions

diff --git a/media/libstagefright/rtsp/ARTSPConnection.cpp b/media/libstagefright/rtsp/ARTSPConnection.cpp
index b398c9d..bd0e491 100644
--- a/media/libstagefright/rtsp/ARTSPConnection.cpp
+++ b/media/libstagefright/rtsp/ARTSPConnection.cpp
@@ -56,6 +56,9 @@ ARTSPConnection::ARTSPConnection(bool uidValid, uid_t uid)
 ARTSPConnection::~ARTSPConnection() {
     if (mSocket >= 0) {
         LOGE("Connection is still open, closing the socket.");
+        if (mUIDValid) {
+            HTTPBase::UnRegisterSocketUserTag(mSocket);
+        }
         close(mSocket);
         mSocket = -1;
     }
@@ -202,6 +205,9 @@ void ARTSPConnection::onConnect(const sp<AMessage> &msg) {
     ++mConnectionID;
 
     if (mState != DISCONNECTED) {
+        if (mUIDValid) {
+            HTTPBase::UnRegisterSocketUserTag(mSocket);
+        }
         close(mSocket);
         mSocket = -1;
 
@@ -251,7 +257,8 @@ void ARTSPConnection::onConnect(const sp<AMessage> &msg) {
     mSocket = socket(AF_INET, SOCK_STREAM, 0);
 
     if (mUIDValid) {
-        HTTPBase::RegisterSocketUser(mSocket, mUID);
+        HTTPBase::RegisterSocketUserTag(mSocket, mUID,
+                                        (uint32_t)*(uint32_t*) "RTSP");
     }
 
     MakeSocketBlocking(mSocket, false);
@@ -279,6 +286,9 @@ void ARTSPConnection::onConnect(const sp<AMessage> &msg) {
         reply->setInt32("result", -errno);
         mState = DISCONNECTED;
 
+        if (mUIDValid) {
+            HTTPBase::UnRegisterSocketUserTag(mSocket);
+        }
         close(mSocket);
         mSocket = -1;
     } else {
@@ -294,6 +304,9 @@ void ARTSPConnection::onConnect(const sp<AMessage> &msg) {
 
 void ARTSPConnection::onDisconnect(const sp<AMessage> &msg) {
     if (mState == CONNECTED || mState == CONNECTING) {
+        if (mUIDValid) {
+            HTTPBase::UnRegisterSocketUserTag(mSocket);
+        }
         close(mSocket);
         mSocket = -1;
 
@@ -358,6 +371,9 @@ void ARTSPConnection::onCompleteConnection(const sp<AMessage> &msg) {
         reply->setInt32("result", -err);
 
         mState = DISCONNECTED;
+        if (mUIDValid) {
+            HTTPBase::UnRegisterSocketUserTag(mSocket);
+        }
         close(mSocket);
         mSocket = -1;
     } else {
diff --git a/media/libstagefright/rtsp/MyHandler.h b/media/libstagefright/rtsp/MyHandler.h
index 71d68f6..8128813 100644
--- a/media/libstagefright/rtsp/MyHandler.h
+++ b/media/libstagefright/rtsp/MyHandler.h
@@ -545,6 +545,12 @@ struct MyHandler : public AHandler {
                 if (result != OK) {
                     if (track) {
                         if (!track->mUsingInterleavedTCP) {
+                            // Clear the tag
+                            if (mUIDValid) {
+                                HTTPBase::UnRegisterSocketUserTag(track->mRTPSocket);
+                                HTTPBase::UnRegisterSocketUserTag(track->mRTCPSocket);
+                            }
+
                             close(track->mRTPSocket);
                             close(track->mRTCPSocket);
                         }
@@ -618,6 +624,12 @@ struct MyHandler : public AHandler {
                     if (!info->mUsingInterleavedTCP) {
                         mRTPConn->removeStream(info->mRTPSocket, info->mRTCPSocket);
 
+                        // Clear the tag
+                        if (mUIDValid) {
+                            HTTPBase::UnRegisterSocketUserTag(info->mRTPSocket);
+                            HTTPBase::UnRegisterSocketUserTag(info->mRTCPSocket);
+                        }
+
                         close(info->mRTPSocket);
                         close(info->mRTCPSocket);
                     }
@@ -1181,8 +1193,10 @@ private:
                     &info->mRTPSocket, &info->mRTCPSocket, &rtpPort);
 
             if (mUIDValid) {
-                HTTPBase::RegisterSocketUser(info->mRTPSocket, mUID);
-                HTTPBase::RegisterSocketUser(info->mRTCPSocket, mUID);
+                HTTPBase::RegisterSocketUserTag(info->mRTPSocket, mUID,
+                                                (uint32_t)*(uint32_t*) "RTP_");
+                HTTPBase::RegisterSocketUserTag(info->mRTCPSocket, mUID,
+                                                (uint32_t)*(uint32_t*) "RTP_");
             }
 
             request.append("Transport: RTP/AVP/UDP;unicast;client_port=");