Skip to content

Commit

Permalink
feat(core): restrict using the read function on the Worker
Browse files Browse the repository at this point in the history
Fixes #3126
  • Loading branch information
loicmathieu committed Mar 5, 2024
1 parent 46e3fa0 commit b002fd6
Show file tree
Hide file tree
Showing 2 changed files with 31 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import io.kestra.core.storages.StorageContext;
import io.kestra.core.storages.StorageInterface;
import io.kestra.core.utils.Slugify;
import io.micronaut.context.annotation.Value;
import io.pebbletemplates.pebble.error.PebbleException;
import io.pebbletemplates.pebble.extension.Function;
import io.pebbletemplates.pebble.template.EvaluationContext;
Expand All @@ -25,13 +26,20 @@ public class ReadFileFunction implements Function {
@Inject
private StorageInterface storageInterface;

@Value("${kestra.server-type}") // default to STANDALONE which should only be the case in tests
private String serverType;

@Override
public List<String> getArgumentNames() {
return List.of("path");
}

@Override
public Object execute(Map<String, Object> args, PebbleTemplate self, EvaluationContext context, int lineNumber) {
if (!calledOnWorker()) {
throw new PebbleException(null, "The 'read' function can only be used in the Worker as it access the internal storage.", lineNumber, self.getName());
}

if (!args.containsKey("path")) {
throw new PebbleException(null, ERROR_MESSAGE, lineNumber, self.getName());
}
Expand All @@ -45,6 +53,17 @@ public Object execute(Map<String, Object> args, PebbleTemplate self, EvaluationC
}
}

private boolean calledOnWorker() {
if ("WORKER".equals(serverType)) {
return true;
}
if ("STANDALONE".equals(serverType)) {
// check that it's called inside a worker thread
return Thread.currentThread().getClass().getName().equals("io.kestra.core.runners.Worker$WorkerThread");
}
return false;
}

@SuppressWarnings("unchecked")
private String readFromNamespaceFile(EvaluationContext context, String path) throws IOException {
Map<String, String> flow = (Map<String, String>) context.getVariable("flow");
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,9 @@
import io.kestra.core.storages.StorageContext;
import io.kestra.core.storages.StorageInterface;
import io.kestra.core.utils.IdUtils;
import io.micronaut.context.annotation.Property;
import io.micronaut.test.extensions.junit5.annotation.MicronautTest;
import io.pebbletemplates.pebble.error.PebbleException;
import jakarta.inject.Inject;
import org.junit.jupiter.api.Test;

Expand All @@ -16,10 +18,12 @@
import java.util.Map;

import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.containsString;
import static org.hamcrest.Matchers.is;
import static org.junit.jupiter.api.Assertions.assertThrows;

@MicronautTest
@MicronautTest(rebuildContext = true)
@Property(name="kestra.server-type", value="WORKER")
class ReadFileFunctionTest {
@Inject
VariableRenderer variableRenderer;
Expand Down Expand Up @@ -131,4 +135,11 @@ void readUnauthorizedInternalStorageFile() throws IOException {
exception = assertThrows(IllegalArgumentException.class, () -> variableRenderer.render("{{ read('" + internalStorageFile + "') }}", triggerVariables));
assertThat(exception.getMessage(), is("Unable to read the file '" + internalStorageFile + "' as it didn't belong to the current execution"));
}

@Test
@Property(name="kestra.server-type", value="EXECUTOR")
void readFailOnNonWorkerNodes() {
IllegalVariableEvaluationException exception = assertThrows(IllegalVariableEvaluationException.class, () -> variableRenderer.render("{{ read('unknown.txt') }}", Map.of("flow", Map.of("namespace", "io.kestra.tests"))));
assertThat(exception.getMessage(), containsString("The 'read' function can only be used in the Worker as it access the internal storage."));
}
}

0 comments on commit b002fd6

Please sign in to comment.