Skip to content

Commit ef50b5a

Browse files
yyubinyyubin
committed
fix(websocket): NIO 기반 핸드셰이크 및 프로토콜 감지 전면 수정
- Blocking I/O 제거 후 ByteBuffer 기반 NIO 읽기 방식으로 일원화 - 중복 ByteBuffer 읽기로 인한 -1 반환 문제 해결 - WebSocket Upgrade 요청 시 HTTP Detector가 처리하지 않도록 로직 수정 - Connection 헤더 "Upgrade" 포함 여부로 검증 (RFC 6455 준수) - 핸드셰이크 응답을 SocketChannel 기반으로 변경 - WebSocketBenchmarkHandler에 @component 추가해 정상 등록 - /ws/benchmark 요청이 HTTP로 잘못 처리되던 문제 해결 전체 WebSocket 요청 처리 흐름이 NIO 모드에서 안정적으로 동작하도록 리팩토링됨.
1 parent 7ffb03b commit ef50b5a

13 files changed

Lines changed: 676 additions & 112 deletions

File tree

‎README.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -172,7 +172,7 @@ public class ChatSocket {
172172
**Reports:** [Tests](https://yyubin.github.io/sprout/tests/) ·
173173
[Coverage](https://yyubin.github.io/sprout/coverage/)
174174

175-
**673 tests, 0 failures (100% pass, Gradle 8.10.1 · 2025‑10-27)**
175+
**687 tests, 0 failures (100% pass, Gradle 8.10.1 · 2025‑10-27)**
176176

177177
**Test Coverage (Jacoco):**
178178
- **Line Coverage: 85%**

‎README_ko.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -174,7 +174,7 @@ public class ChatSocket {
174174
**리포트:** [테스트](https://yyubin.github.io/sprout/tests/) ·
175175
[커버리지](https://yyubin.github.io/sprout/coverage/)
176176

177-
**673개 테스트, 0개 실패 (100% 통과, Gradle 8.10.1 · 2025‑10‑21)**
177+
**687개 테스트, 0개 실패 (100% 통과, Gradle 8.10.1 · 2025‑10‑21)**
178178

179179
**테스트 커버리지 (Jacoco):**
180180
- **라인 커버리지: 85%**

‎src/main/java/app/benchmark/WebSocketBenchmarkHandler.java‎

Lines changed: 173 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,173 @@
1+
package app.benchmark;
2+
3+
import sprout.beans.annotation.Component;
4+
import sprout.server.argument.annotation.Payload;
5+
import sprout.server.websocket.CloseCode;
6+
import sprout.server.websocket.WebSocketSession;
7+
import sprout.server.websocket.annotation.*;
8+
9+
import java.io.IOException;
10+
import java.util.Map;
11+
import java.util.concurrent.ConcurrentHashMap;
12+
13+
/**
14+
* WebSocket 벤치마크용 핸들러
15+
*
16+
* 지원 기능:
17+
* - /echo: 메시지 에코
18+
* - /broadcast: 모든 연결된 클라이언트에 브로드캐스트
19+
* - /chat: 채팅방 시뮬레이션
20+
* - /ping-pong: 간단한 핑퐁 응답
21+
*/
22+
@Component
23+
@WebSocketHandler("/ws/benchmark")
24+
public class WebSocketBenchmarkHandler {
25+
26+
// 연결된 모든 세션 관리
27+
private static final Map<String, WebSocketSession> sessions = new ConcurrentHashMap<>();
28+
29+
// 통계 정보
30+
private static long totalMessagesReceived = 0;
31+
private static long totalMessagesSent = 0;
32+
private static long totalConnections = 0;
33+
34+
@OnOpen
35+
public void onOpen(WebSocketSession session) {
36+
sessions.put(session.getId(), session);
37+
totalConnections++;
38+
System.out.println("[WebSocket Benchmark] 연결 열림: " + session.getId() +
39+
" (총 연결: " + sessions.size() + ", 누적: " + totalConnections + ")");
40+
}
41+
42+
@OnClose
43+
public void onClose(WebSocketSession session, CloseCode closeCode) {
44+
sessions.remove(session.getId());
45+
System.out.println("[WebSocket Benchmark] 연결 닫힘: " + session.getId() +
46+
" (코드: " + closeCode.getCode() + ", 남은 연결: " + sessions.size() + ")");
47+
}
48+
49+
@OnError
50+
public void onError(WebSocketSession session, Throwable error) {
51+
System.err.println("[WebSocket Benchmark] 에러 발생: " + session.getId() +
52+
" - " + error.getMessage());
53+
error.printStackTrace();
54+
}
55+
56+
/**
57+
* Echo: 받은 메시지를 그대로 반환
58+
*/
59+
@MessageMapping("/echo")
60+
public void handleEcho(WebSocketSession session, @Payload String message) throws IOException {
61+
totalMessagesReceived++;
62+
session.sendText(createResponse("/echo", message));
63+
totalMessagesSent++;
64+
}
65+
66+
/**
67+
* Broadcast: 모든 연결된 클라이언트에 메시지 전송
68+
*/
69+
@MessageMapping("/broadcast")
70+
public void handleBroadcast(WebSocketSession session, @Payload String message) throws IOException {
71+
totalMessagesReceived++;
72+
String response = createResponse("/broadcast", "From " + session.getId() + ": " + message);
73+
74+
for (WebSocketSession s : sessions.values()) {
75+
if (s.isOpen()) {
76+
try {
77+
s.sendText(response);
78+
totalMessagesSent++;
79+
} catch (IOException e) {
80+
System.err.println("브로드캐스트 실패: " + s.getId() + " - " + e.getMessage());
81+
}
82+
}
83+
}
84+
}
85+
86+
/**
87+
* Chat: 채팅 메시지 처리
88+
*/
89+
@MessageMapping("/chat")
90+
public void handleChat(WebSocketSession session, @Payload String message) throws IOException {
91+
totalMessagesReceived++;
92+
String username = (String) session.getUserProperties().get("username");
93+
if (username == null) {
94+
username = "User-" + session.getId().substring(0, 8);
95+
session.getUserProperties().put("username", username);
96+
}
97+
98+
String chatMessage = username + ": " + message;
99+
String response = createResponse("/chat", chatMessage);
100+
101+
// 채팅방의 모든 사용자에게 전송
102+
for (WebSocketSession s : sessions.values()) {
103+
if (s.isOpen()) {
104+
try {
105+
s.sendText(response);
106+
totalMessagesSent++;
107+
} catch (IOException e) {
108+
System.err.println("채팅 메시지 전송 실패: " + s.getId());
109+
}
110+
}
111+
}
112+
}
113+
114+
/**
115+
* Ping-Pong: 간단한 응답
116+
*/
117+
@MessageMapping("/ping")
118+
public void handlePing(WebSocketSession session, String message) throws IOException {
119+
totalMessagesReceived++;
120+
session.sendText(createResponse("/pong", "pong"));
121+
totalMessagesSent++;
122+
}
123+
124+
/**
125+
* Stats: 통계 정보 반환
126+
*/
127+
@MessageMapping("/stats")
128+
public void handleStats(WebSocketSession session, String message) throws IOException {
129+
totalMessagesReceived++;
130+
String stats = String.format(
131+
"연결: %d, 수신: %d, 송신: %d, 누적 연결: %d",
132+
sessions.size(), totalMessagesReceived, totalMessagesSent, totalConnections
133+
);
134+
session.sendText(createResponse("/stats", stats));
135+
totalMessagesSent++;
136+
}
137+
138+
/**
139+
* JSON 응답 메시지 생성
140+
*/
141+
private String createResponse(String destination, String payload) {
142+
// JSON 형식: {"destination": "...", "payload": "..."}
143+
return String.format("{\"destination\":\"%s\",\"payload\":\"%s\"}",
144+
destination, escapeJson(payload));
145+
}
146+
147+
/**
148+
* JSON 문자열 이스케이프
149+
*/
150+
private String escapeJson(String str) {
151+
if (str == null) return "";
152+
return str.replace("\\", "\\\\")
153+
.replace("\"", "\\\"")
154+
.replace("\n", "\\n")
155+
.replace("\r", "\\r")
156+
.replace("\t", "\\t");
157+
}
158+
159+
// 통계 초기화 (테스트용)
160+
public static void resetStats() {
161+
totalMessagesReceived = 0;
162+
totalMessagesSent = 0;
163+
totalConnections = 0;
164+
}
165+
166+
// 통계 조회 (테스트용)
167+
public static String getStats() {
168+
return String.format(
169+
"Sessions: %d, Received: %d, Sent: %d, Total Connections: %d",
170+
sessions.size(), totalMessagesReceived, totalMessagesSent, totalConnections
171+
);
172+
}
173+
}

‎src/main/java/sprout/server/builtins/HttpProtocolDetector.java‎

Lines changed: 16 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -21,19 +21,28 @@ public String detect(ByteBuffer buffer) throws Exception {
2121
// 버퍼의 현재 위치를 기록
2222
buffer.mark();
2323

24-
int readLimit = Math.min(buffer.remaining(), HTTP_HEADER_LENGTH);
25-
byte[] headerBytes = new byte[readLimit];
26-
buffer.get(headerBytes);
24+
// 전체 헤더 읽기 (WebSocket 감지를 위해)
25+
byte[] fullHeaderBytes = new byte[buffer.remaining()];
26+
buffer.get(fullHeaderBytes);
2727

2828
// 버퍼의 위치를 원래대로 되돌림
2929
buffer.reset();
3030

31-
String prefix = new String(headerBytes, StandardCharsets.UTF_8);
31+
String fullHeader = new String(fullHeaderBytes, StandardCharsets.UTF_8);
3232

33-
if (HTTP_METHODS.stream().anyMatch(prefix::startsWith)) {
34-
return "HTTP/1.1";
33+
// HTTP 메서드 체크
34+
if (!HTTP_METHODS.stream().anyMatch(fullHeader::startsWith)) {
35+
return "UNKNOWN";
3536
}
3637

37-
return "UNKNOWN";
38+
System.out.println("full header is " + fullHeader);
39+
40+
// WebSocket Upgrade 요청은 UNKNOWN 반환 (WebSocketProtocolDetector가 처리하도록)
41+
if (fullHeader.contains("Upgrade: websocket") ||
42+
fullHeader.contains("Upgrade: WebSocket")) {
43+
return "UNKNOWN";
44+
}
45+
46+
return "HTTP/1.1";
3847
}
3948
}

‎src/main/java/sprout/server/builtins/WebSocketProtocolHandler.java‎

Lines changed: 71 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,6 @@
1616
import sprout.server.websocket.message.WebSocketMessageParser;
1717

1818
import java.io.*;
19-
import java.lang.reflect.Method;
20-
import java.net.Socket;
2119
import java.nio.ByteBuffer;
2220
import java.nio.channels.SelectionKey;
2321
import java.nio.channels.Selector;
@@ -72,32 +70,29 @@ public boolean supports(String protocol) {
7270

7371
@Override
7472
public void accept(SocketChannel channel, Selector selector, ByteBuffer byteBuffer) throws Exception {
75-
Socket socket = channel.socket();
76-
77-
BufferedReader httpReader = new BufferedReader(new InputStreamReader(socket.getInputStream())); // 초기 HTTP 파싱용
78-
BufferedWriter httpWriter = new BufferedWriter(new OutputStreamWriter(socket.getOutputStream()));
79-
8073
// 1. 초기 HTTP 요청 파싱 (웹소켓 핸드셰이크 요청)
81-
String rawHttpRequest = readRawHttpRequestContent(httpReader);
74+
// ByteBuffer를 사용하여 NIO non-blocking 방식으로 읽기
75+
String rawHttpRequest = readRawHttpRequestContent(channel, byteBuffer);
8276
HttpRequest<?> request = httpRequestParser.parse(rawHttpRequest);
8377
if (!request.isValid()) {
8478
System.out.println("Empty or invalid HTTP request for websocket handshake. Closing socket.");
85-
socket.close();
79+
channel.close();
8680
return;
8781
}
8882

8983
// 2. 웹소켓 엔드포인트 찾기
9084
String requestPath = request.getPath();
85+
System.out.println("WebSocket handshake request received for path: " + requestPath + ". Trying to find matching endpoint.");
9186
WebSocketEndpointInfo endpointInfo = endpointRegistry.getEndpointInfo(requestPath);
9287

9388
if (endpointInfo == null) {
94-
sendHttpResponse(httpWriter, 404, "Not Found", "No WebSocket endpoint found for " + requestPath);
95-
socket.close();
89+
sendHttpResponse(channel, 404, "Not Found", "No WebSocket endpoint found for " + requestPath);
90+
channel.close();
9691
return;
9792
}
9893

9994
// 3. 핸드셰이크 수행
100-
boolean handshakeSuccess = handshakeHandler.performHandshake(request, httpWriter); // httpWriter 사용
95+
boolean handshakeSuccess = handshakeHandler.performHandshake(request, channel);
10196
if (!handshakeSuccess) {
10297
System.out.println("WebSocket handshake failed. Closing socket.");
10398
channel.close();
@@ -117,56 +112,76 @@ public void accept(SocketChannel channel, Selector selector, ByteBuffer byteBuff
117112
wsSession.callOnOpenMethod();
118113
}
119114

120-
private String readRawHttpRequestContent(BufferedReader in) throws IOException {
115+
private String readRawHttpRequestContent(SocketChannel channel, ByteBuffer buffer) throws IOException {
121116
StringBuilder sb = new StringBuilder();
122-
String line;
123-
int contentLength = 0;
124-
// 헤더 끝을 나타내는 플래그
125-
boolean headersDone = false;
126-
127-
// HTTP 요청 라인 + 헤더 읽기
128-
// readLine()이 null을 반환하면 클라이언트가 연결을 끊은 것
129-
// line.isEmpty()는 헤더 끝의 빈 줄을 의미
130-
while ((line = in.readLine()) != null) {
131-
if (line.isEmpty()) { // 빈 줄은 헤더의 끝을 의미 (CRLFCRLF 또는 LF LF)
132-
headersDone = true;
133-
break;
134-
}
135-
sb.append(line).append("\r\n"); // HTTP 규격에 맞게 CRLF 추가
136-
if (line.toLowerCase().startsWith("content-length:")) {
137-
try {
138-
contentLength = Integer.parseInt(line.substring(line.indexOf(':') + 1).trim());
139-
} catch (NumberFormatException e) {
140-
System.err.println("Warning: Invalid Content-Length header: " + line);
141-
contentLength = 0; // 파싱 실패 시 0으로 설정
117+
118+
// 1) 이미 읽은 buffer의 데이터를 먼저 추가
119+
if (buffer != null && buffer.hasRemaining()) {
120+
byte[] arr = new byte[buffer.remaining()];
121+
buffer.get(arr);
122+
sb.append(new String(arr, StandardCharsets.UTF_8));
123+
}
124+
125+
// 2) 이미 완전한 HTTP 요청인지 확인
126+
String current = sb.toString();
127+
if (current.contains("\r\n\r\n")) {
128+
return current;
129+
}
130+
131+
// 3) 불완전한 경우, 추가로 읽기 (blocking 모드로 전환)
132+
boolean wasBlocking = channel.isBlocking();
133+
try {
134+
channel.configureBlocking(true);
135+
136+
ByteBuffer readBuffer = ByteBuffer.allocate(8192);
137+
138+
// HTTP 헤더 끝(\r\n\r\n)까지 읽기
139+
while (!sb.toString().contains("\r\n\r\n")) {
140+
readBuffer.clear();
141+
int bytesRead = channel.read(readBuffer);
142+
143+
if (bytesRead == -1) {
144+
return ""; // 연결 종료
145+
}
146+
147+
if (bytesRead == 0) {
148+
break;
149+
}
150+
151+
readBuffer.flip();
152+
byte[] bytes = new byte[readBuffer.remaining()];
153+
readBuffer.get(bytes);
154+
sb.append(new String(bytes, StandardCharsets.UTF_8));
155+
156+
// 너무 큰 요청은 거부 (10KB 제한)
157+
if (sb.length() > 10240) {
158+
throw new IOException("HTTP request too large");
142159
}
143160
}
144-
}
145-
sb.append("\r\n"); // 헤더와 바디 구분자 (readLine()이 빈 줄을 이미 제거했을 수도 있지만, 안전을 위해 추가)
146-
147-
// HTTP 바디 읽기 (Content-Length가 있고, 헤더가 끝났을 경우에만)
148-
if (contentLength > 0 && headersDone) {
149-
char[] body = new char[contentLength];
150-
int totalRead = 0;
151-
int read;
152-
// Content-Length만큼 정확히 읽으려고 시도
153-
// read()는 모든 바이트를 한 번에 읽지 않을 수 있으므로 루프 필요
154-
while (totalRead < contentLength && (read = in.read(body, totalRead, contentLength - totalRead)) != -1) {
155-
totalRead += read;
161+
162+
return sb.toString();
163+
164+
} finally {
165+
// 원래 blocking 모드로 복원
166+
if (!wasBlocking) {
167+
channel.configureBlocking(false);
156168
}
157-
sb.append(body, 0, totalRead); // 읽은 만큼만 추가
158169
}
159-
return sb.toString();
160170
}
161171

162-
163-
// HTTP 응답을 보내는 헬퍼 메서드 (핸드셰이크 실패 또는 엔드포인트 없을 때)
164-
private void sendHttpResponse(BufferedWriter out, int statusCode, String statusText, String message) throws IOException {
165-
out.write("HTTP/1.1 " + statusCode + " " + statusText + "\r\n");
166-
out.write("Content-Type: text/plain;charset=UTF-8\r\n");
167-
out.write("Content-Length: " + message.getBytes(StandardCharsets.UTF_8).length + "\r\n");
168-
out.write("\r\n");
169-
out.write(message);
170-
out.flush();
172+
/**
173+
* NIO 방식으로 HTTP 응답 전송
174+
*/
175+
private void sendHttpResponse(SocketChannel channel, int statusCode, String statusText, String message) throws IOException {
176+
String response = "HTTP/1.1 " + statusCode + " " + statusText + "\r\n" +
177+
"Content-Type: text/plain;charset=UTF-8\r\n" +
178+
"Content-Length: " + message.getBytes(StandardCharsets.UTF_8).length + "\r\n" +
179+
"\r\n" +
180+
message;
181+
182+
ByteBuffer buffer = ByteBuffer.wrap(response.getBytes(StandardCharsets.UTF_8));
183+
while (buffer.hasRemaining()) {
184+
channel.write(buffer);
185+
}
171186
}
172187
}

0 commit comments

Comments
 (0)