1 files changed, 285 insertions, 0 deletions
diff --git a/src/query.cc b/src/query.cc
new file mode 100644
index 0000000..1335631
--- /dev/null
+++ b/src/query.cc
@@ -0,0 +1,285 @@
+#include "login.hh"
+
+// TODO(fusion): Support windows eventually?
+#if OS_LINUX
+#	include <errno.h>
+#	include <netdb.h>
+#	include <sys/socket.h>
+#	include <unistd.h>
+#else
+#	error "Operating system not currently supported."
+#endif
+
+static TQueryManagerConnection *g_QueryManagerConnection;
+
+bool ResolveHostName(const char *HostName, in_addr_t *OutAddr){
+	ASSERT(HostName != NULL && OutAddr != NULL);
+	addrinfo *Result = NULL;
+	addrinfo Hints = {};
+	Hints.ai_family = AF_INET;
+	Hints.ai_socktype = SOCK_STREAM;
+	int ErrCode = getaddrinfo(HostName, NULL, &Hints, &Result);
+	if(ErrCode != 0){
+		LOG_ERR("Failed to resolve hostname \"%s\": %s", HostName, gai_strerror(ErrCode));
+		return false;
+	}
+
+	bool Resolved = false;
+	for(addrinfo *AddrInfo = Result;
+			AddrInfo != NULL;
+			AddrInfo = AddrInfo->ai_next){
+		if(AddrInfo->ai_family == AF_INET && AddrInfo->ai_socktype == SOCK_STREAM){
+			ASSERT(AddrInfo->ai_addrlen == sizeof(sockaddr_in));
+			*OutAddr = ((sockaddr_in*)AddrInfo->ai_addr)->sin_addr.s_addr;
+			Resolved = true;
+			break;
+		}
+	}
+	freeaddrinfo(Result);
+	return Resolved;
+}
+
+bool Connect(TQueryManagerConnection *Connection){
+	if(Connection->Socket != -1){
+		LOG_ERR("Already connected");
+		return false;
+	}
+
+	in_addr_t Addr;
+	if(!ResolveHostName(g_QueryManagerHost, &Addr)){
+		LOG_ERR("Failed to resolve query manager's host name \"%s\"", g_QueryManagerHost);
+		return false;
+	}
+
+	Connection->Socket = socket(AF_INET, SOCK_STREAM, 0);
+	if(Connection->Socket == -1){
+		LOG_ERR("Failed to create socket: (%d) %s", errno, strerrordesc_np(errno));
+		return false;
+	}
+
+	sockaddr_in QueryManagerAddress = {};
+	QueryManagerAddress.sin_family = AF_INET;
+	QueryManagerAddress.sin_port = htons((uint16)g_QueryManagerPort);
+	QueryManagerAddress.sin_addr.s_addr = Addr;
+	if(connect(Connection->Socket, (sockaddr*)&QueryManagerAddress, sizeof(QueryManagerAddress)) == -1){
+		LOG_ERR("Failed to connect: (%d) %s", errno, strerrordesc_np(errno));
+		Disconnect(Connection);
+		return false;
+	}
+
+	TWriteBuffer WriteBuffer = PrepareQuery(Connection, 0);
+	WriteBuffer.Write8((uint8)APPLICATION_TYPE_LOGIN);
+	WriteBuffer.WriteString(g_QueryManagerPassword);
+	int Status = ExecuteQuery(Connection, false, &WriteBuffer, NULL);
+	if(Status != QUERY_STATUS_OK){
+		LOG_ERR("Failed to login to query manager (%d)", Status);
+		Disconnect(Connection);
+		return false;
+	}
+
+	return true;
+}
+
+void Disconnect(TQueryManagerConnection *Connection){
+	if(Connection->Socket != -1){
+		close(Connection->Socket);
+		Connection->Socket = -1;
+	}
+}
+
+bool IsConnected(TQueryManagerConnection *Connection){
+	return Connection->Socket != -1;
+}
+
+TWriteBuffer PrepareQuery(TQueryManagerConnection *Connection, int QueryType){
+	TWriteBuffer WriteBuffer(Connection->Buffer, sizeof(Connection->Buffer));
+	WriteBuffer.Write16(0); // Request Size
+	WriteBuffer.Write8((uint8)QueryType);
+	return WriteBuffer;
+}
+
+static bool WriteExact(int Fd, const uint8 *Buffer, int Size){
+	int BytesToWrite = Size;
+	const uint8 *WritePtr = Buffer;
+	while(BytesToWrite > 0){
+		int Ret = (int)write(Fd, WritePtr, BytesToWrite);
+		if(Ret == -1){
+			return false;
+		}
+		BytesToWrite -= Ret;
+		WritePtr += Ret;
+	}
+	return true;
+}
+
+static bool ReadExact(int Fd, uint8 *Buffer, int Size){
+	int BytesToRead = Size;
+	uint8 *ReadPtr = Buffer;
+	while(BytesToRead > 0){
+		int Ret = (int)read(Fd, ReadPtr, BytesToRead);
+		if(Ret == -1 || Ret == 0){
+			return false;
+		}
+		BytesToRead -= Ret;
+		ReadPtr += Ret;
+	}
+	return true;
+}
+
+int ExecuteQuery(TQueryManagerConnection *Connection, bool AutoReconnect,
+		TWriteBuffer *WriteBuffer, TReadBuffer *OutReadBuffer){
+	ASSERT(WriteBuffer != NULL
+		&& WriteBuffer->Buffer == Connection->Buffer
+		&& WriteBuffer->Size == sizeof(Connection->Buffer)
+		&& WriteBuffer->Position > 2);
+
+	int RequestSize = WriteBuffer->Position - 2;
+	if(RequestSize < 0xFFFF){
+		WriteBuffer->Rewrite16(0, (uint16)RequestSize);
+	}else{
+		WriteBuffer->Rewrite16(0, 0xFFFF);
+		WriteBuffer->Insert32(2, (uint32)RequestSize);
+	}
+
+	if(WriteBuffer->Overflowed()){
+		LOG_ERR("Write buffer overflowed when writing request");
+		return QUERY_STATUS_FAILED;
+	}
+
+	const int MaxAttempts = 2;
+	for(int Attempts = 0; true; Attempts += 1){
+		int WriteSize = WriteBuffer->Position;
+		if(!IsConnected(Connection)){
+			if(!AutoReconnect){
+				return QUERY_STATUS_FAILED;
+			}
+
+			// IMPORTANT(fusion): There is no way around this. `Connect` will
+			// use the connection buffer to send the login query so we need to
+			// save and restore it to not lose any data. One improvement here
+			// would be to use a stack or statically allocated buffer, although
+			// using malloc/free should have no real impact on performance as
+			// we're already doing BLOCKING I/O here.
+			uint8 *TempBuffer = (uint8*)malloc(WriteSize);
+			memcpy(TempBuffer, Connection->Buffer, WriteSize);
+			bool Reconnected = Connect(Connection);
+			memcpy(Connection->Buffer, TempBuffer, WriteSize);
+			free(TempBuffer);
+
+			if(!Reconnected){
+				return QUERY_STATUS_FAILED;
+			}
+		}
+
+		if(!WriteExact(Connection->Socket, Connection->Buffer, WriteSize)){
+			Disconnect(Connection);
+			if(Attempts >= MaxAttempts){
+				LOG_ERR("Failed to write request");
+				return QUERY_STATUS_FAILED;
+			}
+			continue;
+		}
+
+		uint8 Help[4];
+		if(!ReadExact(Connection->Socket, Help, 2)){
+			Disconnect(Connection);
+			if(Attempts >= MaxAttempts){
+				LOG_ERR("Failed to read response size");
+				return QUERY_STATUS_FAILED;
+			}
+			continue;
+		}
+
+		int ResponseSize = BufferRead16LE(Help);
+		if(ResponseSize == 0xFFFF){
+			if(!ReadExact(Connection->Socket, Help, 4)){
+				Disconnect(Connection);
+				LOG_ERR("Failed to read response extended size");
+				return QUERY_STATUS_FAILED;
+			}
+			ResponseSize = BufferRead32LE(Help);
+		}
+
+		if(ResponseSize <= 0 || ResponseSize > (int)sizeof(Connection->Buffer)){
+			Disconnect(Connection);
+			LOG_ERR("Invalid response size %d (BufferSize: %d)",
+					ResponseSize, (int)sizeof(Connection->Buffer));
+			return QUERY_STATUS_FAILED;
+		}
+
+		if(!ReadExact(Connection->Socket, Connection->Buffer, ResponseSize)){
+			Disconnect(Connection);
+			LOG_ERR("Failed to read response");
+			return QUERY_STATUS_FAILED;
+		}
+
+		TReadBuffer ReadBuffer(Connection->Buffer, ResponseSize);
+		int Status = ReadBuffer.Read8();
+		if(OutReadBuffer){
+			*OutReadBuffer = ReadBuffer;
+		}
+		return Status;
+	}
+}
+
+int LoginAccount(int AccountID, const char *Password, const char *IPAddress,
+		int MaxCharacters, int *NumCharacters, TCharacterLoginData *Characters,
+		int *PremiumDays){
+	TWriteBuffer WriteBuffer = PrepareQuery(g_QueryManagerConnection, 11);
+	WriteBuffer.Write32((uint32)AccountID);
+	WriteBuffer.WriteString(Password);
+	WriteBuffer.WriteString(IPAddress);
+
+	TReadBuffer ReadBuffer;
+	int Status = ExecuteQuery(g_QueryManagerConnection, true, &WriteBuffer, &ReadBuffer);
+	int Result = (Status == QUERY_STATUS_OK ? 0 : -1);
+	if(Status == QUERY_STATUS_OK){
+		*NumCharacters = ReadBuffer.Read8();
+		if(*NumCharacters > MaxCharacters){
+			LOG_ERR("Too many characters");
+			return -1;
+		}
+
+		for(int i = 0; i < *NumCharacters; i += 1){
+			ReadBuffer.ReadString(Characters[i].Name, sizeof(Characters[i].Name));
+			ReadBuffer.ReadString(Characters[i].WorldName, sizeof(Characters[i].WorldName));
+			Characters[i].WorldAddress = ReadBuffer.Read32BE();
+			Characters[i].WorldPort = ReadBuffer.Read16();
+		}
+
+		*PremiumDays = ReadBuffer.Read16();
+	}else if(Status == QUERY_STATUS_ERROR){
+		int ErrorCode = ReadBuffer.Read8();
+		if(ErrorCode >= 1 && ErrorCode <= 6){
+			Result = ErrorCode;
+		}else{
+			LOG_ERR("Invalid error code %d", ErrorCode);
+		}
+	}else{
+		LOG_ERR("Request failed");
+	}
+	return Result;
+}
+
+bool InitQuery(void){
+	ASSERT(g_QueryManagerConnection == NULL);
+
+	LOG("QueryManagerHost: %s", g_QueryManagerHost);
+	LOG("QueryManagerPort: %d", g_QueryManagerPort);
+
+	g_QueryManagerConnection = (TQueryManagerConnection*)calloc(1, sizeof(TQueryManagerConnection));
+	g_QueryManagerConnection->Socket = -1;
+	if(!Connect(g_QueryManagerConnection)){
+		LOG_ERR("Failed to connect to query manager");
+		return false;
+	}
+	return true;
+}
+
+void ExitQuery(void){
+	if(g_QueryManagerConnection != NULL){
+		Disconnect(g_QueryManagerConnection);
+		free(g_QueryManagerConnection);
+		g_QueryManagerConnection = NULL;
+	}
+}