From f770c284b6220ff18b37e9ff61a6e8cfae9cfab5 Mon Sep 17 00:00:00 2001 From: Ken Sipe Date: Wed, 24 Aug 2016 10:22:14 -0500 Subject: [PATCH] task arguments (#717) * using arguments if provided otherwise using the job arguments * fixing tests * adding additional tests for argument override * command injection fixed for this feature --- .../scheduler/api/JobManagementResource.scala | 3 ++- .../chronos/scheduler/jobs/TaskUtils.scala | 9 ++++++--- .../jobs/JobSchedulerIntegrationTest.scala | 2 +- .../scheduler/jobs/TaskUtilsSpec.scala | 19 +++++++++++++++++++ 4 files changed, 28 insertions(+), 5 deletions(-) diff --git a/src/main/scala/org/apache/mesos/chronos/scheduler/api/JobManagementResource.scala b/src/main/scala/org/apache/mesos/chronos/scheduler/api/JobManagementResource.scala index ee44c2827..5a858ac58 100644 --- a/src/main/scala/org/apache/mesos/chronos/scheduler/api/JobManagementResource.scala +++ b/src/main/scala/org/apache/mesos/chronos/scheduler/api/JobManagementResource.scala @@ -151,7 +151,8 @@ class JobManagementResource @Inject()(val jobScheduler: JobScheduler, require(jobGraph.lookupVertex(jobName).isDefined, "Job '%s' not found".format(jobName)) val job = jobGraph.getJobForName(jobName).get log.info("Manually triggering job:" + jobName) - jobScheduler.taskManager.enqueue(TaskUtils.getTaskId(job, DateTime.now(DateTimeZone.UTC), 0), job.highPriority) + jobScheduler.taskManager.enqueue(TaskUtils.getTaskId(job, DateTime.now(DateTimeZone.UTC), 0, Option(arguments).filter(_.trim.nonEmpty)) + , job.highPriority) Response.noContent().build } catch { case ex: IllegalArgumentException => diff --git a/src/main/scala/org/apache/mesos/chronos/scheduler/jobs/TaskUtils.scala b/src/main/scala/org/apache/mesos/chronos/scheduler/jobs/TaskUtils.scala index 0bb7617d8..3e5e2745e 100644 --- a/src/main/scala/org/apache/mesos/chronos/scheduler/jobs/TaskUtils.scala +++ b/src/main/scala/org/apache/mesos/chronos/scheduler/jobs/TaskUtils.scala @@ -2,8 +2,8 @@ package org.apache.mesos.chronos.scheduler.jobs import java.util.logging.Logger -import org.apache.mesos.chronos.scheduler.state.PersistenceStore import org.apache.mesos.Protos.{TaskID, TaskState, TaskStatus} +import org.apache.mesos.chronos.scheduler.state.PersistenceStore import org.joda.time.{DateTime, DateTimeZone} import scala.collection.mutable @@ -22,6 +22,8 @@ object TaskUtils { val taskIdTemplate = "ct:%d:%d:%s:%s" val argumentsPattern = """(.*)?""".r val taskIdPattern = """ct:(\d+):(\d+):%s:?%s""".format(JobUtils.jobNamePattern, argumentsPattern).r + val commandInjectionFilter = ";".toSet + private[this] val log = Logger.getLogger(getClass.getName) def getTaskStatus(job: BaseJob, due: DateTime, attempt: Int = 0): TaskStatus = { @@ -99,8 +101,9 @@ object TaskUtils { }) } - def getTaskId(job: BaseJob, due: DateTime, attempt: Int = 0): String = { - taskIdTemplate.format(due.getMillis, attempt, job.name, job.arguments.mkString(" ")) + def getTaskId(job: BaseJob, due: DateTime, attempt: Int = 0, arguments: Option[String] = None): String = { + val args: String = arguments.getOrElse(job.arguments.mkString(" ")).filterNot(commandInjectionFilter) + taskIdTemplate.format(due.getMillis, attempt, job.name, args) } def getDueTimes(tasks: Map[String, Array[Byte]]): Map[String, (BaseJob, Long, Int)] = { diff --git a/src/test/scala/org/apache/mesos/chronos/scheduler/jobs/JobSchedulerIntegrationTest.scala b/src/test/scala/org/apache/mesos/chronos/scheduler/jobs/JobSchedulerIntegrationTest.scala index 090294f08..5c18029cd 100644 --- a/src/test/scala/org/apache/mesos/chronos/scheduler/jobs/JobSchedulerIntegrationTest.scala +++ b/src/test/scala/org/apache/mesos/chronos/scheduler/jobs/JobSchedulerIntegrationTest.scala @@ -114,7 +114,7 @@ class JobSchedulerIntegrationTest extends SpecificationWithJUnit with Mockito { jobMarkedSuccess.successCount must_== 1 jobMarkedSuccess.errorsSinceLastSuccess must_== 0 val lastSuccess = DateTime.parse(jobMarkedSuccess.lastSuccess) - there was one(mockTaskManager).enqueue(TaskUtils.getTaskId(dependentJob, lastSuccess, 0), + there was one(mockTaskManager).enqueue(TaskUtils.getTaskId(dependentJob, lastSuccess, 0, None), highPriority = false) scheduler.handleStartedTask(TaskUtils.getTaskStatus(dependentJob, lastSuccess, 0)) scheduler.handleFinishedTask(TaskUtils.getTaskStatus(dependentJob, lastSuccess, 0)) diff --git a/src/test/scala/org/apache/mesos/chronos/scheduler/jobs/TaskUtilsSpec.scala b/src/test/scala/org/apache/mesos/chronos/scheduler/jobs/TaskUtilsSpec.scala index f87ba2f05..699f3f0ad 100644 --- a/src/test/scala/org/apache/mesos/chronos/scheduler/jobs/TaskUtilsSpec.scala +++ b/src/test/scala/org/apache/mesos/chronos/scheduler/jobs/TaskUtilsSpec.scala @@ -10,16 +10,22 @@ class TaskUtilsSpec extends SpecificationWithJUnit with Mockito { "Get taskId" in { val schedule = "R/2012-01-01T00:00:01.000Z/P1M" val arguments = "-a 1 -b 2" + val cmdArgs = "-c 1 -d 2" val job1 = new ScheduleBasedJob(schedule, "sample-name", "sample-command", arguments = List(arguments)) val job2 = new ScheduleBasedJob(schedule, "sample-name", "sample-command") + val job3 = new ScheduleBasedJob(schedule, "sample-name", "sample-command", arguments = List(arguments)) val ts = 1420843781398L val due = new DateTime(ts) val taskIdOne = TaskUtils.getTaskId(job1, due, 0) val taskIdTwo = TaskUtils.getTaskId(job2, due, 0) + val taskIdThree = TaskUtils.getTaskId(job3, due, 0, Option(cmdArgs)) + val taskIdFour = TaskUtils.getTaskId(job2, due, 0, Option(cmdArgs)) taskIdOne must_== "ct:1420843781398:0:sample-name:" + arguments taskIdTwo must_== "ct:1420843781398:0:sample-name:" + taskIdThree must_== "ct:1420843781398:0:sample-name:" + cmdArgs // test override + taskIdFour must_== "ct:1420843781398:0:sample-name:" + cmdArgs // test adding args } "Get job arguments for taskId" in { @@ -30,6 +36,19 @@ class TaskUtilsSpec extends SpecificationWithJUnit with Mockito { jobArguments must_== arguments } + "Disable command injection" in { + val schedule = "R/2012-01-01T00:00:01.000Z/P1M" + val cmdArgs = "-c 1 ; ./evil.sh" + val expectedArgs = "-c 1 ./evil.sh" + val job1 = new ScheduleBasedJob(schedule, "sample-name", "sample-command") + val ts = 1420843781398L + val due = new DateTime(ts) + + val taskIdOne = TaskUtils.getTaskId(job1, due, 0, Option(cmdArgs)) + + taskIdOne must_== "ct:1420843781398:0:sample-name:" + expectedArgs + } + "Parse taskId" in { val arguments = "-a 1 -b 2" val arguments2 = "-a 1:2 --B test"