Allow opt-in for skipping trace acknowledgment

Author: Steve Salas

agent/src/main/java/com/codedx/codepulse/agent/agent/DefaultTraceAgent.java

@@ -37,7 +37,7 @@
 import com.codedx.codepulse.agent.message.MessageSenderManager;
 import com.codedx.codepulse.agent.message.PooledBufferService;
 import com.codedx.codepulse.agent.protocol.ProtocolVersion;
-import com.codedx.codepulse.agent.protocol.ProtocolVersion3;
+import com.codedx.codepulse.agent.protocol.ProtocolVersion4;
 import com.codedx.codepulse.agent.trace.TraceDataCollector;
 import com.codedx.codepulse.agent.util.ShutdownHook;
 import com.codedx.codepulse.agent.util.SocketFactory;
@@ -65,7 +65,7 @@ public class DefaultTraceAgent implements TraceAgent
 	private RuntimeAgentConfigurationV1 config;
 
 	private final Semaphore startMutex = new Semaphore(0);
-	private final ProtocolVersion protocol = new ProtocolVersion3();
+	private final ProtocolVersion protocol = new ProtocolVersion4();
 	private MinlogListener logger = null;
 	private ClassIdentifier classIdentifier = new ClassIdentifier();
 	private MethodIdentifier methodIdentifier = new MethodIdentifier();
@@ -197,7 +197,7 @@ public boolean connect(int timeout) throws InterruptedException
 			SocketConnection controlConnection = new SocketConnection(controlSocket, true, true);
 
 			error = "Failed to get configuration from HQ";
-			config = protocol.getControlConnectionHandshake().performHandshake(controlConnection);
+			config = protocol.getControlConnectionHandshake().performHandshake(controlConnection, staticConfig.getProjectId());
 			if (config == null)
 			{
 				controlSocket.close();

agent/src/main/java/com/codedx/codepulse/agent/init/ControlConnectionHandshake.java

@@ -24,5 +24,5 @@
 
 public interface ControlConnectionHandshake
 {
-	RuntimeAgentConfigurationV1 performHandshake(Connection connection) throws IOException;
+	RuntimeAgentConfigurationV1 performHandshake(Connection connection, int projectId) throws IOException;
 }

agent/src/main/java/com/codedx/codepulse/agent/init/ControlConnectionHandshakeV1.java

@@ -21,6 +21,7 @@
 import java.io.DataOutputStream;
 import java.io.IOException;
 
+import com.codedx.codepulse.agent.common.message.NotSupportedException;
 import com.codedx.codepulse.agent.control.ConfigurationReader;
 import com.codedx.codepulse.agent.errors.ErrorHandler;
 import com.codedx.codepulse.agent.common.config.RuntimeAgentConfigurationV1;
@@ -29,7 +30,7 @@
 import com.codedx.codepulse.agent.common.message.MessageProtocol;
 
 /**
- * Implements control connection handshake for protocol version 1.
+ * Implements control connection handshake for protocol version 4.
  * @author RobertF
  */
 public class ControlConnectionHandshakeV1 implements ControlConnectionHandshake
@@ -44,13 +45,22 @@ public ControlConnectionHandshakeV1(MessageProtocol protocol, ConfigurationReade
 	}
 
 	@Override
-	public RuntimeAgentConfigurationV1 performHandshake(Connection connection) throws IOException
+	public RuntimeAgentConfigurationV1 performHandshake(Connection connection, int projectId) throws IOException
 	{
 		DataOutputStream outStream = connection.output();
 		DataInputStream inStream = connection.input();
 
 		// say hello!
-		protocol.writeHello(outStream);
+		try {
+			if (projectId == 0) {
+				protocol.writeHello(outStream);
+			} else {
+				protocol.writeProjectHello(outStream, projectId);
+			}
+		} catch (NotSupportedException e) {
+			ErrorHandler.handleError("protocol error: project-hello message is not supported");
+			return null;
+		}
 		outStream.flush();
 
 		byte reply = inStream.readByte();

agent/src/main/java/com/codedx/codepulse/agent/protocol/ProtocolVersion4.java

@@ -0,0 +1,41 @@
+/*
+ * Copyright 2018 Secure Decisions, a division of Applied Visions, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ * This material is based on research sponsored by the Department of Homeland
+ * Security (DHS) Science and Technology Directorate, Cyber Security Division
+ * (DHS S&T/CSD) via contract number HHSP233201600058C.
+ */
+
+package com.codedx.codepulse.agent.protocol;
+
+import com.codedx.codepulse.agent.common.message.MessageProtocolV4;
+import com.codedx.codepulse.agent.control.ConfigurationReaderV2;
+import com.codedx.codepulse.agent.init.ControlConnectionHandshakeV1;
+import com.codedx.codepulse.agent.init.DataConnectionHandshakeV1;
+
+/**
+ * ProtocolVersion implementation for version 4.
+ * @author ssalas
+ */
+public class ProtocolVersion4 extends ProtocolVersionBase
+{
+	public ProtocolVersion4()
+	{
+		messageProtocol = new MessageProtocolV4();
+		configurationReader = new ConfigurationReaderV2();
+		controlConnectionHandshake = new ControlConnectionHandshakeV1(messageProtocol, configurationReader);
+		dataConnectionHandshake = new DataConnectionHandshakeV1(messageProtocol);
+	}
+}

codepulse/src/main/resources/application.conf

@@ -3,6 +3,9 @@ cp {
 		# The port used to listen for agent connections
 		tracePort = 8765,
 		tracePort=${?CODE_PULSE_TRACE_PORT}
+		# Determines whether Code Pulse can skip user acknowledgment of a trace that specifies a project ID
+		skipUserAcknowledgment=false
+		skipUserAcknowledgment=${?CODE_PULSE_TRACE_SKIP_ACK}
 		symbolService {
 			port = "49582"
 			port = ${?SYMBOL_SERVICE_PORT}

codepulse/src/main/scala/com/secdec/codepulse/package.scala

@@ -80,7 +80,7 @@ package object codepulse {
 		private var config =
 			ConfigFactory
 				.parseFile(configFile)
-  		  .resolve()
+				.resolve()
 				.withFallback(ConfigFactory.load())
 				.withOnlyPath("cp.userSettings") // do not return systemSettings, which could get persisted to disk later on
 
@@ -93,6 +93,8 @@ package object codepulse {
 
 		def symbolServicePort = config.getString("cp.userSettings.symbolService.port")
 
+		def skipUserAcknowledgment = config.getBoolean("cp.userSettings.skipUserAcknowledgment")
+
 		def secdecLoggingLevel: Option[Level] = {
 			getLogLevel(config, "cp.userSettings.logging.secdecLoggingLevel")
 		}
@@ -147,10 +149,14 @@ package object codepulse {
 
 			// preserve option to override port for symbol service via an environment variable
 			val symbolServicePortPattern = """(port="?\d+"?)""".r
-			val outputWithServicePortPattern = symbolServicePortPattern.replaceFirstIn(outputWithTracePort, "$1\n" + " " * 16 + "port=\\${?SYMBOL_SERVICE_PORT}")
+			val outputWithServicePort = symbolServicePortPattern.replaceFirstIn(outputWithTracePort, "$1\n" + " " * 16 + "port=\\${?SYMBOL_SERVICE_PORT}")
+
+			// preserve option to skip user ack via an environment variable
+			val allowAgentToAcceptTraceConnectionPattern = """(skipUserAcknowledgment="?(?:false|true|yes|no)"?)""".r
+			val outputWithSkipAck = allowAgentToAcceptTraceConnectionPattern.replaceFirstIn(outputWithServicePort, "$1\n" + " " * 12 + "skipUserAcknowledgment=\\${?CODE_PULSE_SKIP_ACK}")
 
 			val writer = new FileWriter(configFile)
-			writer.write(outputWithServicePortPattern)
+			writer.write(outputWithSkipAck)
 			writer.close()
 		}
 	}
@@ -161,4 +167,4 @@ package object codepulse {
 		def symbolServiceBinary = config.getString("cp.systemSettings.symbolService.binary")
 		def symbolServiceLocation = config.getString("cp.systemSettings.symbolService.location")
 	}
-}
+}

codepulse/src/main/scala/com/secdec/codepulse/tracer/TraceConnectionLooper.scala

@@ -44,6 +44,7 @@ import akka.actor.actorRef2Scala
 import akka.pattern.ask
 import akka.pattern.pipe
 import akka.util.Timeout
+import com.secdec.codepulse.userSettings
 import net.liftweb.common.Loggable
 import reactive.EventSource
 import reactive.EventStream
@@ -289,6 +290,24 @@ class TraceConnectionLooper(acknowledger: TraceConnectionAcknowledger) extends A
 
 		val acknowledgment = acknowledger.getTraceAcknowledgment(trace)
 
+		if (trace.projectId.isDefined) {
+			val projectId = trace.projectId.get
+
+			if (userSettings.skipUserAcknowledgment) {
+				logger.debug(s"Skipping user acknowledgement for project ID $projectId...")
+
+				val project = com.secdec.codepulse.tracer.projectManager().getProject(com.secdec.codepulse.data.model.ProjectId(projectId))
+				if (project.isDefined) {
+					com.secdec.codepulse.tracer.traceConnectionAcknowledger().acknowledgeCurrentTrace(project.get)
+					logger.debug("Skipped user acknowledgement")
+				} else {
+					logger.error(s"A trace specified project ID $projectId, but there is no project with that ID.")
+				}
+			} else {
+				logger.error(s"A trace specified project ID $projectId, but the value of cp.userSettings.skipUserAcknowledgment prevents skipping user acknowledgement.")
+			}
+		}
+
 		// if the trace 'completes' before it has been acknowledged, we need to look for a new one
 		// (this might happen if the agent gets ctrl+C'd before the ack)
 		for {

dotnet-tracer/main/CodePulse.Client.Test/MessageProtocolTests.cs

@@ -42,6 +42,18 @@ public void WhenWriteHelloMessageCreated()
                 });
         }
 
+        [TestMethod]
+        public void WhenWriteProjectHelloMessageCreated()
+        {
+            TestMessage((writer, protocol) => protocol.WriteProjectHello(writer, 42),
+                (reader, protocol) =>
+                {
+                    Assert.AreEqual(MessageTypes.ProjectHello, reader.ReadByte());
+                    Assert.AreEqual(protocol.ProtocolVersion, reader.ReadByte());
+                    Assert.AreEqual(42, reader.ReadInt32BigEndian());
+                });
+        }
+
         [TestMethod]
         public void WhenWriteDataHelloMessageCreated()
         {

dotnet-tracer/main/CodePulse.Client/Agent/DefaultTraceAgent.cs

@@ -139,7 +139,7 @@ public bool Connect()
             {
                 try
                 {
-                    RuntimeAgentConfiguration = _protocolVersion.ControlConnectionHandshake.PerformHandshake(socketConnection);
+                    RuntimeAgentConfiguration = _protocolVersion.ControlConnectionHandshake.PerformHandshake(socketConnection, StaticAgentConfiguration.ProjectId);
                 }
                 catch (HandshakeException ex)
                 {

dotnet-tracer/main/CodePulse.Client/Config/StaticAgentConfiguration.cs

@@ -35,12 +35,23 @@ public class StaticAgentConfiguration
 
         public int ConnectTimeout { get; }
 
+        public int ProjectId { get; }
+
         public ILog Logger { get; }
 
         public StaticAgentConfiguration(int hqPort,
             string hqHost,
             int connectTimeout,
             ILog logger)
+            : this(hqPort, hqHost, connectTimeout, 0, logger)
+        {
+        }
+
+        public StaticAgentConfiguration(int hqPort,
+            string hqHost,
+            int connectTimeout,
+            int projectId,
+            ILog logger)
         {
             if (hqHost == null)
             {
@@ -54,11 +65,16 @@ public StaticAgentConfiguration(int hqPort,
             {
                 throw new ArgumentOutOfRangeException(nameof(connectTimeout));
             }
+            if (projectId < 0)
+            {
+                throw new ArgumentOutOfRangeException(nameof(projectId));
+            }
 
             HqPort = hqPort;
             HqHost = hqHost;
             ConnectTimeout = connectTimeout;
             Logger = logger ?? throw new ArgumentNullException(nameof(logger));
+            ProjectId = projectId;
         }
     }
 }

dotnet-tracer/main/CodePulse.Client/Init/ControlConnectionHandshake.cs

@@ -41,6 +41,11 @@ public ControlConnectionHandshake(IMessageProtocol messageProtocol, IConfigurati
         }
 
         public RuntimeAgentConfiguration PerformHandshake(IConnection connection)
+        {
+            return PerformHandshake(connection, 0);
+        }
+
+        public RuntimeAgentConfiguration PerformHandshake(IConnection connection, int projectId)
         {
             if (connection == null)
             {
@@ -49,8 +54,16 @@ public RuntimeAgentConfiguration PerformHandshake(IConnection connection)
 
             var outputWriter = connection.OutputWriter;
 
-            _messageProtocol.WriteHello(outputWriter);
-            outputWriter.FlushAndLog("WriteHello");
+            if (projectId == 0)
+            {
+                _messageProtocol.WriteHello(outputWriter);
+                outputWriter.FlushAndLog("WriteHello");
+            }
+            else
+            {
+                _messageProtocol.WriteProjectHello(outputWriter, projectId);
+                outputWriter.FlushAndLog("WriteProjectHello");
+            }
 
             var inputReader = connection.InputReader;
             var reply = inputReader.ReadByte();

dotnet-tracer/main/CodePulse.Client/Init/IControlConnectionHandshake.cs

@@ -28,5 +28,7 @@ namespace CodePulse.Client.Init
     public interface IControlConnectionHandshake
     {
         RuntimeAgentConfiguration PerformHandshake(IConnection connection);
+
+        RuntimeAgentConfiguration PerformHandshake(IConnection connection, int projectId);
     }
 }

dotnet-tracer/main/CodePulse.Client/Message/IMessageProtocol.cs

@@ -29,6 +29,7 @@ public interface IMessageProtocol
         byte ProtocolVersion { get; }
 
         void WriteHello(BinaryWriter writer);
+        void WriteProjectHello(BinaryWriter writer, int projectId);
         void WriteDataHello(BinaryWriter writer, byte runId);
         void WriteError(BinaryWriter writer, string error);
         void WriteHeartbeat(BinaryWriter writer, AgentOperationMode mode, ushort sendBufferSize);

dotnet-tracer/main/CodePulse.Client/Message/MessageProtocol.cs

@@ -28,14 +28,21 @@ namespace CodePulse.Client.Message
 {
     public class MessageProtocol : IMessageProtocol
     {
-        public byte ProtocolVersion { get; } = 3;
+        public byte ProtocolVersion { get; } = 4;
 
         public void WriteHello(BinaryWriter writer)
         {
             writer.Write(MessageTypes.Hello);
             writer.Write(ProtocolVersion);
         }
 
+        public void WriteProjectHello(BinaryWriter writer, int projectId)
+        {
+            writer.Write(MessageTypes.ProjectHello);
+            writer.Write(ProtocolVersion);
+            writer.WriteBigEndian(projectId);
+        }
+
         public void WriteDataHello(BinaryWriter writer, byte runId)
         {
             writer.Write(MessageTypes.DataHello);

dotnet-tracer/main/CodePulse.Client/Message/MessageTypes.cs

@@ -41,6 +41,8 @@ public abstract class MessageTypes
 	    public const byte MapSourceLocation = 13;
 	    public const byte MethodVisit = 14;
 
+        public const byte ProjectHello = 16;
+
         public const byte MethodEntry = 20;
         public const byte MethodExit = 21;
         public const byte Exception = 22;

dotnet-tracer/main/CodePulse.Console.Test/CommandLineParserTests.cs

@@ -281,6 +281,45 @@ public void WhenAppModeWithLogAllSpecifiedSuccess()
             ValidateCommandLineArguments(parameters);
         }
 
+        [TestMethod]
+        public void WhenProjectIdInvalidExceptionOccurs()
+        {
+            // arrange
+            var parameters = new[]
+            {
+                "-Target:file",
+                "-SendVisitPointsTimerInterval:45",
+                "-ProjectId:not-an-id"
+            };
+            ValidateCommandLineArguments(parameters, "The Code Pulse project ID must be a valid identifier: not-an-id is not a valid value for Int32.");
+        }
+
+        [TestMethod]
+        public void WhenProjectIdSpecifiedCorrectly()
+        {
+            // arrange
+            var parameters = new[]
+            {
+                "-Target:file",
+                "-SendVisitPointsTimerInterval:45",
+                "-ProjectId:1"
+            };
+            ValidateCommandLineArguments(parameters);
+        }
+
+        [TestMethod]
+        public void WhenProjectIdNegativeExceptionOccurs()
+        {
+            // arrange
+            var parameters = new[]
+            {
+                "-Target:file",
+                "-SendVisitPointsTimerInterval:45",
+                "-ProjectId:-11"
+            };
+            ValidateCommandLineArguments(parameters, "The Code Pulse project ID must be a valid identifier: The argument projectid must be between 0 and 2147483647");
+        }
+
         private static void ValidateCommandLineArguments(string[] parameters, string expectedError = "")
         {
             // arrange