From 7d66360ea62bfc0354a78f91c784c447ce3be545 Mon Sep 17 00:00:00 2001
From: Eugene Rysaj <eugene@timecode.co>
Date: Tue, 29 Apr 2014 17:29:56 +0300
Subject: [PATCH] fix GPUImageRawDataOutput memory access issues when working
 with CVPixelBuffer-backed framebuffer

---
 framework/Source/GPUImageFramebuffer.h   |  2 ++
 framework/Source/GPUImageFramebuffer.m   | 38 +++++++++++++++++++++---
 framework/Source/GPUImageRawDataOutput.m |  3 ++
 3 files changed, 39 insertions(+), 4 deletions(-)

diff --git a/framework/Source/GPUImageFramebuffer.h b/framework/Source/GPUImageFramebuffer.h
index 683494df7..5cf20dd3e 100644
--- a/framework/Source/GPUImageFramebuffer.h
+++ b/framework/Source/GPUImageFramebuffer.h
@@ -50,6 +50,8 @@ typedef struct GPUTextureOptions {
 - (void)restoreRenderTarget;
 
 // Raw data bytes
+- (void)lockForReading;
+- (void)unlockAfterReading;
 - (NSUInteger)bytesPerRow;
 - (GLubyte *)byteBuffer;
 
diff --git a/framework/Source/GPUImageFramebuffer.m b/framework/Source/GPUImageFramebuffer.m
index 2868c17fe..048a30768 100644
--- a/framework/Source/GPUImageFramebuffer.m
+++ b/framework/Source/GPUImageFramebuffer.m
@@ -7,6 +7,7 @@ @interface GPUImageFramebuffer()
 #if TARGET_IPHONE_SIMULATOR || TARGET_OS_IPHONE
     CVPixelBufferRef renderTarget;
     CVOpenGLESTextureRef renderTexture;
+    NSUInteger readLockCount;
 #else
 #endif
     NSUInteger framebufferReferenceCount;
@@ -331,7 +332,7 @@ - (CGImageRef)newCGImageFromFramebufferContents;
             //        glFlush();
             glFinish();
             CFRetain(renderTarget); // I need to retain the pixel buffer here and release in the data source callback to prevent its bytes from being prematurely deallocated during a photo write operation
-            CVPixelBufferLockBaseAddress(renderTarget, 0);
+            [self lockForReading];
             rawImagePixels = (GLubyte *)CVPixelBufferGetBaseAddress(renderTarget);
             dataProvider = CGDataProviderCreateWithData((__bridge_retained void*)self, rawImagePixels, paddedBytesForImage, dataProviderUnlockCallback);
             [[GPUImageContext sharedFramebufferCache] addFramebufferToActiveImageCaptureList:self]; // In case the framebuffer is swapped out on the filter, need to have a strong reference to it somewhere for it to hang on while the image is in existence
@@ -373,7 +374,7 @@ - (CGImageRef)newCGImageFromFramebufferContents;
 - (void)restoreRenderTarget;
 {
 #if TARGET_IPHONE_SIMULATOR || TARGET_OS_IPHONE
-    CVPixelBufferUnlockBaseAddress(renderTarget, 0);
+    [self unlockAfterReading];
     CFRelease(renderTarget);
 #else
 #endif
@@ -382,6 +383,35 @@ - (void)restoreRenderTarget;
 #pragma mark -
 #pragma mark Raw data bytes
 
+- (void)lockForReading
+{
+#if TARGET_IPHONE_SIMULATOR || TARGET_OS_IPHONE
+    if ([GPUImageContext supportsFastTextureUpload])
+    {
+        if (readLockCount == 0)
+        {
+            CVPixelBufferLockBaseAddress(renderTarget, 0);
+        }
+        readLockCount++;
+    }
+#endif
+}
+
+- (void)unlockAfterReading
+{
+#if TARGET_IPHONE_SIMULATOR || TARGET_OS_IPHONE
+    if ([GPUImageContext supportsFastTextureUpload])
+    {
+        NSAssert(readLockCount > 0, @"Unbalanced call to -[GPUImageFramebuffer unlockAfterReading]");
+        readLockCount--;
+        if (readLockCount == 0)
+        {
+            CVPixelBufferUnlockBaseAddress(renderTarget, 0);
+        }
+    }
+#endif
+}
+
 - (NSUInteger)bytesPerRow;
 {
     if ([GPUImageContext supportsFastTextureUpload])
@@ -401,9 +431,9 @@ - (NSUInteger)bytesPerRow;
 - (GLubyte *)byteBuffer;
 {
 #if TARGET_IPHONE_SIMULATOR || TARGET_OS_IPHONE
-    CVPixelBufferLockBaseAddress(renderTarget, 0);
+    [self lockForReading];
     GLubyte * bufferBytes = CVPixelBufferGetBaseAddress(renderTarget);
-    CVPixelBufferUnlockBaseAddress(renderTarget, 0);
+    [self unlockAfterReading];
     return bufferBytes;
 #else
     return NULL; // TODO: do more with this on the non-texture-cache side
diff --git a/framework/Source/GPUImageRawDataOutput.m b/framework/Source/GPUImageRawDataOutput.m
index f2cdb785b..18101e2cd 100755
--- a/framework/Source/GPUImageRawDataOutput.m
+++ b/framework/Source/GPUImageRawDataOutput.m
@@ -106,6 +106,8 @@ - (void)renderAtInternalSize;
     if(lockNextFramebuffer)
     {
         retainedFramebuffer = outputFramebuffer;
+        [retainedFramebuffer lock];
+        [retainedFramebuffer lockForReading];
         lockNextFramebuffer = NO;
     }
 
@@ -297,6 +299,7 @@ - (void)lockFramebufferForReading;
 
 - (void)unlockFramebufferAfterReading;
 {
+    [retainedFramebuffer unlockAfterReading];
     [retainedFramebuffer unlock];
     retainedFramebuffer = nil;
 }