diff --git a/media/ndk/NdkMediaCodec.cpp b/media/ndk/NdkMediaCodec.cpp
index 3d1eca1..22c90e2 100644
--- a/media/ndk/NdkMediaCodec.cpp
+++ b/media/ndk/NdkMediaCodec.cpp
@@ -30,10 +30,10 @@
 
 #include <media/stagefright/foundation/ALooper.h>
 #include <media/stagefright/foundation/AMessage.h>
-#include <media/stagefright/foundation/ABuffer.h>
 
 #include <media/stagefright/MediaCodec.h>
 #include <media/stagefright/MediaErrors.h>
+#include <media/MediaCodecBuffer.h>
 
 using namespace android;
 
@@ -268,13 +268,17 @@
 
 EXPORT
 uint8_t* AMediaCodec_getInputBuffer(AMediaCodec *mData, size_t idx, size_t *out_size) {
-    android::Vector<android::sp<android::ABuffer> > abufs;
+    android::Vector<android::sp<android::MediaCodecBuffer> > abufs;
     if (mData->mCodec->getInputBuffers(&abufs) == 0) {
         size_t n = abufs.size();
         if (idx >= n) {
             ALOGE("buffer index %zu out of range", idx);
             return NULL;
         }
+        if (abufs[idx] == NULL) {
+            ALOGE("buffer index %zu is NULL", idx);
+            return NULL;
+        }
         if (out_size != NULL) {
             *out_size = abufs[idx]->capacity();
         }
@@ -286,7 +290,7 @@
 
 EXPORT
 uint8_t* AMediaCodec_getOutputBuffer(AMediaCodec *mData, size_t idx, size_t *out_size) {
-    android::Vector<android::sp<android::ABuffer> > abufs;
+    android::Vector<android::sp<android::MediaCodecBuffer> > abufs;
     if (mData->mCodec->getOutputBuffers(&abufs) == 0) {
         size_t n = abufs.size();
         if (idx >= n) {