some changes

Author: jamshid-elmurodov

cli/src/main/java/uz/Main.java

@@ -22,6 +22,7 @@
 import java.nio.file.Path;
 import java.time.LocalTime;
 import java.time.format.DateTimeFormatter;
+import java.util.Arrays;
 import java.util.Objects;
 import java.util.concurrent.*;
 
@@ -223,14 +224,22 @@ void requestSender(Request request, WebSocket webSocket) {
                     forwardInfo.getPath()
             ));
 
+            if (responseBodyString.length() > 10000) {
+                printError("Response body is too large. Truncating to 10000 characters.");
+                return;
+            }
+
             Response send = new Response(request.getId(), statusCode, responseBodyString);
-            ByteBuffer buffer = ByteBuffer.wrap(objectMapper.writeValueAsBytes(send));
 
-            webSocket.sendBinary(buffer, true);
+            webSocket.sendText(
+                    objectMapper.writeValueAsString(send), true
+            );
         } catch (ConnectException e) {
             printError("Connection refused. Make sure the local server is running on port " + port);
         } catch (IOException e) {
             printError("I/O error occurred while processing request." + e.getMessage());
+        } catch (Exception e){
+            printError("An error occurred while processing request." + e.getMessage());
         }
     }
 

server/src/main/java/uz/server/config/WebSocketConfig.java

@@ -26,8 +26,7 @@ public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
     @Bean
     public ServletServerContainerFactoryBean createWebSocketContainer() {
         ServletServerContainerFactoryBean container = new ServletServerContainerFactoryBean();
-        container.setMaxTextMessageBufferSize(100 * 1024);
-        container.setMaxBinaryMessageBufferSize(100 * 1024);
+        container.setMaxTextMessageBufferSize(100 * 3072);
         return container;
     }
 }

server/src/main/java/uz/server/controller/ForwardController.java

@@ -21,11 +21,12 @@
 public class ForwardController {
     private final Forwarder forwarder;
 
-    @RequestMapping(value = "/**", headers = {"Upgrade!=websocket", "x-subdomain=*"})
+    @RequestMapping(value = "/**", headers = {"Upgrade!=websocket"})
     public ResponseEntity<String> handleRequest(
             @RequestBody(required = false) String body,
             HttpServletRequest servletRequest) {
-        String subdomain = servletRequest.getHeader("x-subdomain");
+        String host = servletRequest.getHeader("Host");
+        String subdomain = host.split("\\.")[0];
 
         String requestUri = servletRequest.getRequestURI();
 

server/src/main/java/uz/server/repository/TunnelRepository.java

@@ -5,12 +5,11 @@
 import uz.server.domain.entity.User;
 
 import java.util.Optional;
-import java.util.UUID;
 
 public interface TunnelRepository extends JpaRepository<Tunnel, Long> {
-    Integer countByUserAndActiveTrue(User user);
+    Integer countByUser(User user);
 
-    Optional<Tunnel> findBySubdomain(String subdomain);
+    Optional<Tunnel> findBySubdomainIgnoreCase(String subdomain);
 
     Optional<Tunnel> findBySessionId(String id);
 }

server/src/main/java/uz/server/service/TunnelService.java

@@ -9,7 +9,6 @@
 import uz.server.repository.TunnelRepository;
 
 import java.security.SecureRandom;
-import java.util.Base64;
 import java.util.Optional;
 
 @Service
@@ -20,7 +19,7 @@ public class TunnelService {
 
     public String create(String sessionId, User user) {
         log.info("Creating tunnel: userId={}", user.getId());
-        Integer count = repo.countByUserAndActiveTrue(user);
+        Integer count = repo.countByUser(user);
 
         if (count == 3){
             log.error("User can't create more than 3 tunnels: userId={}", user.getId());
@@ -29,7 +28,7 @@ public String create(String sessionId, User user) {
 
         Tunnel save = repo.save(Tunnel.builder()
                 .sessionId(sessionId)
-                .subdomain(getUniqueString())
+                .subdomain(generateUniqueString())
                 .user(user)
                 .build());
 
@@ -38,10 +37,17 @@ public String create(String sessionId, User user) {
         return save.getSubdomain();
     }
 
-    private static String getUniqueString() {
-        byte[] randomBytes = new byte[8];
-        new SecureRandom().nextBytes(randomBytes);
-        return Base64.getUrlEncoder().withoutPadding().encodeToString(randomBytes);
+    private String generateUniqueString() {
+        long timestamp = System.currentTimeMillis() % 100000;
+        StringBuilder sb = new StringBuilder();
+        SecureRandom RANDOM = new SecureRandom();
+        String CHARACTERS = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789";
+
+        sb.append(timestamp);
+        for (int i = 0; i < 8; i++) {
+            sb.append(CHARACTERS.charAt(RANDOM.nextInt(CHARACTERS.length())));
+        }
+        return sb.toString();
     }
 
     public void deactivate(Long tunnelId) {
@@ -60,7 +66,7 @@ public void deactivate(Long tunnelId) {
 
     public Tunnel getTunnelBySubdomain(String subdomain) {
         log.info("Getting tunnel by subdomain: subdomain={}", subdomain);
-        return repo.findBySubdomain(subdomain).orElseThrow(() -> new BaseException("Tunnel not found!"));
+        return repo.findBySubdomainIgnoreCase(subdomain).orElseThrow(() -> new BaseException("Tunnel not found!"));
     }
 
     public Tunnel getTunnelBySessionId(String id) {

server/src/main/java/uz/server/ws/ConnectionHandler.java

@@ -5,13 +5,14 @@
 import org.springframework.stereotype.Component;
 import org.springframework.web.socket.BinaryMessage;
 import org.springframework.web.socket.CloseStatus;
+import org.springframework.web.socket.TextMessage;
 import org.springframework.web.socket.WebSocketSession;
-import org.springframework.web.socket.handler.BinaryWebSocketHandler;
+import org.springframework.web.socket.handler.TextWebSocketHandler;
 
 @Component
 @RequiredArgsConstructor
 @Slf4j
-public class ConnectionHandler extends BinaryWebSocketHandler {
+public class ConnectionHandler extends TextWebSocketHandler {
     private final EventManager eventManager;
 
     @Override
@@ -27,8 +28,8 @@ public void afterConnectionClosed(WebSocketSession session, CloseStatus status)
     }
 
     @Override
-    protected void handleBinaryMessage(WebSocketSession session, BinaryMessage message) {
+    protected void handleTextMessage(WebSocketSession session, TextMessage message) {
         log.info("Message received: {}", message.getPayload());
-        eventManager.onResponseReceived(message);
+        eventManager.onResponseReceived(message, session.getId());
     }
 }

server/src/main/java/uz/server/ws/EventManager.java

@@ -6,6 +6,7 @@
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.stereotype.Component;
 import org.springframework.web.socket.BinaryMessage;
+import org.springframework.web.socket.TextMessage;
 import org.springframework.web.socket.WebSocketSession;
 import uz.server.domain.entity.Tunnel;
 import uz.server.domain.entity.User;
@@ -35,6 +36,7 @@ public class EventManager {
     private final ObjectMapper objectMapper;
     private final RequestHolder requestHolder;
     private final TunnelService tunnelService;
+    private final ResponseHolder responseHolder;
 
     public void onConnectionEstablished(WebSocketSession session) {
         log.info("WebSocket connection established: sessionId={}", session.getId());
@@ -105,27 +107,31 @@ public void onConnectionClosed(WebSocketSession session) {
         sessionHolder.removeSession(session.getId());
     }
 
-    public void onResponseReceived(BinaryMessage message) {
-        ByteBuffer payload = message.getPayload();
+    public void onResponseReceived(TextMessage message, String sessionId) {
+        log.info("Response received: sessionId={}", sessionId);
+        String payload = message.getPayload();
 
-        byte[] bytes = new byte[payload.remaining()];
-        payload.get(bytes);
+        responseHolder.add(sessionId, payload);
 
-        ObjectMapper objectMapper = new ObjectMapper();
-        Response response;
+        if (message.isLast()){
+            log.info("Response is last: sessionId={}", sessionId);
+            ObjectMapper objectMapper = new ObjectMapper();
+            Response response;
 
-        try {
-            response = objectMapper.readValue(bytes, Response.class);
-        } catch (IOException e) {
-            log.error("Error while deserializing response: {}", e.getMessage());
-            return;
-        }
+            try {
+                response = objectMapper.readValue(responseHolder.get(sessionId), Response.class);
+            } catch (IOException e) {
+                log.error("Error while deserializing response: {}", e.getMessage());
+                return;
+            }
 
-        log.info("Response received: reqId={}, status={}, bodyLength={}",
-                response.getRequestId(), response.getStatus(),
-                response.getBody() != null ? response.getBody().length() : 0);
+            log.info("Response received: reqId={}, status={}, bodyLength={}",
+                    response.getRequestId(), response.getStatus(),
+                    response.getBody() != null ? response.getBody().length() : 0);
 
-        requestHolder.complete(response);
+            responseHolder.remove(sessionId);
+            requestHolder.complete(response);
+        }
     }
 
     public Response sendRequestToCLI(String subdomain, Request request) {

server/src/main/java/uz/server/ws/ResponseHolder.java

@@ -0,0 +1,36 @@
+package uz.server.ws;
+
+import lombok.extern.slf4j.Slf4j;
+import org.springframework.stereotype.Component;
+
+import java.util.concurrent.ConcurrentHashMap;
+
+@Slf4j
+@Component
+public class ResponseHolder {
+    private static final ConcurrentHashMap<String, StringBuilder> responses = new ConcurrentHashMap<>();
+
+    public void add(String id, String response) {
+        log.info("Adding response with id: {}", id);
+        if (responses.containsKey(id)) {
+            log.info("Appending response with id: {}", id);
+            StringBuilder builder = responses.get(id);
+            builder.append(response);
+
+            return;
+        }
+
+        log.info("Creating response with id: {}", id);
+        responses.put(id, new StringBuilder(response));
+    }
+
+    public String get(String id) {
+        log.info("Getting response with id: {}", id);
+        return responses.get(id).toString();
+    }
+
+    public void remove(String id) {
+        log.info("Removing response with id: {}", id);
+        responses.remove(id);
+    }
+}