From 29e52981780fdd84e066cbfda857d6b79489297a Mon Sep 17 00:00:00 2001
From: Julian Winkler <julian.winkler1@web.de>
Date: Fri, 19 Jul 2024 22:08:21 +0200
Subject: [PATCH] MediaCodec: handle end-of-file properly

---
 .../android_media_MediaCodec.h                |  2 ++
 .../media/android_media_MediaCodec.c          | 31 +++++++++++++------
 src/api-impl/android/media/MediaCodec.java    | 17 ++++++++--
 3 files changed, 39 insertions(+), 11 deletions(-)

diff --git a/src/api-impl-jni/generated_headers/android_media_MediaCodec.h b/src/api-impl-jni/generated_headers/android_media_MediaCodec.h
index f23fd5cf..61533b99 100644
--- a/src/api-impl-jni/generated_headers/android_media_MediaCodec.h
+++ b/src/api-impl-jni/generated_headers/android_media_MediaCodec.h
@@ -7,6 +7,8 @@
 #ifdef __cplusplus
 extern "C" {
 #endif
+#undef android_media_MediaCodec_BUFFER_FLAG_END_OF_STREAM
+#define android_media_MediaCodec_BUFFER_FLAG_END_OF_STREAM 4L
 /*
  * Class:     android_media_MediaCodec
  * Method:    native_constructor
diff --git a/src/api-impl-jni/media/android_media_MediaCodec.c b/src/api-impl-jni/media/android_media_MediaCodec.c
index 20133a97..5a71e3e6 100644
--- a/src/api-impl-jni/media/android_media_MediaCodec.c
+++ b/src/api-impl-jni/media/android_media_MediaCodec.c
@@ -425,22 +425,26 @@ JNIEXPORT void JNICALL Java_android_media_MediaCodec_native_1start(JNIEnv *env,
 
 JNIEXPORT jint JNICALL Java_android_media_MediaCodec_native_1queueInputBuffer(JNIEnv *env, jobject this, jlong codec, jobject buffer, jlong presentationTimeUs)
 {
-	jarray array_ref;
-	jbyte *array;
+	jarray array_ref = NULL;
+	jbyte *array = NULL;
 	int ret;
 	struct ATL_codec_context *ctx = _PTR(codec);
 	AVCodecContext *codec_ctx = ctx->codec;
-	AVPacket *pkt = av_packet_alloc();
-
-	pkt->size = get_nio_buffer_size(env, buffer);
-	pkt->data = get_nio_buffer(env, buffer, &array_ref, &array);
-	pkt->pts = presentationTimeUs;
+	AVPacket *pkt = NULL;
+	if (buffer) {  // buffer can be null if we're sending EOF
+		pkt = av_packet_alloc();
+		pkt->size = get_nio_buffer_size(env, buffer);
+		pkt->data = get_nio_buffer(env, buffer, &array_ref, &array);
+		pkt->pts = presentationTimeUs;
+	}
 	ret = avcodec_send_packet(codec_ctx, pkt);
 	if (ret < 0 && ret != AVERROR(EAGAIN)) {
 		fprintf(stderr, "Error while sending packet: %d = %s\n", ret, av_err2str(ret));
 	}
-	release_nio_buffer(env, array_ref, array);
-	av_packet_free(&pkt);
+	if (buffer) {
+		release_nio_buffer(env, array_ref, array);
+		av_packet_free(&pkt);
+	}
 	return ret;
 }
 
@@ -455,6 +459,14 @@ JNIEXPORT jint JNICALL Java_android_media_MediaCodec_native_1dequeueOutputBuffer
 
 	ret = avcodec_receive_frame(codec_ctx, frame);
 	if (ret < 0) {
+		if (ret == AVERROR_EOF) {
+			_SET_INT_FIELD(buffer_info, "flags", android_media_MediaCodec_BUFFER_FLAG_END_OF_STREAM);
+			_SET_INT_FIELD(buffer_info, "offset", 0);
+			_SET_INT_FIELD(buffer_info, "size", 0);
+			_SET_LONG_FIELD(buffer_info, "presentationTimeUs", 0);
+			av_frame_free(&frame);
+			return 0;
+		}
 		if (ret != AVERROR(EAGAIN)) {
 			printf("avcodec_receive_frame returned %d\n", ret);
 			printf("frame->data = %p frame->nb_samples = %d\n", frame->data[0], frame->nb_samples);
@@ -462,6 +474,7 @@ JNIEXPORT jint JNICALL Java_android_media_MediaCodec_native_1dequeueOutputBuffer
 		av_frame_free(&frame);
 		return INFO_TRY_AGAIN_LATER;
 	}
+	_SET_INT_FIELD(buffer_info, "flags", 0);
 	_SET_LONG_FIELD(buffer_info, "presentationTimeUs", frame->pts);
 
 	if (codec_ctx->codec_type == AVMEDIA_TYPE_AUDIO) {
diff --git a/src/api-impl/android/media/MediaCodec.java b/src/api-impl/android/media/MediaCodec.java
index 534cdf58..1c89a839 100644
--- a/src/api-impl/android/media/MediaCodec.java
+++ b/src/api-impl/android/media/MediaCodec.java
@@ -9,6 +9,8 @@ import android.view.Surface;
 
 public class MediaCodec {
 
+	public static final int BUFFER_FLAG_END_OF_STREAM = 0x4;
+
 	private String codecName;
 	private ByteBuffer[] inputBuffers;
 	private ByteBuffer[] outputBuffers;
@@ -104,10 +106,16 @@ public class MediaCodec {
 	private void tryProcessInputBuffer() {
 		Integer index = queuedInputBuffers.peek();
 		if (index != null) {
-			int ret = native_queueInputBuffer(native_codec, inputBuffers[index], inputBufferTimestamps[index]);
+			int ret;
+			if (index == -1) {  // end of stream
+				ret = native_queueInputBuffer(native_codec, null, 0);
+			} else {
+				ret = native_queueInputBuffer(native_codec, inputBuffers[index], inputBufferTimestamps[index]);
+			}
 			if (ret == 0) {
 				queuedInputBuffers.remove(index);
-				freeInputBuffers.add(index);
+				if (index != -1)
+					freeInputBuffers.add(index);
 			}
 		}
 	}
@@ -123,6 +131,11 @@ public class MediaCodec {
 	}
 
 	public void queueInputBuffer(int index, int offset, int size, long presentationTimeUs, int flags) {
+		if ((flags & BUFFER_FLAG_END_OF_STREAM) != 0) {
+			queuedInputBuffers.add(-1);
+			tryProcessInputBuffer();
+			return;
+		}
 		inputBufferTimestamps[index] = presentationTimeUs;
 		queuedInputBuffers.add(index);
 		tryProcessInputBuffer();