1 files changed, 83 insertions, 13 deletions

diff --git a/src/connections.cc b/src/connections.cc
index a1111df..339e6dc 100644
--- a/src/connections.cc
+++ b/src/connections.cc
@@ -129,6 +129,7 @@ static TConnection *AssignConnection(int Socket, uint32 Addr, uint16 Port){
 		Connection->State = CONNECTION_READING;
 		Connection->Socket = Socket;
 		Connection->IPAddress = (int)Addr;
+		Connection->RemotePort = (int)Port;
 		Connection->StartTime = GetMonotonicUptime();
 		Connection->RandomSeed = (uint32)rand();
 		StringBufFormat(Connection->RemoteAddress,
@@ -137,7 +138,7 @@ static TConnection *AssignConnection(int Socket, uint32 Addr, uint16 Port){
 				((Connection->IPAddress >> 16) & 0xFF),
 				((Connection->IPAddress >>  8) & 0xFF),
 				((Connection->IPAddress >>  0) & 0xFF),
-				(int)Port);
+				Connection->RemotePort);
 		LOG("Connection %s assigned to slot %d",
 				Connection->RemoteAddress, ConnectionIndex);
 	}
@@ -184,21 +185,87 @@ static void CheckConnectionInput(TConnection *Connection, int Events){
 
 		Connection->RWPosition += BytesRead;
 		if(Connection->RWPosition >= ReadSize){
-			if(Connection->RWSize != 0){
-				Connection->State = CONNECTION_PROCESSING;
-				break;
-			}else if(Connection->RWPosition == 2){
-				int PayloadSize = (int)BufferRead16LE(Connection->Buffer);
-				if(PayloadSize <= 0 || PayloadSize > (int)sizeof(Connection->Buffer)){
-					CloseConnection(Connection);
-					break;
+			ASSERT(Connection->RWPosition == ReadSize);
+			if(Connection->RWSize == 0){
+#if ALLOW_LOCAL_PROXY
+				if(Connection->Buffer[0] == 0x0D && Connection->Buffer[1] == 0x0A){
+					ASSERT(sizeof(Connection->Buffer) >= 28);
+					Connection->RWSize = 28;
+					Connection->ReadingProxyHeader = true;
+				}else if(Connection->Buffer[0] == 0xD9 && Connection->Buffer[1] == 0xFF){
+					ASSERT(sizeof(Connection->Buffer) >= 6);
+					Connection->RWSize = 6;
+					Connection->ReadingProxyHeader = true;
+				}else
+#endif
+				{
+					int PayloadSize = (int)BufferRead16LE(Connection->Buffer);
+					if(PayloadSize <= 0 || PayloadSize > (int)sizeof(Connection->Buffer)){
+						CloseConnection(Connection);
+						break;
+					}
+
+					Connection->RWSize = PayloadSize;
+					Connection->RWPosition = 0;
+				}
+			}
+#if ALLOW_LOCAL_PROXY
+			else if(Connection->ReadingProxyHeader){
+				uint8 *SourceAddr = NULL;
+				if(Connection->Buffer[0] == 0x0D && Connection->Buffer[1] == 0x0A){
+					ASSERT(ReadSize == 28);
+
+					// NOTE(fusion): HAProxy-V2 header.
+					static const uint8_t HAProxyV2Signature[] = {
+						0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51,
+						0x55, 0x49, 0x54, 0x0A,
+					};
+
+					if(memcmp(Connection->Buffer, HAProxyV2Signature, 12) != 0){
+						LOG_ERR("Invalid proxy header from %s", Connection->RemoteAddress);
+						CloseConnection(Connection);
+						break;
+					}
+
+					SourceAddr = Connection->Buffer + 16;
+				}else if(Connection->Buffer[0] == 0xD9 && Connection->Buffer[1] == 0xFF){
+					ASSERT(ReadSize == 6);
+					SourceAddr = Connection->Buffer + 2;
 				}
 
-				Connection->RWSize = PayloadSize;
+				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",
+							(int)Connection->Buffer[0], (int)Connection->Buffer[1],
+							Connection->RemoteAddress);
+					}
+				}else{
+					LOG_ERR("Ignoring unknown proxy header (%02X, %02X) from %s",
+							(int)Connection->Buffer[0], (int)Connection->Buffer[1],
+							Connection->RemoteAddress);
+				}
+
+				Connection->RWSize = 0;
 				Connection->RWPosition = 0;
-			}else{
-				PANIC("Invalid input state (State: %d, RWSize: %d, RWPosition: %d)",
-						Connection->State, Connection->RWSize, Connection->RWPosition);
+				Connection->ReadingProxyHeader = false;
+			}
+#endif
+			else{
+				Connection->State = CONNECTION_PROCESSING;
+				break;
 			}
 		}
 	}
@@ -254,8 +321,11 @@ static void CheckConnectionOutput(TConnection *Connection, int Events){
 			break;
 		}
 
+		// NOTE(fusion): Both login and status protocols will serve a single
+		// request so we always close the connection after it is served.
 		Connection->RWPosition += BytesWritten;
 		if(Connection->RWPosition >= Connection->RWSize){
+			ASSERT(Connection->RWPosition == Connection->RWSize);
 			CloseConnection(Connection);
 			break;
 		}