1 files changed, 28 insertions, 23 deletions

diff --git a/src/connections.cc b/src/connections.cc
index 339e6dc..6f78ef2 100644
--- a/src/connections.cc
+++ b/src/connections.cc
@@ -221,43 +221,48 @@ static void CheckConnectionInput(TConnection *Connection, int Events){
 						0x55, 0x49, 0x54, 0x0A,
 					};
 
-					if(memcmp(Connection->Buffer, HAProxyV2Signature, 12) != 0){
-						LOG_ERR("Invalid proxy header from %s", Connection->RemoteAddress);
-						CloseConnection(Connection);
-						break;
+					if(memcmp(Connection->Buffer, HAProxyV2Signature, 12) == 0){
+						SourceAddr = Connection->Buffer + 16;
 					}
 
-					SourceAddr = Connection->Buffer + 16;
 				}else if(Connection->Buffer[0] == 0xD9 && Connection->Buffer[1] == 0xFF){
+					// NOTE(fusion): Our custom proxy.
 					ASSERT(ReadSize == 6);
 					SourceAddr = Connection->Buffer + 2;
 				}
 
-				if(SourceAddr != NULL){
-					if(Connection->IPAddress == INADDR_LOOPBACK){
-						LOG("Connection proxy header %s -> %d.%d.%d.%d:%d",
-								Connection->RemoteAddress,
-								(int)SourceAddr[0], (int)SourceAddr[1],
-								(int)SourceAddr[2], (int)SourceAddr[3],
-								Connection->RemotePort);
-
-						Connection->IPAddress = BufferRead32BE(SourceAddr);
-						StringBufFormat(Connection->RemoteAddress,
-								"%d.%d.%d.%d:%d",
-								(int)SourceAddr[0], (int)SourceAddr[1],
-								(int)SourceAddr[2], (int)SourceAddr[3],
-								Connection->RemotePort);
-					}else{
-						LOG_ERR("Received proxy header (%02X, %02X) from unknown host %s",
+				if(SourceAddr == NULL || Connection->IPAddress != INADDR_LOOPBACK){
+					ASSERT(ReadSize >= 2);
+
+					if(SourceAddr == NULL){
+						LOG_ERR("Invalid proxy header (%02X, %02X) from host %s",
 							(int)Connection->Buffer[0], (int)Connection->Buffer[1],
 							Connection->RemoteAddress);
 					}
-				}else{
-					LOG_ERR("Ignoring unknown proxy header (%02X, %02X) from %s",
+
+					if(Connection->IPAddress != INADDR_LOOPBACK){
+						LOG_ERR("Received proxy header (%02X, %02X) from unknown host %s",
 							(int)Connection->Buffer[0], (int)Connection->Buffer[1],
 							Connection->RemoteAddress);
+					}
+
+					CloseConnection(Connection);
+					break;
 				}
 
+				LOG("Connection proxy header %s -> %d.%d.%d.%d:%d",
+						Connection->RemoteAddress,
+						(int)SourceAddr[0], (int)SourceAddr[1],
+						(int)SourceAddr[2], (int)SourceAddr[3],
+						Connection->RemotePort);
+
+				Connection->IPAddress = BufferRead32BE(SourceAddr);
+				StringBufFormat(Connection->RemoteAddress,
+						"%d.%d.%d.%d:%d",
+						(int)SourceAddr[0], (int)SourceAddr[1],
+						(int)SourceAddr[2], (int)SourceAddr[3],
+						Connection->RemotePort);
+
 				Connection->RWSize = 0;
 				Connection->RWPosition = 0;
 				Connection->ReadingProxyHeader = false;