Make copyRect function safer and add DX8Wrapper::Bytes_Per_Pixel function

Author: xezon

Core/GameEngineDevice/Source/W3DDevice/GameClient/W3DSmudge.cpp

@@ -134,62 +134,88 @@ void W3DSmudgeManager::ReAcquireResources(void)
 }
 
 /*Copies a portion of the current render target into a specified buffer*/
-Int copyRect(unsigned char *buf, Int bufSize, int oX, int oY, int width, int height)
+Int copyRect(unsigned char *buf, Int bufSize, Int oX, Int oY, Int width, Int height)
 {
- 	IDirect3DSurface8 *surface=NULL;	///<previous render target
- 	IDirect3DSurface8 *tempSurface=NULL;
+	IDirect3DSurface8 *surface=NULL;	///<previous render target
+	IDirect3DSurface8 *tempSurface=NULL;
 	Int result = 0;
-	HRESULT hr = S_OK;
 
- 	LPDIRECT3DDEVICE8 m_pDev=DX8Wrapper::_Get_D3D_Device8();
+	LPDIRECT3DDEVICE8 m_pDev=DX8Wrapper::_Get_D3D_Device8();
 
 	if (!m_pDev)
 		goto error;
 
- 	m_pDev->GetRenderTarget(&surface);
+	m_pDev->GetRenderTarget(&surface);
 
 	if (!surface)
 		goto error;
 
- 	D3DSURFACE_DESC desc;
-
- 	surface->GetDesc(&desc);
+	D3DSURFACE_DESC desc;
+	surface->GetDesc(&desc);
 
 	RECT srcRect;
 	srcRect.left=oX;
 	srcRect.top=oY;
 	srcRect.right=oX+width;
 	srcRect.bottom=oY+height;
+	DEBUG_ASSERTCRASH(srcRect.left >= 0, ("copyRect - Invalid size"));
+	DEBUG_ASSERTCRASH(srcRect.top >= 0, ("copyRect - Invalid size"));
+	DEBUG_ASSERTCRASH(srcRect.right < (LONG)desc.Width, ("copyRect - Invalid size"));
+	DEBUG_ASSERTCRASH(srcRect.bottom < (LONG)desc.Height, ("copyRect - Invalid size"));
 
 	POINT dstPoint;
 	dstPoint.x=0;
 	dstPoint.y=0;
 
- 	hr=m_pDev->CreateImageSurface(  width, height, desc.Format, &tempSurface);
+	HRESULT hr;
+	hr=m_pDev->CreateImageSurface(width, height, desc.Format, &tempSurface);
 
 	if (hr != S_OK)
 		goto error;
 
- 	hr=m_pDev->CopyRects(surface,&srcRect,1,tempSurface,&dstPoint);
+	hr=m_pDev->CopyRects(surface,&srcRect,1,tempSurface,&dstPoint);
 
 	if (hr != S_OK)
 		goto error;
 
- 	D3DLOCKED_RECT lrect;
+	D3DLOCKED_RECT lrect;
 
- 	hr=tempSurface->LockRect(&lrect,NULL,D3DLOCK_READONLY);
+	hr=tempSurface->LockRect(&lrect,NULL,D3DLOCK_READONLY);
 
 	if (hr != S_OK)
 		goto error;
 
- 	tempSurface->GetDesc(&desc);
+	tempSurface->GetDesc(&desc);
+	UnsignedInt bytesPerPixel;
+	bytesPerPixel = DX8Wrapper::Bytes_Per_Pixel(desc.Format);
+
+	if (bytesPerPixel == 0)
+		goto error;
+
+	Int rowSize;
+	Int totalSize;
+	rowSize = width * bytesPerPixel;
+	totalSize = height * rowSize;
 
-	if (desc.Size < bufSize)
-		bufSize = desc.Size;
+	if (totalSize > bufSize)
+		totalSize = bufSize;
 
-	memcpy(buf,lrect.pBits,bufSize);
-	result = bufSize;
+	UnsignedByte* dst;
+	UnsignedByte* src;
+	Int rowCount;
+	Int row;
+	dst = buf;
+	src = (UnsignedByte*)lrect.pBits;
+	rowCount = totalSize / rowSize;
+
+	for (row = 0; row < rowCount; ++row)
+	{
+		memcpy(dst, src, rowSize);
+		dst += rowSize;
+		src += lrect.Pitch;
+	}
 
+	result = totalSize;
 	tempSurface->UnlockRect();
 
 error:

GeneralsMD/Code/Libraries/Source/WWVegas/WW3D2/dx8wrapper.h

@@ -455,6 +455,8 @@ class DX8Wrapper
 	static void _Enable_Triangle_Draw(bool enable) { _EnableTriangleDraw=enable; }
 	static bool _Is_Triangle_Draw_Enabled() { return _EnableTriangleDraw; }
 
+	static unsigned int Bytes_Per_Pixel(D3DFORMAT format);
+
 	/*
 	** Additional swap chain interface
 	**
@@ -1161,6 +1163,85 @@ WWINLINE void DX8Wrapper::Set_Alpha (const float alpha, unsigned int &color)
 	component [3] = 255.0f * alpha;
 }
 
+WWINLINE unsigned int DX8Wrapper::Bytes_Per_Pixel(D3DFORMAT format)
+{
+	switch (format)
+	{
+	// 8-bit formats
+	case D3DFMT_A8:
+	case D3DFMT_P8:
+	case D3DFMT_L8:
+		return 1;
+
+	// 16-bit formats
+	case D3DFMT_R5G6B5:
+	case D3DFMT_X1R5G5B5:
+	case D3DFMT_A1R5G5B5:
+	case D3DFMT_A4R4G4B4:
+	case D3DFMT_X4R4G4B4:
+	case D3DFMT_A8L8:
+	case D3DFMT_A4L4:
+	case D3DFMT_V8U8:
+	case D3DFMT_L6V5U5:
+		return 2;
+
+	// 24-bit formats
+	case D3DFMT_R8G8B8:
+		return 3;
+
+	// 32-bit formats
+	case D3DFMT_A8R8G8B8:
+	case D3DFMT_X8R8G8B8:
+	case D3DFMT_A8R3G3B2:
+	case D3DFMT_A2B10G10R10:
+	case D3DFMT_G16R16:
+	case D3DFMT_X8L8V8U8:
+	case D3DFMT_Q8W8V8U8:
+	case D3DFMT_W11V11U10:
+	case D3DFMT_A2W10V10U10:
+	case D3DFMT_V16U16:
+		return 4;
+
+	// Palette-based alpha
+	case D3DFMT_A8P8:
+		return 0;
+
+	// Packed YUV
+	case D3DFMT_UYVY:
+	case D3DFMT_YUY2:
+		return 0;
+
+	// Block-compressed
+	case D3DFMT_DXT1:
+	case D3DFMT_DXT2:
+	case D3DFMT_DXT3:
+	case D3DFMT_DXT4:
+	case D3DFMT_DXT5:
+		return 0;
+
+	// Depth / stencil
+	case D3DFMT_D16_LOCKABLE:
+	case D3DFMT_D16:
+	case D3DFMT_D15S1:
+	case D3DFMT_D24S8:
+	case D3DFMT_D24X8:
+	case D3DFMT_D24X4S4:
+	case D3DFMT_D32:
+		return 0;
+
+	// Non-image data
+	case D3DFMT_VERTEXDATA:
+	case D3DFMT_INDEX16:
+	case D3DFMT_INDEX32:
+		return 0;
+
+	// Unknown / invalid
+	case D3DFMT_UNKNOWN:
+	default:
+		return 0;
+	}
+}
+
 WWINLINE void DX8Wrapper::Get_Render_State(RenderStateStruct& state)
 {
 	state=render_state;