Fix: Java agent: handle partial payload read
[lttng-ust.git] / liblttng-ust-java-agent / java / lttng-ust-agent-common / org / lttng / ust / agent / client / LttngTcpSessiondClient.java
index 6177feac7a849cefb44e959197e30894b295f529..0f979a4f88da7d7fa97da19c6786f0dcc9fe1e90 100644 (file)
@@ -21,14 +21,16 @@ package org.lttng.ust.agent.client;
 import java.io.BufferedReader;
 import java.io.DataInputStream;
 import java.io.DataOutputStream;
+import java.io.FileInputStream;
 import java.io.FileNotFoundException;
-import java.io.FileReader;
 import java.io.IOException;
+import java.io.InputStreamReader;
 import java.lang.management.ManagementFactory;
 import java.net.Socket;
 import java.net.UnknownHostException;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
+import java.nio.charset.Charset;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 
@@ -44,6 +46,7 @@ public class LttngTcpSessiondClient implements Runnable {
        private static final String SESSION_HOST = "127.0.0.1";
        private static final String ROOT_PORT_FILE = "/var/run/lttng/agent.port";
        private static final String USER_PORT_FILE = "/.lttng/agent.port";
+       private static final Charset PORT_FILE_ENCODING = Charset.forName("UTF-8");
 
        private static final int PROTOCOL_MAJOR_VERSION = 2;
        private static final int PROTOCOL_MINOR_VERSION = 0;
@@ -111,14 +114,14 @@ public class LttngTcpSessiondClient implements Runnable {
                                /*
                                 * Connect to the session daemon before anything else.
                                 */
-                               LttngUstAgentLogger.log(getClass(), "Connecting to sessiond");
+                               log("Connecting to sessiond");
                                connectToSessiond();
 
                                /*
                                 * Register to the session daemon as the Java component of the
                                 * UST application.
                                 */
-                               LttngUstAgentLogger.log(getClass(), "Registering to sessiond");
+                               log("Registering to sessiond");
                                registerToSessiond();
 
                                /*
@@ -126,7 +129,7 @@ public class LttngTcpSessiondClient implements Runnable {
                                 * session daemon. This will return if and only if there is a
                                 * fatal error or the socket closes.
                                 */
-                               LttngUstAgentLogger.log(getClass(), "Waiting on sessiond commands...");
+                               log("Waiting on sessiond commands...");
                                handleSessiondCmd();
                        } catch (UnknownHostException uhe) {
                                uhe.printStackTrace();
@@ -144,7 +147,7 @@ public class LttngTcpSessiondClient implements Runnable {
         * Dispose this client and close any socket connection it may hold.
         */
        public void close() {
-               LttngUstAgentLogger.log(getClass(), "Closing client");
+               log("Closing client");
                this.quit = true;
 
                try {
@@ -157,29 +160,42 @@ public class LttngTcpSessiondClient implements Runnable {
        }
 
        private void connectToSessiond() throws IOException {
-               int port;
+               int rootPort = getPortFromFile(ROOT_PORT_FILE);
+               int userPort = getPortFromFile(getHomePath() + USER_PORT_FILE);
+
+               /*
+                * Check for the edge case of both files existing but pointing to the
+                * same port. In this case, let the root client handle it.
+                */
+               if ((rootPort != 0) && (rootPort == userPort) && (!isRoot)) {
+                       log("User and root config files both point to port " + rootPort +
+                                       ". Letting the root client handle it.");
+                       throw new IOException();
+               }
 
-               if (this.isRoot) {
-                       port = getPortFromFile(ROOT_PORT_FILE);
-                       if (port == 0) {
-                               /* No session daemon available. Stop and retry later. */
-                               throw new IOException();
-                       }
-               } else {
-                       port = getPortFromFile(getHomePath() + USER_PORT_FILE);
-                       if (port == 0) {
-                               /* No session daemon available. Stop and retry later. */
-                               throw new IOException();
-                       }
+               int portToUse = (isRoot ? rootPort : userPort);
+
+               if (portToUse == 0) {
+                       /* No session daemon available. Stop and retry later. */
+                       throw new IOException();
                }
 
-               this.sessiondSock = new Socket(SESSION_HOST, port);
+               this.sessiondSock = new Socket(SESSION_HOST, portToUse);
                this.inFromSessiond = new DataInputStream(sessiondSock.getInputStream());
                this.outToSessiond = new DataOutputStream(sessiondSock.getOutputStream());
        }
 
        private static String getHomePath() {
-               return System.getProperty("user.home");
+               /*
+                * The environment variable LTTNG_HOME overrides HOME if
+                * defined.
+                */
+               String homePath = System.getenv("LTTNG_HOME");
+
+               if (homePath == null) {
+                       homePath = System.getProperty("user.home");
+               }
+               return homePath;
        }
 
        /**
@@ -188,27 +204,34 @@ public class LttngTcpSessiondClient implements Runnable {
         * @return port value if found else 0.
         */
        private static int getPortFromFile(String path) throws IOException {
-               int port;
                BufferedReader br = null;
 
                try {
-                       br = new BufferedReader(new FileReader(path));
+                       br = new BufferedReader(new InputStreamReader(new FileInputStream(path), PORT_FILE_ENCODING));
                        String line = br.readLine();
-                       port = Integer.parseInt(line, 10);
+                       if (line == null) {
+                               /* File exists but is empty. */
+                               return 0;
+                       }
+
+                       int port = Integer.parseInt(line, 10);
                        if (port < 0 || port > 65535) {
                                /* Invalid value. Ignore. */
                                port = 0;
                        }
+                       return port;
+
+               } catch (NumberFormatException e) {
+                       /* File contained something that was not a number. */
+                       return 0;
                } catch (FileNotFoundException e) {
                        /* No port available. */
-                       port = 0;
+                       return 0;
                } finally {
                        if (br != null) {
                                br.close();
                        }
                }
-
-               return port;
        }
 
        private void registerToSessiond() throws IOException {
@@ -231,7 +254,7 @@ public class LttngTcpSessiondClient implements Runnable {
                /* Data read from the socket */
                byte inputData[] = null;
                /* Reply data written to the socket, sent to the sessiond */
-               byte responseData[] = null;
+               LttngAgentResponse response;
 
                while (true) {
                        /* Get header from session daemon. */
@@ -253,82 +276,83 @@ public class LttngTcpSessiondClient implements Runnable {
                                 * We don't send any reply to the registration done command.
                                 * This just marks the end of the initial session setup.
                                 */
-                               LttngUstAgentLogger.log(getClass(), "Registration done");
+                               log("Registration done");
                                continue;
                        }
                        case CMD_LIST:
                        {
                                SessiondCommand listLoggerCmd = new SessiondListLoggersCommand();
-                               LttngAgentResponse response = listLoggerCmd.execute(logAgent);
-                               responseData = response.getBytes();
-                               LttngUstAgentLogger.log(getClass(), "Received list loggers command");
+                               response = listLoggerCmd.execute(logAgent);
+                               log("Received list loggers command");
                                break;
                        }
                        case CMD_EVENT_ENABLE:
                        {
                                if (inputData == null) {
                                        /* Invalid command */
-                                       responseData = LttngAgentResponse.FAILURE_RESPONSE.getBytes();
+                                       response = LttngAgentResponse.FAILURE_RESPONSE;
                                        break;
                                }
                                SessiondCommand enableEventCmd = new SessiondEnableEventCommand(inputData);
-                               LttngAgentResponse response = enableEventCmd.execute(logAgent);
-                               responseData = response.getBytes();
-                               LttngUstAgentLogger.log(getClass(), "Received enable event command");
+                               response = enableEventCmd.execute(logAgent);
+                               log("Received enable event command: " + enableEventCmd.toString());
                                break;
                        }
                        case CMD_EVENT_DISABLE:
                        {
                                if (inputData == null) {
                                        /* Invalid command */
-                                       responseData = LttngAgentResponse.FAILURE_RESPONSE.getBytes();
+                                       response = LttngAgentResponse.FAILURE_RESPONSE;
                                        break;
                                }
                                SessiondCommand disableEventCmd = new SessiondDisableEventCommand(inputData);
-                               LttngAgentResponse response = disableEventCmd.execute(logAgent);
-                               responseData = response.getBytes();
-                               LttngUstAgentLogger.log(getClass(), "Received disable event command");
+                               response = disableEventCmd.execute(logAgent);
+                               log("Received disable event command: " + disableEventCmd.toString());
                                break;
                        }
                        case CMD_APP_CTX_ENABLE:
                        {
                                if (inputData == null) {
                                        /* This commands expects a payload, invalid command */
-                                       responseData = LttngAgentResponse.FAILURE_RESPONSE.getBytes();
+                                       response = LttngAgentResponse.FAILURE_RESPONSE;
                                        break;
                                }
                                SessiondCommand enableAppCtxCmd = new SessiondEnableAppContextCommand(inputData);
-                               LttngAgentResponse response = enableAppCtxCmd.execute(logAgent);
-                               responseData = response.getBytes();
-                               LttngUstAgentLogger.log(getClass(), "Received enable app-context command");
+                               response = enableAppCtxCmd.execute(logAgent);
+                               log("Received enable app-context command");
                                break;
                        }
                        case CMD_APP_CTX_DISABLE:
                        {
                                if (inputData == null) {
                                        /* This commands expects a payload, invalid command */
-                                       responseData = LttngAgentResponse.FAILURE_RESPONSE.getBytes();
+                                       response = LttngAgentResponse.FAILURE_RESPONSE;
                                        break;
                                }
                                SessiondCommand disableAppCtxCmd = new SessiondDisableAppContextCommand(inputData);
-                               LttngAgentResponse response = disableAppCtxCmd.execute(logAgent);
-                               responseData = response.getBytes();
-                               LttngUstAgentLogger.log(getClass(), "Received disable app-context command");
+                               response = disableAppCtxCmd.execute(logAgent);
+                               log("Received disable app-context command");
                                break;
                        }
                        default:
                        {
                                /* Unknown command, send empty reply */
-                               responseData = new byte[4];
-                               ByteBuffer buf = ByteBuffer.wrap(responseData);
-                               buf.order(ByteOrder.BIG_ENDIAN);
-                               LttngUstAgentLogger.log(getClass(), "Received unknown command, ignoring");
+                               response = null;
+                               log("Received unknown command, ignoring");
                                break;
                        }
                        }
 
                        /* Send response to the session daemon. */
-                       LttngUstAgentLogger.log(getClass(), "Sending response");
+                       byte[] responseData;
+                       if (response == null) {
+                               responseData = new byte[4];
+                               ByteBuffer buf = ByteBuffer.wrap(responseData);
+                               buf.order(ByteOrder.BIG_ENDIAN);
+                       } else {
+                               log("Sending response: " + response.toString());
+                               responseData = response.getBytes();
+                       }
                        this.outToSessiond.write(responseData, 0, responseData.length);
                        this.outToSessiond.flush();
                }
@@ -340,10 +364,17 @@ public class LttngTcpSessiondClient implements Runnable {
         */
        private SessiondCommandHeader recvHeader() throws IOException {
                byte data[] = new byte[SessiondCommandHeader.HEADER_SIZE];
+               int bytesLeft = data.length;
+               int bytesOffset = 0;
 
-               int readLen = this.inFromSessiond.read(data, 0, data.length);
-               if (readLen != data.length) {
-                       throw new IOException();
+               while (bytesLeft != 0) {
+                       int bytesRead = this.inFromSessiond.read(data, bytesOffset, bytesLeft);
+
+                       if (bytesRead < 0) {
+                               throw new IOException();
+                       }
+                       bytesLeft -= bytesRead;
+                       bytesOffset += bytesRead;
                }
                return new SessiondCommandHeader(data);
        }
@@ -357,14 +388,32 @@ public class LttngTcpSessiondClient implements Runnable {
         */
        private byte[] recvPayload(SessiondCommandHeader headerCmd) throws IOException {
                byte payload[] = new byte[(int) headerCmd.getDataSize()];
+               int bytesLeft = payload.length;
+               int bytesOffset = 0;
 
                /* Failsafe check so we don't waste our time reading 0 bytes. */
-               if (payload.length == 0) {
+               if (bytesLeft == 0) {
                        return null;
                }
 
-               this.inFromSessiond.read(payload, 0, payload.length);
+               while (bytesLeft != 0) {
+                       int bytesRead = inFromSessiond.read(payload, bytesOffset, bytesLeft);
+
+                       if (bytesRead < 0) {
+                               throw new IOException();
+                       }
+                       bytesLeft -= bytesRead;
+                       bytesOffset += bytesRead;
+               }
                return payload;
        }
 
+       /**
+        * Wrapper for this class's logging, adds the connection's characteristics
+        * to help differentiate between multiple TCP clients.
+        */
+       private void log(String message) {
+               LttngUstAgentLogger.log(getClass(),
+                               "(root=" + isRoot + ", domain=" + domainValue + ") " + message);
+       }
 }
This page took 0.027529 seconds and 4 git commands to generate.