3 files changed, 39 insertions, 9 deletions

diff --git a/arch/arm/mach-tegra/include/nvreftrack.h b/arch/arm/mach-tegra/include/nvreftrack.h
index 1a5b9c7debc2..69dfdeb1f638 100644
--- a/arch/arm/mach-tegra/include/nvreftrack.h
+++ b/arch/arm/mach-tegra/include/nvreftrack.h
@@ -98,6 +98,7 @@
 typedef enum
 {
     NvRtObjType_NvRm_NvRmMemHandle = 0,
+    NvRtObjType_NvRm_NvRmDmaHandle,
     NvRtObjType_NvRm_Num,
     NvRtObjType_NvRm_ForceWord = 0x7FFFFFFF
 } NvRtObjType_NvRm;
diff --git a/arch/arm/mach-tegra/nvrm/dispatch/nvrm_dma_dispatch.c b/arch/arm/mach-tegra/nvrm/dispatch/nvrm_dma_dispatch.c
index ff246393b5ef..003e0329cfbf 100644
--- a/arch/arm/mach-tegra/nvrm/dispatch/nvrm_dma_dispatch.c
+++ b/arch/arm/mach-tegra/nvrm/dispatch/nvrm_dma_dispatch.c
@@ -260,7 +260,6 @@ static NvError NvRmDmaAbort_dispatch_( void *InBuffer, NvU32 InSize, void *OutBu
 
     p_in = (NvRmDmaAbort_in *)InBuffer;
 
-
     NvRmDmaAbort( p_in->hDma );
 
     return err_;
@@ -300,6 +299,7 @@ static NvError NvRmDmaFree_dispatch_( void *InBuffer, NvU32 InSize, void *OutBuf
 
     p_in = (NvRmDmaFree_in *)InBuffer;
 
+    if (p_in->hDma != NULL) NvRtFreeObjRef(Ctx, NvRtObjType_NvRm_NvRmDmaHandle, p_in->hDma);
 
     NvRmDmaFree( p_in->hDma );
 
@@ -311,12 +311,25 @@ static NvError NvRmDmaAllocate_dispatch_( void *InBuffer, NvU32 InSize, void *Ou
     NvError err_ = NvSuccess;
     NvRmDmaAllocate_in *p_in;
     NvRmDmaAllocate_out *p_out;
+    NvRtObjRefHandle ref_phDma = 0;
 
     p_in = (NvRmDmaAllocate_in *)InBuffer;
     p_out = (NvRmDmaAllocate_out *)((NvU8 *)OutBuffer + OFFSET(NvRmDmaAllocate_params, out) - OFFSET(NvRmDmaAllocate_params, inout));
 
+    err_ = NvRtAllocObjRef(Ctx, &ref_phDma);
+    if (err_ != NvSuccess)
+    {
+        goto clean;
+    }
 
     p_out->ret_ = NvRmDmaAllocate( p_in->hRmDevice, &p_out->phDma, p_in->Enable32bitSwap, p_in->Priority, p_in->DmaRequestorModuleId, p_in->DmaRequestorInstanceId );
+    if ( p_out->ret_ == NvSuccess )
+    {
+        NvRtStoreObjRef(Ctx, ref_phDma, NvRtObjType_NvRm_NvRmDmaHandle, p_out->phDma);
+        ref_phDma = 0;
+    }
+clean:
+    if (ref_phDma) NvRtDiscardObjRef(Ctx, ref_phDma);
 
     return err_;
 }
@@ -330,7 +343,6 @@ static NvError NvRmDmaGetCapabilities_dispatch_( void *InBuffer, NvU32 InSize, v
     p_in = (NvRmDmaGetCapabilities_in *)InBuffer;
     p_out = (NvRmDmaGetCapabilities_out *)((NvU8 *)OutBuffer + OFFSET(NvRmDmaGetCapabilities_params, out) - OFFSET(NvRmDmaGetCapabilities_params, inout));
 
-
     p_out->ret_ = NvRmDmaGetCapabilities( p_in->hDevice, &p_in->pRmDmaCaps );
 
     return err_;
diff --git a/arch/arm/mach-tegra/nvrm_user.c b/arch/arm/mach-tegra/nvrm_user.c
index 93603aa23bc3..85b78b1389a3 100644
--- a/arch/arm/mach-tegra/nvrm_user.c
+++ b/arch/arm/mach-tegra/nvrm_user.c
@@ -34,6 +34,7 @@
 #include "nvassert.h"
 #include "nvos.h"
 #include "nvrm_memmgr.h"
+#include "nvrm_dma.h"
 #include "nvrm_ioctls.h"
 #include "nvrm_power_private.h"
 #include "mach/nvrm_linux.h"
@@ -109,6 +110,8 @@ struct nvrm_file_priv {
 
 static void client_detach(NvRtClientHandle client)
 {
+    void *ptr;
+
     if (NvRtUnregisterClient(s_RtHandle, client))
     {
         NvDispatchCtx dctx;
@@ -116,15 +119,29 @@ static void client_detach(NvRtClientHandle client)
         dctx.Rt = s_RtHandle;
         dctx.Client = client;
         dctx.PackageIdx = 0;
-
         for (;;)
         {
-            void* ptr = NvRtFreeObjRef(&dctx,
-                                       NvRtObjType_NvRm_NvRmMemHandle,
-                                       NULL);
-            if (!ptr) break;
-            NVRT_LEAK("NvRm", "NvRmMemHandle", ptr);
-            NvRmMemHandleFree(ptr);
+            NvBool found = NV_FALSE;
+
+            if ((ptr = NvRtFreeObjRef(&dctx,
+                                      NvRtObjType_NvRm_NvRmMemHandle,
+                                      NULL)) != NULL)
+            {
+                NVRT_LEAK("NvRm", "NvRmMemHandle", ptr);
+                NvRmMemHandleFree(ptr);
+                found = NV_TRUE;
+            }
+            if ((ptr = NvRtFreeObjRef(&dctx,
+                                      NvRtObjType_NvRm_NvRmDmaHandle,
+                                      NULL)) != NULL)
+            {
+                NVRT_LEAK("NvRm", "NvRmDmaHandle", ptr);
+                NvRmDmaFree(ptr);
+                found = NV_TRUE;
+            }
+
+            if (found == NV_FALSE)
+                break;
         }
 
         NvRtUnregisterClient(s_RtHandle, client);