From d4e1aa3b609fbc78c8fb3eea5371d2e7c40c11ec Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Florian=20M=C3=A4rkl?= <info@florianmaerkl.de>
Date: Mon, 21 Dec 2020 17:24:37 +0100
Subject: [PATCH] Fix Audio Parsing for PS5

---
 lib/include/chiaki/takion.h |  7 +++++++
 lib/src/audioreceiver.c     |  4 +++-
 lib/src/takion.c            | 27 ++++++++++++++++++++++++---
 3 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/lib/include/chiaki/takion.h b/lib/include/chiaki/takion.h
index dba5872..3f7e96f 100644
--- a/lib/include/chiaki/takion.h
+++ b/lib/include/chiaki/takion.h
@@ -45,6 +45,8 @@ typedef struct chiaki_takion_av_packet_t
 
 	uint64_t key_pos;
 
+	uint8_t byte_before_audio_data;
+
 	uint8_t *data; // not owned
 	size_t data_size;
 } ChiakiTakionAVPacket;
@@ -226,6 +228,11 @@ CHIAKI_EXPORT ChiakiErrorCode chiaki_takion_send_feedback_history(ChiakiTakion *
 
 CHIAKI_EXPORT ChiakiErrorCode chiaki_takion_v9_av_packet_parse(ChiakiTakionAVPacket *packet, ChiakiKeyState *key_state, uint8_t *buf, size_t buf_size);
 
+#define CHIAKI_TAKION_V12_AV_HEADER_SIZE_VIDEO 0x17
+#define CHIAKI_TAKION_V12_AV_HEADER_SIZE_AUDIO 0x13
+
+CHIAKI_EXPORT ChiakiErrorCode chiaki_takion_v12_av_packet_parse(ChiakiTakionAVPacket *packet, ChiakiKeyState *key_state, uint8_t *buf, size_t buf_size);
+
 #define CHIAKI_TAKION_V7_AV_HEADER_SIZE_BASE					0x12
 #define CHIAKI_TAKION_V7_AV_HEADER_SIZE_VIDEO_ADD				0x3
 #define CHIAKI_TAKION_V7_AV_HEADER_SIZE_NALU_INFO_STRUCTS_ADD	0x3
diff --git a/lib/src/audioreceiver.c b/lib/src/audioreceiver.c
index 9aeaa61..5808e4c 100644
--- a/lib/src/audioreceiver.c
+++ b/lib/src/audioreceiver.c
@@ -78,7 +78,9 @@ CHIAKI_EXPORT void chiaki_audio_receiver_av_packet(ChiakiAudioReceiver *audio_re
 
 	if(packet->data_size != (size_t)unit_size * (size_t)packet->units_in_frame_total)
 	{
-		CHIAKI_LOGE(audio_receiver->log, "Audio AV Packet size mismatch");
+		CHIAKI_LOGE(audio_receiver->log, "Audio AV Packet size mismatch %#llx vs %#llx",
+				(unsigned long long)packet->data_size,
+				(unsigned long long)(unit_size * packet->units_in_frame_total));
 		return;
 	}
 
diff --git a/lib/src/takion.c b/lib/src/takion.c
index 3439410..934f151 100644
--- a/lib/src/takion.c
+++ b/lib/src/takion.c
@@ -189,9 +189,11 @@ CHIAKI_EXPORT ChiakiErrorCode chiaki_takion_connect(ChiakiTakion *takion, Chiaki
 			takion->av_packet_parse = chiaki_takion_v7_av_packet_parse;
 			break;
 		case 9:
-		case 12:
 			takion->av_packet_parse = chiaki_takion_v9_av_packet_parse;
 			break;
+		case 12:
+			takion->av_packet_parse = chiaki_takion_v12_av_packet_parse;
+			break;
 		default:
 			CHIAKI_LOGE(takion->log, "Unknown Takion Protocol Version %u", (unsigned int)takion->version);
 			return CHIAKI_ERR_INVALID_DATA;
@@ -1243,7 +1245,7 @@ static void takion_handle_packet_av(ChiakiTakion *takion, uint8_t base_type, uin
 	}
 }
 
-CHIAKI_EXPORT ChiakiErrorCode chiaki_takion_v9_av_packet_parse(ChiakiTakionAVPacket *packet, ChiakiKeyState *key_state, uint8_t *buf, size_t buf_size)
+static ChiakiErrorCode av_packet_parse(bool v12, ChiakiTakionAVPacket *packet, ChiakiKeyState *key_state, uint8_t *buf, size_t buf_size)
 {
 	memset(packet, 0, sizeof(ChiakiTakionAVPacket));
 
@@ -1261,7 +1263,9 @@ CHIAKI_EXPORT ChiakiErrorCode chiaki_takion_v9_av_packet_parse(ChiakiTakionAVPac
 
 	uint8_t *av = buf+1;
 	size_t av_size = buf_size-1;
-	size_t av_header_size = packet->is_video ? CHIAKI_TAKION_V9_AV_HEADER_SIZE_VIDEO : CHIAKI_TAKION_V9_AV_HEADER_SIZE_AUDIO;
+	size_t av_header_size = v12
+		? (packet->is_video ? CHIAKI_TAKION_V12_AV_HEADER_SIZE_VIDEO : CHIAKI_TAKION_V12_AV_HEADER_SIZE_AUDIO)
+		: (packet->is_video ? CHIAKI_TAKION_V9_AV_HEADER_SIZE_VIDEO : CHIAKI_TAKION_V9_AV_HEADER_SIZE_AUDIO);
 	if(av_size < av_header_size + 1)
 		return CHIAKI_ERR_BUF_TOO_SMALL;
 
@@ -1320,12 +1324,29 @@ CHIAKI_EXPORT ChiakiErrorCode chiaki_takion_v9_av_packet_parse(ChiakiTakionAVPac
 		av_size -= 3;
 	}
 
+	if(v12 && !packet->is_video)
+	{
+		packet->byte_before_audio_data = *av;
+		av += 1;
+		av_size -= 1;
+	}
+
 	packet->data = av;
 	packet->data_size = av_size;
 
 	return CHIAKI_ERR_SUCCESS;
 }
 
+CHIAKI_EXPORT ChiakiErrorCode chiaki_takion_v9_av_packet_parse(ChiakiTakionAVPacket *packet, ChiakiKeyState *key_state, uint8_t *buf, size_t buf_size)
+{
+	return av_packet_parse(false, packet, key_state, buf, buf_size);
+}
+
+CHIAKI_EXPORT ChiakiErrorCode chiaki_takion_v12_av_packet_parse(ChiakiTakionAVPacket *packet, ChiakiKeyState *key_state, uint8_t *buf, size_t buf_size)
+{
+	return av_packet_parse(true, packet, key_state, buf, buf_size);
+}
+
 CHIAKI_EXPORT ChiakiErrorCode chiaki_takion_v7_av_packet_format_header(uint8_t *buf, size_t buf_size, size_t *header_size_out, ChiakiTakionAVPacket *packet)
 {
 	size_t header_size = CHIAKI_TAKION_V7_AV_HEADER_SIZE_BASE;