Add support for forwarding stdin data to server, and getting stdout and stderr from server

src/desktop-trampoline.c

@@ -1,5 +1,7 @@
 #include <errno.h>
+#include <fcntl.h>
 #include <stdarg.h>
+#include <stdint.h>
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
@@ -44,6 +46,42 @@ int isValidEnvVar(char *env) {
   return 0;
 }
 
+/**
+ * Reads a string from the socket, reading first 2 bytes to get its length and
+ * then the string itself.
+ */
+ssize_t readDelimitedString(SOCKET socket, char *buffer, size_t bufferSize) {
+  uint16_t outputLength = 0;
+  if (readSocket(socket, &outputLength, sizeof(uint16_t)) < (int)sizeof(uint16_t)) {
+      printSocketError("ERROR: Error reading from socket");
+      return -1;
+  }
+
+  if (outputLength > bufferSize) {
+    fprintf(stderr, "ERROR: received string is bigger than buffer (%d > %zu)", outputLength, bufferSize);
+    return -1;
+  }
+
+  size_t totalBytesRead = 0;
+  ssize_t bytesRead = 0;
+
+  // Read output from server
+  do {
+    bytesRead = readSocket(socket, buffer + totalBytesRead, outputLength - totalBytesRead);
+
+    if (bytesRead == -1) {
+      printSocketError("ERROR: Error reading from socket");
+      return -1;
+    }
+
+    totalBytesRead += bytesRead;
+  } while (bytesRead > 0);
+
+  buffer[totalBytesRead] = '\0';
+
+  return totalBytesRead;
+}
+
 int runTrampolineClient(SOCKET *outSocket, int argc, char **argv, char **envp) {
   char *desktopPortString;
 
@@ -102,29 +140,57 @@ int runTrampolineClient(SOCKET *outSocket, int argc, char **argv, char **envp) {
     WRITE_STRING_OR_EXIT("environment variable", validEnvVars[idx]);
   }
 
-  // TODO: send stdin stuff?
-
   char buffer[BUFFER_LENGTH + 1];
-  size_t totalBytesRead = 0;
-  ssize_t bytesRead = 0;
+  size_t totalBytesWritten = 0;
+  ssize_t bytesToWrite = 0;
 
-  // Read output from server
+  // Make stdin reading non-blocking, to prevent getting stuck when no data is
+  // provided via stdin.
+  int flags = fcntl(STDIN_FILENO, F_GETFL, 0);
+  fcntl(STDIN_FILENO, F_SETFL, flags | O_NONBLOCK);
+
+  // Send stdin data
   do {
-    bytesRead = readSocket(socket, buffer + totalBytesRead, BUFFER_LENGTH - totalBytesRead);
+    bytesToWrite = read(0, buffer, BUFFER_LENGTH);
+
+    if (bytesToWrite == -1) {
+      if (totalBytesWritten == 0) {
+        // No stdin content found, continuing...
+        break;
+      } else {
+        fprintf(stderr, "ERROR: Error reading stdin data");
+        return 1;
+      }
+    }
 
-    if (bytesRead == -1) {
-      printSocketError("ERROR: Error reading from socket");
+    if (writeSocket(socket, buffer, bytesToWrite) != 0) {
+      printSocketError("ERROR: Couldn't send stdin data");
       return 1;
     }
 
-    totalBytesRead += bytesRead;
-  } while (bytesRead > 0);
+    totalBytesWritten += bytesToWrite;
+  } while (bytesToWrite > 0);
 
-  buffer[totalBytesRead] = '\0';
+  writeSocket(socket, "\0", 1);
+
+  // Read stdout from the server
+  if (readDelimitedString(socket, buffer, BUFFER_LENGTH) == -1) {
+    fprintf(stderr, "ERROR: Couldn't read stdout from socket");
+    return 1;
+  }
 
   // Write that output to stdout
   fprintf(stdout, "%s", buffer);
 
+  // Read stderr from the server
+  if (readDelimitedString(socket, buffer, BUFFER_LENGTH) == -1) {
+    fprintf(stderr, "ERROR: Couldn't read stdout from socket");
+    return 1;
+  }
+
+  // Write that output to stderr
+  fprintf(stderr, "%s", buffer);
+
   return 0;
 }
 

test/index.test.js

@@ -9,6 +9,44 @@ const { createServer } = require('net')
 const trampolinePath = getDesktopTrampolinePath()
 const run = promisify(execFile)
 
+async function startTrampolineServer(output, stdout = '', stderr = '') {
+  const server = createServer(socket => {
+    socket.pipe(split2(/\0/)).on('data', data => {
+      output.push(data.toString('utf8'))
+    })
+
+    const buffer = Buffer.alloc(2, 0)
+
+    // Send stdout
+    buffer.writeUInt16LE(stdout.length, 0)
+    socket.write(buffer)
+    socket.write(stdout)
+
+    // Send stderr
+    buffer.writeUInt16LE(stderr.length, 0)
+    socket.write(buffer)
+    socket.write(stderr)
+
+    // Close the socket
+    socket.end()
+  })
+  server.unref()
+
+  const startServer = async () => {
+    return new Promise((resolve, reject) => {
+      server.on('error', e => reject(e))
+      server.listen(0, '127.0.0.1', () => {
+        resolve(server.address().port)
+      })
+    })
+  }
+
+  return {
+    server,
+    port: await startServer(),
+  }
+}
+
 describe('desktop-trampoline', () => {
   it('exists and is a regular file', async () =>
     expect((await stat(trampolinePath)).isFile()).toBe(true))
@@ -19,48 +57,127 @@ describe('desktop-trampoline', () => {
   it('fails when required environment variables are missing', () =>
     expect(run(trampolinePath, ['Username'])).rejects.toThrow())
 
-  it('forwards arguments and valid environment variables correctly', async () => {
-    const output = []
-    const server = createServer(socket => {
-      socket.pipe(split2(/\0/)).on('data', data => {
-        output.push(data.toString('utf8'))
-      })
+  describe('with a trampoline server', () => {
+    let server = null
+    let output = []
+    let baseEnv = {}
 
-      // Don't send anything and just close the socket after the trampoline is
-      // done forwarding data.
-      socket.end()
+    async function configureTrampolineServer(stdout = '', stderr = '') {
+      output = []
+      const serverInfo = await startTrampolineServer(output, stdout, stderr)
+      server = serverInfo.server
+      baseEnv = {
+        DESKTOP_PORT: serverInfo.port,
+      }
+    }
+
+    afterEach(() => {
+      server.close()
     })
-    server.unref()
 
-    const startTrampolineServer = async () => {
-      return new Promise((resolve, reject) => {
-        server.on('error', e => reject(e))
-        server.listen(0, '127.0.0.1', () => {
-          resolve(server.address().port)
+    it('forwards arguments and valid environment variables correctly without stdin', async () => {
+      await configureTrampolineServer()
+
+      const env = {
+        ...baseEnv,
+        DESKTOP_TRAMPOLINE_IDENTIFIER: '123456',
+        DESKTOP_USERNAME: 'sergiou87',
+        DESKTOP_USERNAME_FAKE: 'fake-user',
+        INVALID_VARIABLE: 'foo bar',
+      }
+      const opts = { env }
+
+      await run(trampolinePath, ['baz'], opts)
+
+      const outputArguments = output.slice(1, 2)
+      expect(outputArguments).toStrictEqual(['baz'])
+      // output[2] is the number of env variables
+      const outputEnv = output.slice(3, output.length - 1)
+      expect(outputEnv).toHaveLength(2)
+      expect(outputEnv).toContain('DESKTOP_TRAMPOLINE_IDENTIFIER=123456')
+      expect(outputEnv).toContain(`DESKTOP_USERNAME=sergiou87`)
+
+      expect(output[output.length - 1]).toStrictEqual('')
+    })
+
+    it('forwards arguments, environment variables and stdin correctly', async () => {
+      await configureTrampolineServer()
+
+      const env = {
+        ...baseEnv,
+        DESKTOP_TRAMPOLINE_IDENTIFIER: '123456',
+        DESKTOP_USERNAME: 'sergiou87',
+      }
+      const opts = {
+        env,
+        stdin: 'This is a test\nWith a multiline\nStandard input text',
+      }
+
+      const run = new Promise((resolve, reject) => {
+        const process = execFile(trampolinePath, ['baz'], opts, function (
+          err,
+          stdout,
+          stderr
+        ) {
+          if (!err) {
+            resolve({ stdout, stderr, exitCode: 0 })
+            return
+          }
+
+          reject(err)
         })
+
+        process.stdin.end(
+          'This is a test\nWith a multiline\nStandard input text',
+          'utf-8'
+        )
       })
-    }
+      await run
 
-    const port = await startTrampolineServer()
-    const env = {
-      DESKTOP_TRAMPOLINE_IDENTIFIER: '123456',
-      DESKTOP_PORT: port,
-      DESKTOP_USERNAME: 'sergiou87',
-      DESKTOP_USERNAME_FAKE: 'fake-user',
-      INVALID_VARIABLE: 'foo bar',
-    }
-    const opts = { env }
+      const outputArguments = output.slice(1, 2)
+      expect(outputArguments).toStrictEqual(['baz'])
+      // output[2] is the number of env variables
+      const outputEnv = output.slice(3, output.length - 1)
+      expect(outputEnv).toHaveLength(2)
+      expect(outputEnv).toContain('DESKTOP_TRAMPOLINE_IDENTIFIER=123456')
+      expect(outputEnv).toContain(`DESKTOP_USERNAME=sergiou87`)
+
+      expect(output[output.length - 1]).toBe(
+        'This is a test\nWith a multiline\nStandard input text'
+      )
+    })
 
-    await run(trampolinePath, ['baz'], opts)
+    it('outputs the same stdout received from the server, when no stderr is specified', async () => {
+      await configureTrampolineServer('This is the command stdout', '')
 
-    const outputArguments = output.slice(1, 2)
-    expect(outputArguments).toStrictEqual(['baz'])
-    // output[2] is the number of env variables
-    const outputEnv = output.slice(3)
-    expect(outputEnv).toHaveLength(2)
-    expect(outputEnv).toContain('DESKTOP_TRAMPOLINE_IDENTIFIER=123456')
-    expect(outputEnv).toContain(`DESKTOP_USERNAME=sergiou87`)
+      const opts = { env: baseEnv }
+      const result = await run(trampolinePath, ['baz'], opts)
 
-    server.close()
+      expect(result.stdout).toStrictEqual('This is the command stdout')
+      expect(result.stderr).toStrictEqual('')
+    })
+
+    it('outputs the same stderr received from the server, when no stdout is specified', async () => {
+      await configureTrampolineServer('', 'This is the command stderr')
+
+      const opts = { env: baseEnv }
+      const result = await run(trampolinePath, ['baz'], opts)
+
+      expect(result.stdout).toStrictEqual('')
+      expect(result.stderr).toStrictEqual('This is the command stderr')
+    })
+
+    it('outputs the same stdout and stderr received from the server, when both are specified', async () => {
+      await configureTrampolineServer(
+        'This is the command stdout',
+        'This is the command stderr'
+      )
+
+      const opts = { env: baseEnv }
+      const result = await run(trampolinePath, ['baz'], opts)
+
+      expect(result.stdout).toStrictEqual('This is the command stdout')
+      expect(result.stderr).toStrictEqual('This is the command stderr')
+    })
   })
 })