Merge pull request #2359 from Fluorohydride/patch-replay

update class Replay

gframe/replay.cpp

@@ -5,11 +5,13 @@
 
 namespace ygo {
 
-Replay::Replay() {
-	is_recording = false;
-	is_replaying = false;
-	replay_data = new unsigned char[0x20000];
-	comp_data = new unsigned char[0x2000];
+Replay::Replay()
+	: fp(nullptr), pheader(), pdata(nullptr), replay_size(0), comp_size(0), is_recording(false), is_replaying(false) {
+	#ifdef _WIN32
+		recording_fp = nullptr;
+	#endif
+	replay_data = new unsigned char[MAX_REPLAY_SIZE];
+	comp_data = new unsigned char[MAX_COMP_SIZE];
 }
 Replay::~Replay() {
 	delete[] replay_data;
@@ -32,6 +34,9 @@ void Replay::BeginRecord() {
 		return;
 #endif
 	pdata = replay_data;
+	replay_size = 0;
+	comp_size = 0;
+	is_replaying = false;
 	is_recording = true;
 }
 void Replay::WriteHeader(ReplayHeader& header) {
@@ -44,9 +49,11 @@ void Replay::WriteHeader(ReplayHeader& header) {
 	fflush(fp);
 #endif
 }
-void Replay::WriteData(const void* data, unsigned int length, bool flush) {
+void Replay::WriteData(const void* data, int length, bool flush) {
 	if(!is_recording)
 		return;
+	if (length < 0 || (pdata - replay_data) + length > MAX_REPLAY_SIZE)
+		return;
 	memcpy(pdata, data, length);
 	pdata += length;
 #ifdef _WIN32
@@ -61,6 +68,8 @@ void Replay::WriteData(const void* data, unsigned int length, bool flush) {
 void Replay::WriteInt32(int data, bool flush) {
 	if(!is_recording)
 		return;
+	if ((pdata - replay_data) + 4 > MAX_REPLAY_SIZE)
+		return;
 	*((int*)(pdata)) = data;
 	pdata += 4;
 #ifdef _WIN32
@@ -75,6 +84,8 @@ void Replay::WriteInt32(int data, bool flush) {
 void Replay::WriteInt16(short data, bool flush) {
 	if(!is_recording)
 		return;
+	if ((pdata - replay_data) + 2 > MAX_REPLAY_SIZE)
+		return;
 	*((short*)(pdata)) = data;
 	pdata += 2;
 #ifdef _WIN32
@@ -89,6 +100,8 @@ void Replay::WriteInt16(short data, bool flush) {
 void Replay::WriteInt8(char data, bool flush) {
 	if(!is_recording)
 		return;
+	if ((pdata - replay_data) + 1 > MAX_REPLAY_SIZE)
+		return;
 	*pdata = data;
 	pdata++;
 #ifdef _WIN32
@@ -116,11 +129,19 @@ void Replay::EndRecord() {
 #else
 	fclose(fp);
 #endif
-	pheader.datasize = pdata - replay_data;
+	if(pdata - replay_data > 0 && pdata - replay_data <= MAX_REPLAY_SIZE)
+		replay_size = pdata - replay_data;
+	else
+		replay_size = 0;
+	pheader.datasize = replay_size;
 	pheader.flag |= REPLAY_COMPRESSED;
 	size_t propsize = 5;
-	comp_size = 0x1000;
-	LzmaCompress(comp_data, &comp_size, replay_data, pdata - replay_data, pheader.props, &propsize, 5, 1 << 24, 3, 0, 2, 32, 1);
+	comp_size = MAX_COMP_SIZE;
+	int ret = LzmaCompress(comp_data, &comp_size, replay_data, replay_size, pheader.props, &propsize, 5, 1 << 24, 3, 0, 2, 32, 1);
+	if (ret != SZ_OK) {
+		*((int*)(comp_data)) = ret;
+		comp_size = sizeof(ret);
+	}
 	is_recording = false;
 }
 void Replay::SaveReplay(const wchar_t* name) {
@@ -162,22 +183,33 @@ bool Replay::OpenReplay(const wchar_t* name) {
 	}
 	if(!fp)
 		return false;
+
+	pdata = replay_data;
+	is_recording = false;
+	is_replaying = false;
+	replay_size = 0;
+	comp_size = 0;
 	if(fread(&pheader, sizeof(pheader), 1, fp) < 1) {
 		fclose(fp);
 		return false;
 	}
 	if(pheader.flag & REPLAY_COMPRESSED) {
-		comp_size = fread(comp_data, 1, 0x1000, fp);
+		comp_size = fread(comp_data, 1, MAX_COMP_SIZE, fp);
 		fclose(fp);
+		if ((int)pheader.datasize < 0 && (int)pheader.datasize > MAX_REPLAY_SIZE)
+			return false;
 		replay_size = pheader.datasize;
-		if(LzmaUncompress(replay_data, &replay_size, comp_data, &comp_size, pheader.props, 5) != SZ_OK)
+		if (LzmaUncompress(replay_data, &replay_size, comp_data, &comp_size, pheader.props, 5) != SZ_OK)
+			return false;
+		if (replay_size != pheader.datasize) {
+			replay_size = 0;
 			return false;
+		}
 	} else {
-		comp_size = fread(replay_data, 1, 0x20000, fp);
+		replay_size = fread(replay_data, 1, MAX_REPLAY_SIZE, fp);
 		fclose(fp);
-		replay_size = comp_size;
+		comp_size = 0;
 	}
-	pdata = replay_data;
 	is_replaying = true;
 	return true;
 }
@@ -196,7 +228,7 @@ bool Replay::CheckReplay(const wchar_t* name) {
 	ReplayHeader rheader;
 	size_t count = fread(&rheader, sizeof(ReplayHeader), 1, rfp);
 	fclose(rfp);
-	return count == 1 && rheader.id == 0x31707279 && rheader.version >= 0x12d0;
+	return count == 1 && rheader.id == 0x31707279 && rheader.version >= 0x12d0u;
 }
 bool Replay::DeleteReplay(const wchar_t* name) {
 	wchar_t fname[256];
@@ -245,7 +277,7 @@ void Replay::ReadName(wchar_t* data) {
 	ReadData(buffer, 40);
 	BufferIO::CopyWStr(buffer, data, 20);
 }
-void Replay::ReadData(void* data, unsigned int length) {
+void Replay::ReadData(void* data, int length) {
 	if(!is_replaying)
 		return;
 	memcpy(data, pdata, length);

gframe/replay.h

@@ -6,11 +6,16 @@
 
 namespace ygo {
 
+// replay flag
 #define REPLAY_COMPRESSED	0x1
 #define REPLAY_TAG			0x2
 #define REPLAY_DECODED		0x4
 #define REPLAY_SINGLE_MODE	0x8
 
+// max size
+#define MAX_REPLAY_SIZE	0x20000
+#define MAX_COMP_SIZE	0x2000
+
 struct ReplayHeader {
 	unsigned int id;
 	unsigned int version;
@@ -19,21 +24,28 @@ struct ReplayHeader {
 	unsigned int datasize;
 	unsigned int hash;
 	unsigned char props[8];
+
+	ReplayHeader()
+		: id(0), version(0), flag(0), seed(0), datasize(0), hash(0), props{ 0 } {}
 };
 
 class Replay {
 public:
 	Replay();
 	~Replay();
+
+	// record
 	void BeginRecord();
 	void WriteHeader(ReplayHeader& header);
-	void WriteData(const void* data, unsigned int length, bool flush = true);
+	void WriteData(const void* data, int length, bool flush = true);
 	void WriteInt32(int data, bool flush = true);
 	void WriteInt16(short data, bool flush = true);
 	void WriteInt8(char data, bool flush = true);
 	void Flush();
 	void EndRecord();
 	void SaveReplay(const wchar_t* name);
+
+	// play
 	bool OpenReplay(const wchar_t* name);
 	static bool CheckReplay(const wchar_t* name);
 	static bool DeleteReplay(const wchar_t* name);
@@ -41,22 +53,25 @@ public:
 	bool ReadNextResponse(unsigned char resp[64]);
 	void ReadName(wchar_t* data);
 	void ReadHeader(ReplayHeader& header);
-	void ReadData(void* data, unsigned int length);
+	void ReadData(void* data, int length);
 	int ReadInt32();
 	short ReadInt16();
 	char ReadInt8();
 	void Rewind();
 
 	FILE* fp;
-	ReplayHeader pheader;
 #ifdef _WIN32
 	HANDLE recording_fp;
 #endif
+
+	ReplayHeader pheader;
 	unsigned char* replay_data;
 	unsigned char* comp_data;
-	unsigned char* pdata;
 	size_t replay_size;
 	size_t comp_size;
+
+private:
+	unsigned char* pdata;
 	bool is_recording;
 	bool is_replaying;
 };