1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
|
#include "Common.h"
#include "IO.h"
#include <memory>
#ifdef WIN32
#include <Windows.h>
#include <winsock.h>
#else
#include <arpa/inet.h>
#include <netinet/in.h>
#include <sys/socket.h>
#include <sys/unistd.h>
#endif
#include <string>
[[noreturn]] void PrintErrorMessageAndExit(StringView message,
bool print_last_error) {
SendOutput(CRUT("{}\n"), message);
if (print_last_error) {
#ifdef WIN32
auto error_code = WSAGetLastError();
SendOutput(OutputType::Error, CRUT("Error code is {}.\n"), error_code);
wchar_t buffer[500];
if (!FormatMessageW(FORMAT_MESSAGE_FROM_SYSTEM |
FORMAT_MESSAGE_ARGUMENT_ARRAY |
FORMAT_MESSAGE_IGNORE_INSERTS,
nullptr, error_code, 0, buffer, 500, nullptr)) {
SendOutput(OutputType::Error, CRUT("Failed to format error message.\n"));
} else {
SendOutput(OutputType::Error, CRUT("{}\n"), buffer);
}
#else
#endif
}
BeforeExit();
std::exit(1);
}
#ifdef WIN32
namespace {
void InitWSA() {
WSADATA wsa_data;
if (WSAStartup(MAKEWORD(2, 2), &wsa_data)) { // initialize wsa
PrintErrorMessageAndExit(CRUT("Failed to initialize wsa."));
}
}
} // namespace
#endif
int CloseSocket(int socket) {
#ifdef WIN32
return closesocket(socket);
#else
return close(socket);
#endif
}
void BeforeExit() {
#ifdef WIN32
WSACleanup();
#endif
SignalAndWaitForOutputThreadStop();
}
bool SafeSend(int socket, std::string_view buffer) {
const int total_byte_count = buffer.size();
int byte_count_sent = 0;
int retry_count = 0;
while (true) {
// Now we have sent all data.
if (byte_count_sent == total_byte_count)
return true;
auto byte_actually_sent = send(socket, buffer.data() + byte_count_sent,
buffer.size() - byte_count_sent, 0);
// send failed
if (byte_actually_sent == -1) {
return false;
}
byte_count_sent += byte_actually_sent;
}
}
bool SafeReadUntil(int socket, char c, std::string &data, std::string &rest) {
data = rest;
const int buffer_size = 100;
char buffer[buffer_size];
while (true) {
int received_number = recv(socket, buffer, buffer_size, 0);
if (received_number == -1) {
return false;
}
bool end = false;
for (int i = 0; i < received_number; i++) {
if (buffer[i] == c) {
data.append(buffer, i);
rest = std::string(buffer + i + 1, received_number - i - 1);
end = true;
break;
}
}
if (end)
return true;
data.append(buffer, received_number);
}
}
int main() {
#ifdef WIN32
HANDLE h = GetStdHandle(STD_OUTPUT_HANDLE);
DWORD mode;
GetConsoleMode(h, &mode);
mode |= ENABLE_VIRTUAL_TERMINAL_PROCESSING;
SetConsoleMode(h, mode);
InitWSA();
#endif
int c = Main();
BeforeExit();
return c;
}
|
