diff --git a/.github/workflows/deployment-operator-cd-sentinel-harness.yaml b/.github/workflows/deployment-operator-cd-sentinel-harness.yaml index dc95878d3e..d95e9c2e41 100644 --- a/.github/workflows/deployment-operator-cd-sentinel-harness.yaml +++ b/.github/workflows/deployment-operator-cd-sentinel-harness.yaml @@ -202,6 +202,6 @@ jobs: cache-from: type=gha cache-to: type=gha,mode=max build-args: | - GO_VERSION=1.26.5 + GO_VERSION=1.26.6 SENTINEL_HARNESS_BASE_IMAGE_REPO=ghcr.io/pluralsh/sentinel-harness-base - SENTINEL_HARNESS_BASE_IMAGE_TAG=${{ needs.publish-base-image.outputs.version }} \ No newline at end of file + SENTINEL_HARNESS_BASE_IMAGE_TAG=${{ needs.publish-base-image.outputs.version }} diff --git a/.gitignore b/.gitignore index f8987b5140..7b38b9ee87 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,7 @@ /.idea /.vscode/**/* !/.vscode/settings.json +/.codex # System .DS_Store diff --git a/charts/console/values.yaml b/charts/console/values.yaml index 4627a1aa02..5d83e5a2d9 100644 --- a/charts/console/values.yaml +++ b/charts/console/values.yaml @@ -637,8 +637,10 @@ cloudQuery: resources: requests: - memory: 250Mi + memory: 512Mi cpu: 100m + limits: + memory: 1Gi livenessProbe: ~ diff --git a/go/ai-proxy/go.mod b/go/ai-proxy/go.mod index 5d6094334d..44032a7463 100644 --- a/go/ai-proxy/go.mod +++ b/go/ai-proxy/go.mod @@ -1,6 +1,6 @@ module github.com/pluralsh/console/go/ai-proxy -go 1.26.5 +go 1.26.6 replace github.com/pluralsh/console/go/polly => ../polly @@ -36,14 +36,26 @@ require ( github.com/aws/aws-sdk-go-v2/service/sts v1.42.2 // indirect github.com/aws/smithy-go v1.27.1 // indirect github.com/bahlo/generic-list-go v0.2.0 // indirect + github.com/beorn7/perks v1.0.1 // indirect github.com/buger/jsonparser v1.1.2 // indirect github.com/cenkalti/backoff v2.2.1+incompatible // indirect + github.com/cespare/xxhash/v2 v2.3.0 // indirect github.com/go-logr/logr v1.4.3 // indirect - github.com/kr/pretty v0.3.1 // indirect + github.com/golang-jwt/jwt/v5 v5.3.1 // indirect + github.com/jpillora/backoff v1.0.0 // indirect github.com/mailru/easyjson v0.9.1 // indirect + github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect + github.com/mwitkow/go-conntrack v0.0.0-20190716064945-2f068394615f // indirect + github.com/prometheus/client_golang v1.23.2 // indirect + github.com/prometheus/client_model v0.6.2 // indirect + github.com/prometheus/common v0.67.5 // indirect + github.com/prometheus/procfs v0.20.1 // indirect github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect + go.yaml.in/yaml/v2 v2.4.4 // indirect golang.org/x/crypto v0.53.0 // indirect + golang.org/x/net v0.56.0 // indirect golang.org/x/sys v0.46.0 // indirect golang.org/x/text v0.39.0 // indirect + google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af // indirect gopkg.in/yaml.v3 v3.0.1 // indirect ) diff --git a/go/ai-proxy/go.sum b/go/ai-proxy/go.sum index d2aaae9143..d74c0e622a 100644 --- a/go/ai-proxy/go.sum +++ b/go/ai-proxy/go.sum @@ -36,37 +36,44 @@ github.com/aws/smithy-go v1.27.1 h1:4T340VFndXtADGF52gYa1POyL7s9E4Z1OeZ1hCscIw8= github.com/aws/smithy-go v1.27.1/go.mod h1:YE2RhdIuDbA5E5bTdciG9KrW3+TiEONeUWCqxX9i1Fc= github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk= github.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg= +github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM= github.com/buger/jsonparser v1.1.2 h1:frqHqw7otoVbk5M8LlE/L7HTnIq2v9RX6EJ48i9AxJk= github.com/buger/jsonparser v1.1.2/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0= github.com/cenkalti/backoff v2.2.1+incompatible h1:tNowT99t7UNflLxfYYSlKYsBpXdEet03Pg2g16Swow4= github.com/cenkalti/backoff v2.2.1+incompatible/go.mod h1:90ReRw6GdpyfrHakVjL/QHaoyV4aDUVVkXQJJJ3NXXM= -github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= +github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI= github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= +github.com/golang-jwt/jwt/v5 v5.3.1 h1:kYf81DTWFe7t+1VvL7eS+jKFVWaUnK9cB1qbwn63YCY= github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/gorilla/mux v1.8.1 h1:TuBL49tXwgrFYWhqrNgrUNEY92u81SPhu7sTdzQEiWY= github.com/gorilla/mux v1.8.1/go.mod h1:AKf9I4AEqPTmMytcMc0KkNouC66V3BtZ4qD5fmWSiMQ= +github.com/jpillora/backoff v1.0.0 h1:uvFg412JmmHBHw7iwprIxkPMI+sGQ4kzOWsMeHnm2EA= github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= github.com/mailru/easyjson v0.9.1 h1:LbtsOm5WAswyWbvTEOqhypdPeZzHavpZx96/n553mR8= github.com/mailru/easyjson v0.9.1/go.mod h1:1+xMtQp2MRNVL/V1bOzuP3aP8VNwRW55fQUto+XFtTU= +github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA= +github.com/mwitkow/go-conntrack v0.0.0-20190716064945-2f068394615f h1:KUppIJq7/+SVif2QVs3tOP0zanoHgBEVAwHxUSIzRqU= github.com/ollama/ollama v0.21.1 h1:kF0PLEBucnoRDYADS1NmcQXGqC/B7M9WJy8ZyNIr/NA= github.com/ollama/ollama v0.21.1/go.mod h1:274niu48upWz/M7vL53i1WFe+TJRRw5oo4GiacbIYrA= -github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA= github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U= github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= -github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs= -github.com/rogpeppe/go-internal v1.15.0 h1:D0RCU5rMAp+SpgkiNdrjfJ+LX4J1M32V2NeCY7EJ6hc= -github.com/rogpeppe/go-internal v1.15.0/go.mod h1:DrUVZyrJU+txYW5/1kwtXQSMFio52ZOxX7yM1VHvnxs= +github.com/prometheus/client_golang v1.23.2 h1:Je96obch5RDVy3FDMndoUsjAhG5Edi49h0RJWRi/o0o= +github.com/prometheus/client_model v0.6.2 h1:oBsgwpGs7iVziMvrGhE53c/GrLUsZdHnqNwqPLxwZyk= +github.com/prometheus/common v0.67.5 h1:pIgK94WWlQt1WLwAC5j2ynLaBRDiinoAb86HZHTUGI4= +github.com/prometheus/procfs v0.20.1 h1:XwbrGOIplXW/AU3YhIhLODXMJYyC1isLFfYCsTEycfc= github.com/prometheus/sigv4 v0.4.1 h1:EIc3j+8NBea9u1iV6O5ZAN8uvPq2xOIUPcqCTivHuXs= github.com/prometheus/sigv4 v0.4.1/go.mod h1:eu+ZbRvsc5TPiHwqh77OWuCnWK73IdkETYY46P4dXOU= +github.com/rogpeppe/go-internal v1.15.0 h1:D0RCU5rMAp+SpgkiNdrjfJ+LX4J1M32V2NeCY7EJ6hc= +github.com/rogpeppe/go-internal v1.15.0/go.mod h1:DrUVZyrJU+txYW5/1kwtXQSMFio52ZOxX7yM1VHvnxs= github.com/samber/lo v1.53.0 h1:t975lj2py4kJPQ6haz1QMgtId2gtmfktACxIXArw3HM= github.com/samber/lo v1.53.0/go.mod h1:4+MXEGsJzbKGaUEQFKBq2xtfuznW9oz/WrgyzMzRoM0= github.com/spf13/pflag v1.0.10 h1:4EBh2KAYBwaONj6b2Ye1GiHfwjqyROoF4RwYO+vPwFk= @@ -77,8 +84,11 @@ github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/ github.com/wk8/go-ordered-map/v2 v2.1.8/go.mod h1:5nJHM5DyteebpVlHnWMV0rPz6Zp7+xBAnxjb1X5vnTw= github.com/xyproto/randomstring v1.0.5 h1:YtlWPoRdgMu3NZtP45drfy1GKoojuR7hmRcnhZqKjWU= github.com/xyproto/randomstring v1.0.5/go.mod h1:rgmS5DeNXLivK7YprL0pY+lTuhNQW3iGxZ18UQApw/E= +go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= +go.yaml.in/yaml/v2 v2.4.4 h1:tuyd0P+2Ont/d6e2rl3be67goVK4R6deVxCUX5vyPaQ= golang.org/x/crypto v0.53.0 h1:QZ4Muo8THX6CizN2vPPd5fBGHyogrdK9fG4wLPFUsto= golang.org/x/crypto v0.53.0/go.mod h1:DNLU434OwVakk9PzuwV8w62mAJpRJL3vsgcfp4Qnsio= +golang.org/x/net v0.56.0 h1:Rw8j/hFzGvJUZwNBXnAtf5sVDVt+65SK2C7IxCxZt5o= golang.org/x/oauth2 v0.36.0 h1:peZ/1z27fi9hUOFCAZaHyrpWG5lwe0RJEEEeH0ThlIs= golang.org/x/oauth2 v0.36.0/go.mod h1:YDBUJMTkDnJS+A4BP4eZBjCqtokkg1hODuPjwiGPO7Q= golang.org/x/sys v0.46.0 h1:noSf2Fq6F8DBgS+LysIkx7rIExoNHJsxOAtPp4rthXw= @@ -87,9 +97,11 @@ golang.org/x/term v0.44.0 h1:0rLvDRCtNj0gZkyIXhCyOb2OAzEhLVqc4B+hrsBhrmc= golang.org/x/term v0.44.0/go.mod h1:7ze4MdzUzLXpSAoFP1H0bOI9aXDqveSvatT5vKcFh2Y= golang.org/x/text v0.39.0 h1:UbZz4pLOvn600D6Oh6GGEI6VAmndrEBLv8/6BEXzyus= golang.org/x/text v0.39.0/go.mod h1:3UwRclnC2g0TU9x8PZiyfOajCd1zaUNHF9cvqcQZ+ZM= +google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af h1:+5/Sw3GsDNlEmu7TfklWKPdQ0Ykja5VEmq2i817+jbI= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= +gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= k8s.io/klog/v2 v2.140.0 h1:Tf+J3AH7xnUzZyVVXhTgGhEKnFqye14aadWv7bzXdzc= diff --git a/go/client/go.mod b/go/client/go.mod index b7c2df2fd7..1c7943de74 100644 --- a/go/client/go.mod +++ b/go/client/go.mod @@ -1,6 +1,6 @@ module github.com/pluralsh/console/go/client -go 1.26.5 +go 1.26.6 require ( github.com/99designs/gqlgen v0.17.78 @@ -86,7 +86,7 @@ require ( golang.org/x/sys v0.46.0 // indirect golang.org/x/time v0.15.0 // indirect golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da // indirect - google.golang.org/genproto/googleapis/rpc v0.0.0-20260406210006-6f92a3bedf2d // indirect + google.golang.org/genproto/googleapis/rpc v0.0.0-20260414002931-afd174a4e478 // indirect google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af // indirect gopkg.in/ini.v1 v1.67.2 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect diff --git a/go/client/go.sum b/go/client/go.sum index 34a990bce6..fb4f255000 100644 --- a/go/client/go.sum +++ b/go/client/go.sum @@ -295,8 +295,7 @@ golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da h1:noIWHXmPHxILtqtCOPIhS golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da/go.mod h1:NDW/Ps6MPRej6fsCIbMTohpP40sJ/P/vI1MoTEGwX90= google.golang.org/appengine v1.6.8 h1:IhEN5q69dyKagZPYMSdIjS2HqprW324FRQZJcGqPAsM= google.golang.org/appengine v1.6.8/go.mod h1:1jJ3jBArFh5pcgW8gCtRJnepW8FzD1V44FJffLiz/Ds= -google.golang.org/genproto/googleapis/rpc v0.0.0-20260406210006-6f92a3bedf2d h1:wT2n40TBqFY6wiwazVK9/iTWbsQrgk5ZfCSVFLO9LQA= -google.golang.org/genproto/googleapis/rpc v0.0.0-20260406210006-6f92a3bedf2d/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8= +google.golang.org/genproto/googleapis/rpc v0.0.0-20260414002931-afd174a4e478 h1:RmoJA1ujG+/lRGNfUnOMfhCy5EipVMyvUE+KNbPbTlw= google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af h1:+5/Sw3GsDNlEmu7TfklWKPdQ0Ykja5VEmq2i817+jbI= google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= diff --git a/go/cloud-query/README.md b/go/cloud-query/README.md index c68e0da155..5fc1308704 100644 --- a/go/cloud-query/README.md +++ b/go/cloud-query/README.md @@ -7,16 +7,20 @@ Cloud Query is a service part of the Plural Console ecosystem that provides clou - Query cloud resources across multiple providers - Embedded PostgreSQL database for data storage and retrieval through PostgreSQL FDW steampipe extension - gRPC API for integration with other services +- Sandboxed execution for Lua and Monty's limited Python subset - Containerized deployment for easy scaling CloudQuery provider support includes AWS, Azure, GCP, and VMware vSphere. ## Prerequisites -- Go 1.24.2 or higher +- Go 1.26.6 or higher - Docker (for containerized deployment) - Make +Running the complete service locally requires a writable temporary directory. +The Python worker extracts gomonty's embedded native library there on first use. + ## Getting Started ### Local Development @@ -89,6 +93,14 @@ Cloud-Query also exposes ToolQuery gRPC endpoints for observability tools (metri ToolQuery also supports cloud function invocation via `InvokeLambda` for AWS Lambda, GCP Cloud Run services (Gen2), and Azure Functions using canonical identifiers and cloud connection credentials. +### Sandboxed Python execution + +`RunPython` executes Monty's limited Python subset in a crash-isolated `cloud-query python-worker` subprocess. The request's optional JSON object is available as `input`; `output` starts as an empty dictionary and must remain a JSON-serializable dictionary. The response returns that dictionary as `result_json` and standard-output text from `print()` separately as `stdout`. + +This is not CPython. The sandbox has no host filesystem, environment, network, subprocess, shell, `pip`, third-party packages, or callback access. Two workers start with only `TMPDIR=/tmp`; each request gets a fresh gomonty REPL. Failed or canceled workers are killed and replaced, and healthy workers are periodically recycled. Source is limited to 64 KiB, input and result JSON to 1 MiB, stdout to 64 KiB, interpreter execution to 60 seconds, interpreter-managed memory to 100 MiB, and recursion to 100 frames. Two runs execute concurrently and up to 16 more wait in a FIFO queue. A full queue is rejected. A parent watchdog ends a run after 65 seconds or the caller's earlier deadline. + +The image embeds gomonty `v0.0.14` and its platform-specific glibc library, built against official Monty commit `c9802b5f30d11fecf9f153feb1dfdab3abda070e`. It contains no separate `monty` executable. Both pins are recorded in OCI labels. Monty's memory limit covers interpreter-managed allocations rather than total pod RSS; operators should measure the workload before reducing the default cloud-query memory allocation. + ### Tool Provider Credentials and Permissions - `Dynatrace`: diff --git a/go/cloud-query/api/proto/toolquery.proto b/go/cloud-query/api/proto/toolquery.proto index 29e9cbb977..c4e22f390c 100644 --- a/go/cloud-query/api/proto/toolquery.proto +++ b/go/cloud-query/api/proto/toolquery.proto @@ -300,6 +300,22 @@ message RunLuaOutput { string result_json = 1; } +message RunPythonInput { + // Python source. Write structured results to the global `output` dictionary. + string script = 1; + + // Optional JSON object exposed as the global `input` dictionary. Empty means `{}`. + string input_json = 2; +} + +message RunPythonOutput { + // JSON-encoded contents of the global `output` dictionary. + string result_json = 1; + + // Text emitted by print() during the user script. + string stdout = 2; +} + service ToolQuery { rpc Metrics(MetricsQueryInput) returns (MetricsQueryOutput) {} rpc MetricsSearch(MetricsSearchInput) returns (MetricsSearchOutput) {} @@ -308,4 +324,5 @@ service ToolQuery { rpc Traces(TracesQueryInput) returns (TracesQueryOutput) {} rpc InvokeLambda(InvokeLambdaInput) returns (InvokeLambdaOutput) {} rpc RunLua(RunLuaInput) returns (RunLuaOutput) {} + rpc RunPython(RunPythonInput) returns (RunPythonOutput) {} } diff --git a/go/cloud-query/cmd/main.go b/go/cloud-query/cmd/main.go index a47c08f509..c0c27efb24 100644 --- a/go/cloud-query/cmd/main.go +++ b/go/cloud-query/cmd/main.go @@ -13,6 +13,7 @@ import ( "github.com/pluralsh/console/go/cloud-query/internal/pool" "github.com/pluralsh/console/go/cloud-query/internal/server" "github.com/pluralsh/console/go/cloud-query/internal/service" + pythontools "github.com/pluralsh/console/go/cloud-query/internal/tools/python" ) func startHealthzHandler() { @@ -27,9 +28,20 @@ func startHealthzHandler() { } func main() { + if len(os.Args) == 2 && os.Args[1] == "python-worker" { + if err := pythontools.NewWorker().Run(os.Stdin, os.Stdout); err != nil { + os.Exit(1) + } + return + } + startHealthzHandler() - services := []service.Service{service.NewToolQueryService()} + toolQueryService, err := service.NewToolQueryService(context.Background()) + if err != nil { + klog.Fatalf("failed to initialize tool query service: %v", err) + } + services := []service.Service{toolQueryService} if args.DatabaseEnabled() { p, err := pool.NewConnectionPool(args.DatabaseConnectionTTL()) @@ -67,8 +79,5 @@ func handleShutdown(cancel context.CancelFunc, s *server.Server) { <-signalChan klog.Info("received shutdown signal, shutting down gracefully...") - s.Stop() cancel() - klog.Info("stopped gracefully") - os.Exit(0) } diff --git a/go/cloud-query/docs/api-reference.md b/go/cloud-query/docs/api-reference.md index fd4da9019e..ac0616d967 100644 --- a/go/cloud-query/docs/api-reference.md +++ b/go/cloud-query/docs/api-reference.md @@ -405,6 +405,8 @@ service ToolQuery { rpc Logs(LogsQueryInput) returns (LogsQueryOutput) {} rpc Traces(TracesQueryInput) returns (TracesQueryOutput) {} rpc InvokeLambda(InvokeLambdaInput) returns (InvokeLambdaOutput) {} + rpc RunLua(RunLuaInput) returns (RunLuaOutput) {} + rpc RunPython(RunPythonInput) returns (RunPythonOutput) {} } ``` @@ -1507,6 +1509,24 @@ message JaegerTracesOptions { } ``` +## Run Python + +`RunPython` synchronously executes Monty's limited Python subset in a fresh logical sandbox session. `input_json` is optional but, when present, must encode an object. It is exposed as the global `input`; scripts write their structured response to the global `output` dictionary. Printed text is returned separately. + +```protobuf +message RunPythonInput { + string script = 1; + string input_json = 2; +} + +message RunPythonOutput { + string result_json = 1; + string stdout = 2; +} +``` + +The runtime exposes no host filesystem, environment, network, subprocess, shell, package installation, third-party package, or host-tool callback. It limits source to 64 KiB, input and result JSON to 1 MiB, stdout to 64 KiB, execution to 60 seconds, memory to 100 MiB, recursion to 100 frames, wall time to 65 seconds, and concurrency to two runs per process. Up to 16 additional requests wait in a bounded FIFO queue. It uses gomonty `v0.0.14`, built against official Monty commit `c9802b5f30d11fecf9f153feb1dfdab3abda070e`; it is not CPython. + ## Invoke Lambda `InvokeLambda` invokes serverless functions using canonical provider identifiers only. diff --git a/go/cloud-query/go.mod b/go/cloud-query/go.mod index c628c63a5e..c5524ec57a 100644 --- a/go/cloud-query/go.mod +++ b/go/cloud-query/go.mod @@ -1,6 +1,6 @@ module github.com/pluralsh/console/go/cloud-query -go 1.26.5 +go 1.26.6 require ( github.com/Azure/azure-sdk-for-go/sdk/azcore v1.21.1 @@ -18,6 +18,7 @@ require ( github.com/aws/aws-sdk-go-v2/service/lambda v1.89.1 github.com/aws/aws-sdk-go-v2/service/sts v1.42.2 github.com/elastic/go-elasticsearch/v9 v9.3.1 + github.com/ewhauser/gomonty v0.0.14 github.com/gofrs/uuid v4.4.0+incompatible github.com/lib/pq v1.12.3 github.com/orcaman/concurrent-map/v2 v2.0.1 @@ -58,6 +59,7 @@ require ( github.com/aws/smithy-go v1.27.1 // indirect github.com/beorn7/perks v1.0.1 // indirect github.com/cespare/xxhash/v2 v2.3.0 // indirect + github.com/ebitengine/purego v0.10.0 // indirect github.com/elastic/elastic-transport-go/v8 v8.8.0 // indirect github.com/felixge/httpsnoop v1.0.4 // indirect github.com/go-logr/logr v1.4.3 // indirect @@ -80,6 +82,8 @@ require ( github.com/prometheus/client_model v0.6.2 // indirect github.com/prometheus/procfs v0.20.1 // indirect github.com/rogpeppe/go-internal v1.15.0 // indirect + github.com/vmihailenco/msgpack/v5 v5.4.1 // indirect + github.com/vmihailenco/tagparser/v2 v2.0.0 // indirect github.com/xeipuuv/gojsonpointer v0.0.0-20190905194746-02993c407bfb // indirect github.com/xeipuuv/gojsonreference v0.0.0-20180127040603-bd5ef7bd5415 // indirect github.com/xeipuuv/gojsonschema v1.2.0 // indirect diff --git a/go/cloud-query/go.sum b/go/cloud-query/go.sum index 730935a9b3..e698498192 100644 --- a/go/cloud-query/go.sum +++ b/go/cloud-query/go.sum @@ -78,10 +78,14 @@ github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSs github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/ebitengine/purego v0.10.0 h1:QIw4xfpWT6GWTzaW5XEKy3HXoqrJGx1ijYHzTF0/ISU= +github.com/ebitengine/purego v0.10.0/go.mod h1:iIjxzd6CiRiOG0UyXP+V1+jWqUXVjPKLAI0mRfJZTmQ= github.com/elastic/elastic-transport-go/v8 v8.8.0 h1:7k1Ua+qluFr6p1jfJjGDl97ssJS/P7cHNInzfxgBQAo= github.com/elastic/elastic-transport-go/v8 v8.8.0/go.mod h1:YLHer5cj0csTzNFXoNQ8qhtGY1GTvSqPnKWKaqQE3Hk= github.com/elastic/go-elasticsearch/v9 v9.3.1 h1:v5A9uFw0nLFA0luD3xAqliBXbscfuhch409HIinfhKY= github.com/elastic/go-elasticsearch/v9 v9.3.1/go.mod h1:B5u4H2jo2/v0+PrgbmIUdEyHdenFyavWtjciAFl7TA0= +github.com/ewhauser/gomonty v0.0.14 h1:DM+iSZ/WJzl+huEH5SGONLl8CTooQrfVFbC0xY1ghO4= +github.com/ewhauser/gomonty v0.0.14/go.mod h1:XCLUVfUFX733MIidhDw3iLWwaVNxT23i9xK+ioAAOVc= github.com/felixge/httpsnoop v1.0.4 h1:NFTV2Zj1bL4mc9sqWACXbQFVBBg2W3GPvqp8/ESS2Wg= github.com/felixge/httpsnoop v1.0.4/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U= github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= @@ -163,6 +167,10 @@ github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+ github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= +github.com/vmihailenco/msgpack/v5 v5.4.1 h1:cQriyiUvjTwOHg8QZaPihLWeRAAVoCpE00IUPn0Bjt8= +github.com/vmihailenco/msgpack/v5 v5.4.1/go.mod h1:GaZTsDaehaPpQVyxrf5mtQlH+pc21PIudVV/E3rRQok= +github.com/vmihailenco/tagparser/v2 v2.0.0 h1:y09buUbR+b5aycVFQs/g70pqKVZNBmxwAhO7/IwNM9g= +github.com/vmihailenco/tagparser/v2 v2.0.0/go.mod h1:Wri+At7QHww0WTrCBeu4J6bNtoV6mEfg5OIWRZA9qds= github.com/xeipuuv/gojsonpointer v0.0.0-20180127040702-4e3ac2762d5f/go.mod h1:N2zxlSyiKSe5eX1tZViRH5QA0qijqEDrYZiPEAiq3wU= github.com/xeipuuv/gojsonpointer v0.0.0-20190905194746-02993c407bfb h1:zGWFAtiMcyryUHoUjUJX0/lt1H2+i2Ka2n+D3DImSNo= github.com/xeipuuv/gojsonpointer v0.0.0-20190905194746-02993c407bfb/go.mod h1:N2zxlSyiKSe5eX1tZViRH5QA0qijqEDrYZiPEAiq3wU= diff --git a/go/cloud-query/internal/proto/toolquery/toolquery.pb.go b/go/cloud-query/internal/proto/toolquery/toolquery.pb.go index d004253d7e..f68d0754cf 100644 --- a/go/cloud-query/internal/proto/toolquery/toolquery.pb.go +++ b/go/cloud-query/internal/proto/toolquery/toolquery.pb.go @@ -1,6 +1,6 @@ // Code generated by protoc-gen-go. DO NOT EDIT. // versions: -// protoc-gen-go v1.36.11 +// protoc-gen-go v1.36.11-devel // protoc v6.31.1 // source: toolquery.proto @@ -2918,6 +2918,114 @@ func (x *RunLuaOutput) GetResultJson() string { return "" } +type RunPythonInput struct { + state protoimpl.MessageState `protogen:"open.v1"` + // Python source. Write structured results to the global `output` dictionary. + Script string `protobuf:"bytes,1,opt,name=script,proto3" json:"script,omitempty"` + // Optional JSON object exposed as the global `input` dictionary. Empty means `{}`. + InputJson string `protobuf:"bytes,2,opt,name=input_json,json=inputJson,proto3" json:"input_json,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *RunPythonInput) Reset() { + *x = RunPythonInput{} + mi := &file_toolquery_proto_msgTypes[44] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *RunPythonInput) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*RunPythonInput) ProtoMessage() {} + +func (x *RunPythonInput) ProtoReflect() protoreflect.Message { + mi := &file_toolquery_proto_msgTypes[44] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use RunPythonInput.ProtoReflect.Descriptor instead. +func (*RunPythonInput) Descriptor() ([]byte, []int) { + return file_toolquery_proto_rawDescGZIP(), []int{44} +} + +func (x *RunPythonInput) GetScript() string { + if x != nil { + return x.Script + } + return "" +} + +func (x *RunPythonInput) GetInputJson() string { + if x != nil { + return x.InputJson + } + return "" +} + +type RunPythonOutput struct { + state protoimpl.MessageState `protogen:"open.v1"` + // JSON-encoded contents of the global `output` dictionary. + ResultJson string `protobuf:"bytes,1,opt,name=result_json,json=resultJson,proto3" json:"result_json,omitempty"` + // Text emitted by print() during the user script. + Stdout string `protobuf:"bytes,2,opt,name=stdout,proto3" json:"stdout,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *RunPythonOutput) Reset() { + *x = RunPythonOutput{} + mi := &file_toolquery_proto_msgTypes[45] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *RunPythonOutput) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*RunPythonOutput) ProtoMessage() {} + +func (x *RunPythonOutput) ProtoReflect() protoreflect.Message { + mi := &file_toolquery_proto_msgTypes[45] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use RunPythonOutput.ProtoReflect.Descriptor instead. +func (*RunPythonOutput) Descriptor() ([]byte, []int) { + return file_toolquery_proto_rawDescGZIP(), []int{45} +} + +func (x *RunPythonOutput) GetResultJson() string { + if x != nil { + return x.ResultJson + } + return "" +} + +func (x *RunPythonOutput) GetStdout() string { + if x != nil { + return x.Stdout + } + return "" +} + var File_toolquery_proto protoreflect.FileDescriptor const file_toolquery_proto_rawDesc = "" + @@ -3232,7 +3340,15 @@ const file_toolquery_proto_rawDesc = "" + "\x06script\x18\x01 \x01(\tR\x06script\"/\n" + "\fRunLuaOutput\x12\x1f\n" + "\vresult_json\x18\x01 \x01(\tR\n" + - "resultJson2\x9c\x04\n" + + "resultJson\"G\n" + + "\x0eRunPythonInput\x12\x16\n" + + "\x06script\x18\x01 \x01(\tR\x06script\x12\x1d\n" + + "\n" + + "input_json\x18\x02 \x01(\tR\tinputJson\"J\n" + + "\x0fRunPythonOutput\x12\x1f\n" + + "\vresult_json\x18\x01 \x01(\tR\n" + + "resultJson\x12\x16\n" + + "\x06stdout\x18\x02 \x01(\tR\x06stdout2\xe2\x04\n" + "\tToolQuery\x12H\n" + "\aMetrics\x12\x1c.toolquery.MetricsQueryInput\x1a\x1d.toolquery.MetricsQueryOutput\"\x00\x12P\n" + "\rMetricsSearch\x12\x1d.toolquery.MetricsSearchInput\x1a\x1e.toolquery.MetricsSearchOutput\"\x00\x12_\n" + @@ -3240,7 +3356,8 @@ const file_toolquery_proto_rawDesc = "" + "\x04Logs\x12\x19.toolquery.LogsQueryInput\x1a\x1a.toolquery.LogsQueryOutput\"\x00\x12E\n" + "\x06Traces\x12\x1b.toolquery.TracesQueryInput\x1a\x1c.toolquery.TracesQueryOutput\"\x00\x12M\n" + "\fInvokeLambda\x12\x1c.toolquery.InvokeLambdaInput\x1a\x1d.toolquery.InvokeLambdaOutput\"\x00\x12;\n" + - "\x06RunLua\x12\x16.toolquery.RunLuaInput\x1a\x17.toolquery.RunLuaOutput\"\x00BEZCgithub.com/pluralsh/console/go/cloud-query/internal/proto/toolqueryb\x06proto3" + "\x06RunLua\x12\x16.toolquery.RunLuaInput\x1a\x17.toolquery.RunLuaOutput\"\x00\x12D\n" + + "\tRunPython\x12\x19.toolquery.RunPythonInput\x1a\x1a.toolquery.RunPythonOutput\"\x00BEZCgithub.com/pluralsh/console/go/cloud-query/internal/proto/toolqueryb\x06proto3" var ( file_toolquery_proto_rawDescOnce sync.Once @@ -3254,7 +3371,7 @@ func file_toolquery_proto_rawDescGZIP() []byte { return file_toolquery_proto_rawDescData } -var file_toolquery_proto_msgTypes = make([]protoimpl.MessageInfo, 47) +var file_toolquery_proto_msgTypes = make([]protoimpl.MessageInfo, 49) var file_toolquery_proto_goTypes = []any{ (*ElasticConnection)(nil), // 0: toolquery.ElasticConnection (*OpensearchConnection)(nil), // 1: toolquery.OpensearchConnection @@ -3300,11 +3417,13 @@ var file_toolquery_proto_goTypes = []any{ (*InvokeLambdaOutput)(nil), // 41: toolquery.InvokeLambdaOutput (*RunLuaInput)(nil), // 42: toolquery.RunLuaInput (*RunLuaOutput)(nil), // 43: toolquery.RunLuaOutput - nil, // 44: toolquery.MetricPoint.LabelsEntry - nil, // 45: toolquery.LogEntry.LabelsEntry - nil, // 46: toolquery.TraceSpan.TagsEntry - (*timestamppb.Timestamp)(nil), // 47: google.protobuf.Timestamp - (*cloudquery.Connection)(nil), // 48: cloudquery.Connection + (*RunPythonInput)(nil), // 44: toolquery.RunPythonInput + (*RunPythonOutput)(nil), // 45: toolquery.RunPythonOutput + nil, // 46: toolquery.MetricPoint.LabelsEntry + nil, // 47: toolquery.LogEntry.LabelsEntry + nil, // 48: toolquery.TraceSpan.TagsEntry + (*timestamppb.Timestamp)(nil), // 49: google.protobuf.Timestamp + (*cloudquery.Connection)(nil), // 50: cloudquery.Connection } var file_toolquery_proto_depIdxs = []int32{ 0, // 0: toolquery.ToolConnection.elastic:type_name -> toolquery.ElasticConnection @@ -3318,8 +3437,8 @@ var file_toolquery_proto_depIdxs = []int32{ 10, // 8: toolquery.ToolConnection.azure:type_name -> toolquery.AzureConnection 6, // 9: toolquery.ToolConnection.jaeger:type_name -> toolquery.JaegerConnection 1, // 10: toolquery.ToolConnection.opensearch:type_name -> toolquery.OpensearchConnection - 47, // 11: toolquery.TimeRange.start:type_name -> google.protobuf.Timestamp - 47, // 12: toolquery.TimeRange.end:type_name -> google.protobuf.Timestamp + 49, // 11: toolquery.TimeRange.start:type_name -> google.protobuf.Timestamp + 49, // 12: toolquery.TimeRange.end:type_name -> google.protobuf.Timestamp 11, // 13: toolquery.MetricsQueryInput.connection:type_name -> toolquery.ToolConnection 12, // 14: toolquery.MetricsQueryInput.range:type_name -> toolquery.TimeRange 14, // 15: toolquery.MetricsQueryInput.options:type_name -> toolquery.MetricsOptions @@ -3334,8 +3453,8 @@ var file_toolquery_proto_depIdxs = []int32{ 21, // 24: toolquery.TracesQueryInput.options:type_name -> toolquery.TracesOptions 23, // 25: toolquery.TracesOptions.jaeger:type_name -> toolquery.JaegerTracesOptions 22, // 26: toolquery.JaegerTracesOptions.attributes:type_name -> toolquery.JaegerTraceQueryAttribute - 47, // 27: toolquery.MetricPoint.timestamp:type_name -> google.protobuf.Timestamp - 44, // 28: toolquery.MetricPoint.labels:type_name -> toolquery.MetricPoint.LabelsEntry + 49, // 27: toolquery.MetricPoint.timestamp:type_name -> google.protobuf.Timestamp + 46, // 28: toolquery.MetricPoint.labels:type_name -> toolquery.MetricPoint.LabelsEntry 24, // 29: toolquery.MetricsQueryOutput.metrics:type_name -> toolquery.MetricPoint 11, // 30: toolquery.MetricsSearchInput.connection:type_name -> toolquery.ToolConnection 27, // 31: toolquery.MetricsSearchInput.options:type_name -> toolquery.MetricsSearchOptions @@ -3345,14 +3464,14 @@ var file_toolquery_proto_depIdxs = []int32{ 32, // 35: toolquery.MetricsLabelSearchInput.options:type_name -> toolquery.MetricsLabelSearchOptions 33, // 36: toolquery.MetricsLabelSearchOptions.azure:type_name -> toolquery.AzureMetricsLabelSearchOptions 34, // 37: toolquery.MetricsLabelSearchOutput.results:type_name -> toolquery.MetricsLabelSearchResult - 47, // 38: toolquery.LogEntry.timestamp:type_name -> google.protobuf.Timestamp - 45, // 39: toolquery.LogEntry.labels:type_name -> toolquery.LogEntry.LabelsEntry + 49, // 38: toolquery.LogEntry.timestamp:type_name -> google.protobuf.Timestamp + 47, // 39: toolquery.LogEntry.labels:type_name -> toolquery.LogEntry.LabelsEntry 36, // 40: toolquery.LogsQueryOutput.logs:type_name -> toolquery.LogEntry - 47, // 41: toolquery.TraceSpan.start:type_name -> google.protobuf.Timestamp - 47, // 42: toolquery.TraceSpan.end:type_name -> google.protobuf.Timestamp - 46, // 43: toolquery.TraceSpan.tags:type_name -> toolquery.TraceSpan.TagsEntry + 49, // 41: toolquery.TraceSpan.start:type_name -> google.protobuf.Timestamp + 49, // 42: toolquery.TraceSpan.end:type_name -> google.protobuf.Timestamp + 48, // 43: toolquery.TraceSpan.tags:type_name -> toolquery.TraceSpan.TagsEntry 38, // 44: toolquery.TracesQueryOutput.spans:type_name -> toolquery.TraceSpan - 48, // 45: toolquery.InvokeLambdaInput.connection:type_name -> cloudquery.Connection + 50, // 45: toolquery.InvokeLambdaInput.connection:type_name -> cloudquery.Connection 13, // 46: toolquery.ToolQuery.Metrics:input_type -> toolquery.MetricsQueryInput 26, // 47: toolquery.ToolQuery.MetricsSearch:input_type -> toolquery.MetricsSearchInput 31, // 48: toolquery.ToolQuery.MetricsLabelSearch:input_type -> toolquery.MetricsLabelSearchInput @@ -3360,15 +3479,17 @@ var file_toolquery_proto_depIdxs = []int32{ 20, // 50: toolquery.ToolQuery.Traces:input_type -> toolquery.TracesQueryInput 40, // 51: toolquery.ToolQuery.InvokeLambda:input_type -> toolquery.InvokeLambdaInput 42, // 52: toolquery.ToolQuery.RunLua:input_type -> toolquery.RunLuaInput - 25, // 53: toolquery.ToolQuery.Metrics:output_type -> toolquery.MetricsQueryOutput - 30, // 54: toolquery.ToolQuery.MetricsSearch:output_type -> toolquery.MetricsSearchOutput - 35, // 55: toolquery.ToolQuery.MetricsLabelSearch:output_type -> toolquery.MetricsLabelSearchOutput - 37, // 56: toolquery.ToolQuery.Logs:output_type -> toolquery.LogsQueryOutput - 39, // 57: toolquery.ToolQuery.Traces:output_type -> toolquery.TracesQueryOutput - 41, // 58: toolquery.ToolQuery.InvokeLambda:output_type -> toolquery.InvokeLambdaOutput - 43, // 59: toolquery.ToolQuery.RunLua:output_type -> toolquery.RunLuaOutput - 53, // [53:60] is the sub-list for method output_type - 46, // [46:53] is the sub-list for method input_type + 44, // 53: toolquery.ToolQuery.RunPython:input_type -> toolquery.RunPythonInput + 25, // 54: toolquery.ToolQuery.Metrics:output_type -> toolquery.MetricsQueryOutput + 30, // 55: toolquery.ToolQuery.MetricsSearch:output_type -> toolquery.MetricsSearchOutput + 35, // 56: toolquery.ToolQuery.MetricsLabelSearch:output_type -> toolquery.MetricsLabelSearchOutput + 37, // 57: toolquery.ToolQuery.Logs:output_type -> toolquery.LogsQueryOutput + 39, // 58: toolquery.ToolQuery.Traces:output_type -> toolquery.TracesQueryOutput + 41, // 59: toolquery.ToolQuery.InvokeLambda:output_type -> toolquery.InvokeLambdaOutput + 43, // 60: toolquery.ToolQuery.RunLua:output_type -> toolquery.RunLuaOutput + 45, // 61: toolquery.ToolQuery.RunPython:output_type -> toolquery.RunPythonOutput + 54, // [54:62] is the sub-list for method output_type + 46, // [46:54] is the sub-list for method input_type 46, // [46:46] is the sub-list for extension type_name 46, // [46:46] is the sub-list for extension extendee 0, // [0:46] is the sub-list for field type_name @@ -3420,7 +3541,7 @@ func file_toolquery_proto_init() { GoPackagePath: reflect.TypeOf(x{}).PkgPath(), RawDescriptor: unsafe.Slice(unsafe.StringData(file_toolquery_proto_rawDesc), len(file_toolquery_proto_rawDesc)), NumEnums: 0, - NumMessages: 47, + NumMessages: 49, NumExtensions: 0, NumServices: 1, }, diff --git a/go/cloud-query/internal/proto/toolquery/toolquery_grpc.pb.go b/go/cloud-query/internal/proto/toolquery/toolquery_grpc.pb.go index f990f45cbe..400a79eb5d 100644 --- a/go/cloud-query/internal/proto/toolquery/toolquery_grpc.pb.go +++ b/go/cloud-query/internal/proto/toolquery/toolquery_grpc.pb.go @@ -26,6 +26,7 @@ const ( ToolQuery_Traces_FullMethodName = "/toolquery.ToolQuery/Traces" ToolQuery_InvokeLambda_FullMethodName = "/toolquery.ToolQuery/InvokeLambda" ToolQuery_RunLua_FullMethodName = "/toolquery.ToolQuery/RunLua" + ToolQuery_RunPython_FullMethodName = "/toolquery.ToolQuery/RunPython" ) // ToolQueryClient is the client API for ToolQuery service. @@ -39,6 +40,7 @@ type ToolQueryClient interface { Traces(ctx context.Context, in *TracesQueryInput, opts ...grpc.CallOption) (*TracesQueryOutput, error) InvokeLambda(ctx context.Context, in *InvokeLambdaInput, opts ...grpc.CallOption) (*InvokeLambdaOutput, error) RunLua(ctx context.Context, in *RunLuaInput, opts ...grpc.CallOption) (*RunLuaOutput, error) + RunPython(ctx context.Context, in *RunPythonInput, opts ...grpc.CallOption) (*RunPythonOutput, error) } type toolQueryClient struct { @@ -119,6 +121,16 @@ func (c *toolQueryClient) RunLua(ctx context.Context, in *RunLuaInput, opts ...g return out, nil } +func (c *toolQueryClient) RunPython(ctx context.Context, in *RunPythonInput, opts ...grpc.CallOption) (*RunPythonOutput, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(RunPythonOutput) + err := c.cc.Invoke(ctx, ToolQuery_RunPython_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + // ToolQueryServer is the server API for ToolQuery service. // All implementations must embed UnimplementedToolQueryServer // for forward compatibility. @@ -130,6 +142,7 @@ type ToolQueryServer interface { Traces(context.Context, *TracesQueryInput) (*TracesQueryOutput, error) InvokeLambda(context.Context, *InvokeLambdaInput) (*InvokeLambdaOutput, error) RunLua(context.Context, *RunLuaInput) (*RunLuaOutput, error) + RunPython(context.Context, *RunPythonInput) (*RunPythonOutput, error) mustEmbedUnimplementedToolQueryServer() } @@ -161,6 +174,9 @@ func (UnimplementedToolQueryServer) InvokeLambda(context.Context, *InvokeLambdaI func (UnimplementedToolQueryServer) RunLua(context.Context, *RunLuaInput) (*RunLuaOutput, error) { return nil, status.Error(codes.Unimplemented, "method RunLua not implemented") } +func (UnimplementedToolQueryServer) RunPython(context.Context, *RunPythonInput) (*RunPythonOutput, error) { + return nil, status.Error(codes.Unimplemented, "method RunPython not implemented") +} func (UnimplementedToolQueryServer) mustEmbedUnimplementedToolQueryServer() {} func (UnimplementedToolQueryServer) testEmbeddedByValue() {} @@ -308,6 +324,24 @@ func _ToolQuery_RunLua_Handler(srv interface{}, ctx context.Context, dec func(in return interceptor(ctx, in, info, handler) } +func _ToolQuery_RunPython_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(RunPythonInput) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(ToolQueryServer).RunPython(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: ToolQuery_RunPython_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(ToolQueryServer).RunPython(ctx, req.(*RunPythonInput)) + } + return interceptor(ctx, in, info, handler) +} + // ToolQuery_ServiceDesc is the grpc.ServiceDesc for ToolQuery service. // It's only intended for direct use with grpc.RegisterService, // and not to be introspected or modified (even as a copy) @@ -343,6 +377,10 @@ var ToolQuery_ServiceDesc = grpc.ServiceDesc{ MethodName: "RunLua", Handler: _ToolQuery_RunLua_Handler, }, + { + MethodName: "RunPython", + Handler: _ToolQuery_RunPython_Handler, + }, }, Streams: []grpc.StreamDesc{}, Metadata: "toolquery.proto", diff --git a/go/cloud-query/internal/server/server.go b/go/cloud-query/internal/server/server.go index a2560d3cab..7c018e1abb 100644 --- a/go/cloud-query/internal/server/server.go +++ b/go/cloud-query/internal/server/server.go @@ -5,6 +5,7 @@ import ( "crypto/tls" "fmt" "net" + "sync" "time" "google.golang.org/grpc" @@ -23,6 +24,7 @@ type Server struct { server *grpc.Server services []service.Service stopped chan struct{} // Channel to signal when server is stopped + stopOnce sync.Once } // Start initializes and starts the gRPC server @@ -49,21 +51,32 @@ func (in *Server) Start(ctx context.Context) error { // Stop gracefully shuts down the server func (in *Server) Stop() { - klog.Info("stopping gRPC server") - go func() { - in.server.GracefulStop() - close(in.stopped) - }() + in.stopOnce.Do(func() { + klog.Info("stopping gRPC server") + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + go func() { + in.server.GracefulStop() + close(in.stopped) + }() + + select { + case <-in.stopped: + klog.Info("gRPC server stopped gracefully") + case <-ctx.Done(): + klog.Info("timeout waiting for server to stop, forcing shutdown") + in.server.Stop() + } - // Wait for a graceful stop with a timeout - timeout := time.After(10 * time.Second) - select { - case <-in.stopped: - klog.Info("gRPC server stopped gracefully") - case <-timeout: - klog.Info("timeout waiting for server to stop, forcing shutdown") - in.server.Stop() - } + for _, registered := range in.services { + if closer, ok := registered.(service.Closer); ok { + if err := closer.Close(ctx); err != nil { + klog.ErrorS(err, "failed to close service") + } + } + } + }) } func (in *Server) register() { diff --git a/go/cloud-query/internal/service/service.go b/go/cloud-query/internal/service/service.go index 8c4b92a8e3..d55cd87da2 100644 --- a/go/cloud-query/internal/service/service.go +++ b/go/cloud-query/internal/service/service.go @@ -1,6 +1,8 @@ package service import ( + "context" + "google.golang.org/grpc" ) @@ -8,3 +10,9 @@ import ( type Service interface { Install(server *grpc.Server) } + +// Closer is implemented by services that own background processes or other +// resources which must be released during server shutdown. +type Closer interface { + Close(context.Context) error +} diff --git a/go/cloud-query/internal/service/toolquery.go b/go/cloud-query/internal/service/toolquery.go index 9fc32ed4b4..5a121e7698 100644 --- a/go/cloud-query/internal/service/toolquery.go +++ b/go/cloud-query/internal/service/toolquery.go @@ -16,11 +16,14 @@ import ( "github.com/pluralsh/console/go/cloud-query/internal/tools" lambdatools "github.com/pluralsh/console/go/cloud-query/internal/tools/lambda" luatools "github.com/pluralsh/console/go/cloud-query/internal/tools/lua" + pythontools "github.com/pluralsh/console/go/cloud-query/internal/tools/python" + pythongrpc "github.com/pluralsh/console/go/cloud-query/internal/tools/python/grpc" ) // ToolQueryService implements the toolquery.ToolQueryServer interface. type ToolQueryService struct { toolquery.UnimplementedToolQueryServer + python pythontools.Runner } // Install registers the ToolQuery service with the gRPC server. @@ -29,9 +32,20 @@ func (in *ToolQueryService) Install(server *grpc.Server) { toolquery.RegisterToolQueryServer(server, in) } -// NewToolQueryService creates a new instance of the ToolQuery server. -func NewToolQueryService() Service { - return &ToolQueryService{} +// NewToolQueryService creates a ToolQuery server and verifies its sandboxed +// Python worker pool before the gRPC server starts accepting requests. +func NewToolQueryService(ctx context.Context) (Service, error) { + runner, err := pythontools.NewRunner(ctx, pythontools.RunnerConfig{}) + if err != nil { + return nil, err + } + + return &ToolQueryService{python: runner}, nil +} + +// Close reaps all Python worker subprocesses. +func (in *ToolQueryService) Close(ctx context.Context) error { + return in.python.CloseContext(ctx) } func (in *ToolQueryService) Metrics(ctx context.Context, input *toolquery.MetricsQueryInput) (*toolquery.MetricsQueryOutput, error) { @@ -204,6 +218,28 @@ func (in *ToolQueryService) RunLua(ctx context.Context, input *toolquery.RunLuaI }, nil } +func (in *ToolQueryService) RunPython(ctx context.Context, input *toolquery.RunPythonInput) (*toolquery.RunPythonOutput, error) { + if input == nil { + return nil, status.Error(codes.InvalidArgument, "input is required") + } + output, err := in.python.Run(ctx, pythontools.RunInput{ + Script: input.GetScript(), + InputJSON: input.GetInputJson(), + }) + if err != nil { + if pythontools.CodeOf(err) == pythontools.Internal { + klog.ErrorS(err, "python worker failed") + } + + return nil, pythongrpc.StatusError(err) + } + + return &toolquery.RunPythonOutput{ + ResultJson: output.ResultJSON, + Stdout: output.Stdout, + }, nil +} + func (in *ToolQueryService) validateInput(connection *toolquery.ToolConnection, query string, timeRange *toolquery.TimeRange) error { if err := in.validateSearchInput(connection); err != nil { return err diff --git a/go/cloud-query/internal/service/toolquery_python_test.go b/go/cloud-query/internal/service/toolquery_python_test.go new file mode 100644 index 0000000000..83f5774c65 --- /dev/null +++ b/go/cloud-query/internal/service/toolquery_python_test.go @@ -0,0 +1,77 @@ +package service + +import ( + "context" + "errors" + "os" + "strings" + "testing" + + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + "github.com/pluralsh/console/go/cloud-query/internal/proto/toolquery" + pythontools "github.com/pluralsh/console/go/cloud-query/internal/tools/python" +) + +func TestMain(m *testing.M) { + if len(os.Args) == 2 && os.Args[1] == "python-worker" { + if err := pythontools.NewWorker().Run(os.Stdin, os.Stdout); err != nil { + os.Exit(1) + } + + os.Exit(0) + } + + os.Exit(m.Run()) +} + +type failingPythonRunner struct{ err error } + +func (r failingPythonRunner) Run(context.Context, pythontools.RunInput) (*pythontools.RunOutput, error) { + return nil, r.err +} + +func (failingPythonRunner) CloseContext(context.Context) error { return nil } + +func TestRunPythonValidatesInput(t *testing.T) { + server := &ToolQueryService{python: failingPythonRunner{err: errors.New("runner must not be called")}} + + if _, err := server.RunPython(context.Background(), nil); status.Code(err) != codes.InvalidArgument { + t.Fatalf("nil input code = %s, want %s", status.Code(err), codes.InvalidArgument) + } +} + +func TestRunPythonHidesInternalCause(t *testing.T) { + server := &ToolQueryService{python: failingPythonRunner{err: errors.New("private process detail")}} + + _, err := server.RunPython(context.Background(), &toolquery.RunPythonInput{Script: "output = {}"}) + if got := status.Convert(err).Message(); got != "python runtime failed" { + t.Fatalf("RunPython() message = %q, want safe summary", got) + } +} + +func TestRunPythonPreservesSandboxErrorMessage(t *testing.T) { + runner, err := pythontools.NewRunner(context.Background(), pythontools.RunnerConfig{ + Workers: 1, + QueueSize: 1, + MaxSuccessfulRunsBeforeRecycle: 1, + }) + if err != nil { + t.Fatal(err) + } + defer runner.CloseContext(context.Background()) + + server := &ToolQueryService{python: runner} + _, err = server.RunPython(context.Background(), &toolquery.RunPythonInput{ + Script: "import socket\nsocket.create_connection(('example.com', 443))\noutput = {'reachable': True}", + }) + + got := status.Convert(err) + if got.Code() != codes.FailedPrecondition { + t.Fatalf("status code = %s", got.Code()) + } + if !strings.Contains(strings.ToLower(got.Message()), "socket") { + t.Fatalf("status message = %q", got.Message()) + } +} diff --git a/go/cloud-query/internal/tools/python/README.md b/go/cloud-query/internal/tools/python/README.md new file mode 100644 index 0000000000..168738c000 --- /dev/null +++ b/go/cloud-query/internal/tools/python/README.md @@ -0,0 +1,89 @@ +# Sandboxed Python runner + +This package provides the runner used by the `RunPython` service RPC with +Monty's limited Python subset. It is not CPython. The parent starts four copies +of `cloud-query python-worker` and communicates with them over private +stdin/stdout pipes. Each worker loads gomonty's embedded native library and +creates a fresh REPL for every request. + +## Architecture + +The root package is a facade for runner and worker construction, stable error +codes, and safe public error messages. Its implementation is split into four +private modules: + +- `contract` owns execution types, fixed limits, validation, and errors. +- `protocol` owns strict version 1 framing and message validation. +- `worker` joins the request server to the Monty runtime. +- `pool` owns isolated subprocesses, queueing, recycling, and shutdown. + +Dependencies flow inward: `pool` and `worker` use `protocol`; those modules use +`contract`; the root package joins them. Each worker receives a replacement +environment containing only `TMPDIR=/tmp` by default. + +Errors have a stable code, a bounded public summary, and a private diagnostic +capped at 64 KiB. Use `PublicMessage` before returning an error to an untrusted +caller. Protocol data and process diagnostics must remain internal. + +## Execution contract + +Each run receives Python source and an optional JSON object. The object becomes +the global `input` dictionary. `output` starts as an empty dictionary and must +remain JSON serializable and object shaped. The response returns `output` and +the user script's standard output separately. + +Workers receive only `TMPDIR=/tmp`. The host exposes UTC clock callbacks for +`datetime.now()` and `date.today()` only; `datetime.now()` returns a naive UTC +value, and non-UTC timezone arguments are rejected. It supplies no filesystem, +environment, network, subprocess, shell, package, external function, or +name-lookup capability. + +Limits: + +- Source: 64 KiB +- Input and result JSON: 1 MiB each +- Captured stdout: 64 KiB +- Private diagnostics: 64 KiB +- Monty execution: 60 seconds +- Parent wall clock: 65 seconds or the caller's earlier deadline +- Monty-managed memory: 100 MiB +- Recursion: 100 frames +- Active workers: 4 +- Waiting requests: 16 by default +- Worker recycling: 10 successful requests + +Monty's memory limit does not bound total worker or pod RSS. The parent provides +the hard crash and wall-time limit by killing and replacing a worker process. + +## Private protocol + +Protocol version 1 uses a four-byte big-endian length followed by strict JSON. +Frames are bounded before allocation, unknown fields and trailing JSON values +are rejected, and only `health` and `run` request kinds are accepted. Every +response must repeat the request ID and kind. Error responses carry a stable +code, a bounded sanitized public summary, and a separate bounded private +diagnostic. The parent validates and sanitizes the public summary again before +returning it, while the private diagnostic remains internal. Malformed protocol +data terminates the worker. + +The parent admits requests to a bounded FIFO queue. Caller cancellation applies +while waiting and during execution. A worker is replaced after any execution +error, cancellation that reaches its exchange, protocol or transport failure, +or after ten successful responses. It does not run a health request between +successful executions. Shutdown stops admission, drains queued and active work +until the server shutdown context expires, then kills and reaps the remaining +workers. + +## Native runtime + +The Go module pins gomonty `v0.0.14`, which embeds glibc shared libraries for +Linux amd64 and arm64. That release builds against official Monty commit +`c9802b5f30d11fecf9f153feb1dfdab3abda070e`. The production image contains no +separate `monty` executable. + +Run focused validation with: + +```sh +cd go/cloud-query +go test -race ./internal/tools/python/... ./internal/service +``` diff --git a/go/cloud-query/internal/tools/python/config.go b/go/cloud-query/internal/tools/python/config.go new file mode 100644 index 0000000000..0a24a7930b --- /dev/null +++ b/go/cloud-query/internal/tools/python/config.go @@ -0,0 +1,78 @@ +package python + +import ( + "os" + + "github.com/samber/lo" + + "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/contract" +) + +const ( + defaultWorkers = 4 + defaultQueueSize = 16 + defaultRecycle = 10 +) + +var currentExecutable = os.Executable + +// WorkerProcessConfig configures the command used for each isolated worker. +// Environment replaces the inherited environment instead of extending it. +type WorkerProcessConfig struct { + Executable string + Arguments []string + Environment []string +} + +// RunnerConfig configures worker capacity, queueing, recycling, and process +// startup. Zero-valued capacity settings use the package defaults. +type RunnerConfig struct { + Workers int + QueueSize int + MaxSuccessfulRunsBeforeRecycle int + WorkerProcess WorkerProcessConfig +} + +func resolveRunnerConfig(overlay RunnerConfig) (RunnerConfig, error) { + config := RunnerConfig{ + Workers: defaultWorkers, + QueueSize: defaultQueueSize, + MaxSuccessfulRunsBeforeRecycle: defaultRecycle, + WorkerProcess: WorkerProcessConfig{ + Arguments: []string{"python-worker"}, + Environment: []string{"TMPDIR=/tmp"}, + }, + } + if overlay.Workers != 0 { + config.Workers = overlay.Workers + } + + if overlay.QueueSize != 0 { + config.QueueSize = overlay.QueueSize + } + + if overlay.MaxSuccessfulRunsBeforeRecycle != 0 { + config.MaxSuccessfulRunsBeforeRecycle = overlay.MaxSuccessfulRunsBeforeRecycle + } + + if overlay.WorkerProcess.Executable != "" { + config.WorkerProcess.Executable = overlay.WorkerProcess.Executable + } else { + executable, err := currentExecutable() + if err != nil { + return RunnerConfig{}, contract.UnavailableError("locating the python worker executable", err) + } + + config.WorkerProcess.Executable = executable + } + + if overlay.WorkerProcess.Arguments != nil { + config.WorkerProcess.Arguments = lo.Clone(overlay.WorkerProcess.Arguments) + } + + if overlay.WorkerProcess.Environment != nil { + config.WorkerProcess.Environment = lo.Clone(overlay.WorkerProcess.Environment) + } + + return config, nil +} diff --git a/go/cloud-query/internal/tools/python/config_test.go b/go/cloud-query/internal/tools/python/config_test.go new file mode 100644 index 0000000000..8d9a4cc5a7 --- /dev/null +++ b/go/cloud-query/internal/tools/python/config_test.go @@ -0,0 +1,95 @@ +package python + +import ( + "errors" + "testing" +) + +func TestResolveRunnerConfigDefaultsAndOverlays(t *testing.T) { + previous := currentExecutable + currentExecutable = func() (string, error) { return "/test/python", nil } + t.Cleanup(func() { currentExecutable = previous }) + + defaults, err := resolveRunnerConfig(RunnerConfig{}) + if err != nil { + t.Fatal(err) + } + + if defaults.Workers != 4 || defaults.QueueSize != 16 || defaults.MaxSuccessfulRunsBeforeRecycle != 10 { + t.Fatalf("defaults = %#v", defaults) + } + if defaults.WorkerProcess.Executable != "/test/python" { + t.Fatalf("executable = %q", defaults.WorkerProcess.Executable) + } + if got := defaults.WorkerProcess.Arguments; len(got) != 1 || got[0] != "python-worker" { + t.Fatalf("arguments = %#v", got) + } + if got := defaults.WorkerProcess.Environment; len(got) != 1 || got[0] != "TMPDIR=/tmp" { + t.Fatalf("worker defaults = %#v", defaults.WorkerProcess) + } + + arguments := []string{"child", "worker"} + environment := []string{"TMPDIR=/sandbox"} + resolved, err := resolveRunnerConfig(RunnerConfig{ + Workers: 2, + QueueSize: 3, + MaxSuccessfulRunsBeforeRecycle: 4, + WorkerProcess: WorkerProcessConfig{ + Executable: "/custom/python", + Arguments: arguments, + Environment: environment, + }, + }) + if err != nil { + t.Fatal(err) + } + + arguments[0] = "changed" + environment[0] = "changed=value" + if resolved.Workers != 2 || resolved.QueueSize != 3 || resolved.MaxSuccessfulRunsBeforeRecycle != 4 { + t.Fatalf("numeric overlay = %#v", resolved) + } + if resolved.WorkerProcess.Executable != "/custom/python" { + t.Fatalf("executable overlay = %q", resolved.WorkerProcess.Executable) + } + if resolved.WorkerProcess.Arguments[0] != "child" || + resolved.WorkerProcess.Environment[0] != "TMPDIR=/sandbox" { + t.Fatalf("overlay = %#v", resolved) + } + + empty, err := resolveRunnerConfig(RunnerConfig{WorkerProcess: WorkerProcessConfig{Arguments: []string{}, Environment: []string{}}}) + if err != nil || empty.WorkerProcess.Arguments == nil || empty.WorkerProcess.Environment == nil { + t.Fatalf("explicit empty slices = %#v, %v", empty.WorkerProcess, err) + } + + if len(empty.WorkerProcess.Arguments) != 0 || len(empty.WorkerProcess.Environment) != 0 { + t.Fatalf("explicit empty slices = %#v, %v", empty.WorkerProcess, err) + } +} + +func TestNewRunnerRejectsInvalidValues(t *testing.T) { + previous := currentExecutable + currentExecutable = func() (string, error) { return "/test/python", nil } + t.Cleanup(func() { currentExecutable = previous }) + for _, config := range []RunnerConfig{ + {Workers: -1}, + {QueueSize: -1}, + {MaxSuccessfulRunsBeforeRecycle: -1}, + {WorkerProcess: WorkerProcessConfig{Executable: " "}}, + {WorkerProcess: WorkerProcessConfig{Environment: []string{"missing-equals"}}}, + } { + if _, err := NewRunner(t.Context(), config); CodeOf(err) != InvalidArgument { + t.Fatalf("config %#v: %v", config, err) + } + } + + currentExecutable = func() (string, error) { return "", errors.New("missing executable") } + if _, err := resolveRunnerConfig(RunnerConfig{}); CodeOf(err) != Unavailable { + t.Fatalf("executable error = %v", err) + } + + custom, err := resolveRunnerConfig(RunnerConfig{WorkerProcess: WorkerProcessConfig{Executable: "/custom/python"}}) + if err != nil || custom.WorkerProcess.Executable != "/custom/python" { + t.Fatalf("custom executable with unavailable default = %#v, %v", custom.WorkerProcess, err) + } +} diff --git a/go/cloud-query/internal/tools/python/errors.go b/go/cloud-query/internal/tools/python/errors.go new file mode 100644 index 0000000000..fe15f2d1cd --- /dev/null +++ b/go/cloud-query/internal/tools/python/errors.go @@ -0,0 +1,10 @@ +package python + +import "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/contract" + +// CodeOf returns err's stable runner error code, or Internal for unclassified errors. +func CodeOf(err error) Code { return contract.CodeOf(err) } + +// PublicMessage returns the safe summary suitable for untrusted callers. It +// never includes private worker or process diagnostics. +func PublicMessage(err error) string { return contract.PublicMessage(err) } diff --git a/go/cloud-query/internal/tools/python/grpc/status.go b/go/cloud-query/internal/tools/python/grpc/status.go new file mode 100644 index 0000000000..da8d74c76d --- /dev/null +++ b/go/cloud-query/internal/tools/python/grpc/status.go @@ -0,0 +1,32 @@ +package grpc + +import ( + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + "github.com/pluralsh/console/go/cloud-query/internal/tools/python" +) + +// StatusError translates a runner error into an RPC-safe gRPC status error. +func StatusError(err error) error { + return status.Error(statusCode(python.CodeOf(err)), python.PublicMessage(err)) +} + +func statusCode(code python.Code) codes.Code { + switch code { + case python.InvalidArgument: + return codes.InvalidArgument + case python.FailedPrecondition: + return codes.FailedPrecondition + case python.Canceled: + return codes.Canceled + case python.DeadlineExceeded: + return codes.DeadlineExceeded + case python.ResourceExhausted: + return codes.ResourceExhausted + case python.Unavailable: + return codes.Unavailable + default: + return codes.Internal + } +} diff --git a/go/cloud-query/internal/tools/python/grpc/status_test.go b/go/cloud-query/internal/tools/python/grpc/status_test.go new file mode 100644 index 0000000000..e917925235 --- /dev/null +++ b/go/cloud-query/internal/tools/python/grpc/status_test.go @@ -0,0 +1,43 @@ +package grpc + +import ( + "errors" + "testing" + + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + "github.com/pluralsh/console/go/cloud-query/internal/tools/python" + "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/contract" +) + +func TestStatusCode(t *testing.T) { + for code, want := range map[python.Code]codes.Code{ + python.InvalidArgument: codes.InvalidArgument, + python.FailedPrecondition: codes.FailedPrecondition, + python.Canceled: codes.Canceled, + python.DeadlineExceeded: codes.DeadlineExceeded, + python.ResourceExhausted: codes.ResourceExhausted, + python.Unavailable: codes.Unavailable, + python.Internal: codes.Internal, + } { + if got := statusCode(code); got != want { + t.Errorf("statusCode(%s) = %s, want %s", code, got, want) + } + } +} + +func TestStatusErrorPreservesPublicMessage(t *testing.T) { + err := contract.Failed( + "python ImportError: socket is unavailable", + errors.New("private sandbox diagnostic"), + ) + + got := status.Convert(StatusError(err)) + if got.Code() != codes.FailedPrecondition { + t.Fatalf("status code = %s", got.Code()) + } + if got.Message() != "python ImportError: socket is unavailable" { + t.Fatalf("status message = %q", got.Message()) + } +} diff --git a/go/cloud-query/internal/tools/python/internal/contract/errors.go b/go/cloud-query/internal/tools/python/internal/contract/errors.go new file mode 100644 index 0000000000..7561f2f626 --- /dev/null +++ b/go/cloud-query/internal/tools/python/internal/contract/errors.go @@ -0,0 +1,139 @@ +package contract + +import ( + "context" + "errors" + "strings" +) + +const maxPublicRunes = 512 +const detailTruncationMarker = "\n[python diagnostic truncated]" + +type runError struct { + code Code + summary string + cause error +} + +func (e *runError) Error() string { + if e.cause == nil { + return e.summary + } + return e.summary + ": " + e.cause.Error() +} + +func (e *runError) Unwrap() error { return e.cause } + +// CodeOf returns err's stable code, or Internal for an unclassified error. +func CodeOf(err error) Code { + var run *runError + if errors.As(err, &run) { + return run.code + } + return Internal +} + +// PublicMessage returns err's safe public summary without its diagnostic cause. +func PublicMessage(err error) string { + var run *runError + if errors.As(err, &run) { + return run.summary + } + return "python runtime failed" +} + +// Detail returns err's bounded diagnostic for private logs and protocol traffic. +func Detail(err error) string { return truncateDetail(err.Error()) } + +// New creates an error with a stable code, sanitized public summary, and +// bounded private diagnostic cause. +func New(code Code, summary string, cause error) error { + return &runError{ + code: code, + summary: public(summary), + cause: boundedCause(cause), + } +} + +// Invalid creates an InvalidArgument error. +func Invalid(summary string, cause error) error { + return New(InvalidArgument, summary, cause) +} + +// Failed creates a FailedPrecondition error. +func Failed(summary string, cause error) error { + return New(FailedPrecondition, summary, cause) +} + +// CanceledError creates a Canceled error with the standard public summary. +func CanceledError(cause error) error { + return New(Canceled, "python execution was canceled", cause) +} + +// Deadline creates a DeadlineExceeded error with the standard public summary. +func Deadline(cause error) error { + return New(DeadlineExceeded, "python execution timed out", cause) +} + +// Exhausted creates a ResourceExhausted error. +func Exhausted(summary string, cause error) error { + return New(ResourceExhausted, summary, cause) +} + +// UnavailableError creates an Unavailable error. +func UnavailableError(summary string, cause error) error { + return New(Unavailable, summary, cause) +} + +// InternalError creates an Internal error with the standard public summary. +func InternalError(cause error) error { + return New(Internal, "python runtime failed", cause) +} + +// ContextError maps a context cancellation or deadline error to a runner code. +func ContextError(err error) error { + if errors.Is(err, context.Canceled) { + return CanceledError(err) + } + return Deadline(err) +} + +func public(summary string) string { + summary = strings.Join(strings.Fields(strings.ToValidUTF8(summary, "")), " ") + if summary == "" { + return "python runtime failed" + } + runes := []rune(summary) + if len(runes) > maxPublicRunes { + return string(runes[:maxPublicRunes-1]) + "…" + } + return summary +} + +func boundedCause(cause error) error { + if cause == nil { + return nil + } + if len(cause.Error()) <= MaxDetailBytes { + return cause + } + return truncatedCauseError{detail: truncateDetail(cause.Error()), cause: cause} +} + +func truncateDetail(detail string) string { + if len(detail) <= MaxDetailBytes { + return detail + } + limit := MaxDetailBytes - len(detailTruncationMarker) + + return detail[:limit] + detailTruncationMarker +} + +type truncatedCauseError struct { + detail string + cause error +} + +func (e truncatedCauseError) Error() string { return e.detail } + +func (e truncatedCauseError) Unwrap() error { return e.cause } diff --git a/go/cloud-query/internal/tools/python/internal/contract/errors_test.go b/go/cloud-query/internal/tools/python/internal/contract/errors_test.go new file mode 100644 index 0000000000..4ae97e5e5f --- /dev/null +++ b/go/cloud-query/internal/tools/python/internal/contract/errors_test.go @@ -0,0 +1,18 @@ +package contract + +import ( + "errors" + "strings" + "testing" +) + +func TestErrorKeepsCauseAndBoundsDisclosure(t *testing.T) { + cause := errors.New(strings.Repeat("x", MaxDetailBytes+1)) + err := Invalid(strings.Repeat("summary ", 100), cause) + if CodeOf(err) != InvalidArgument || len([]rune(PublicMessage(err))) > 512 { + t.Fatalf("unexpected classified error: %q", PublicMessage(err)) + } + if !errors.Is(err, cause) || !strings.HasSuffix(Detail(err), detailTruncationMarker) { + t.Fatal("cause or detail truncation was lost") + } +} diff --git a/go/cloud-query/internal/tools/python/internal/contract/types.go b/go/cloud-query/internal/tools/python/internal/contract/types.go new file mode 100644 index 0000000000..f492e734e0 --- /dev/null +++ b/go/cloud-query/internal/tools/python/internal/contract/types.go @@ -0,0 +1,48 @@ +package contract + +const ( + // MaxSourceBytes bounds Python source accepted for one execution. + MaxSourceBytes = 64 << 10 + // MaxInputBytes bounds the JSON input accepted for one execution. + MaxInputBytes = 1 << 20 + // MaxResultBytes bounds the JSON result returned for one execution. + MaxResultBytes = 1 << 20 + // MaxStdoutBytes bounds stdout captured for one execution. + MaxStdoutBytes = 64 << 10 + // MaxDetailBytes bounds private diagnostics retained or sent on the protocol. + MaxDetailBytes = 64 << 10 + // MaxPublicMessageBytes bounds a sanitized error summary sent by a worker. + MaxPublicMessageBytes = 4 << 10 +) + +// RunInput contains validated source and the JSON object supplied as input. +type RunInput struct { + Script string + InputJSON string +} + +// RunOutput contains the JSON object assigned to output and captured stdout. +type RunOutput struct { + ResultJSON string + Stdout string +} + +// Code identifies a stable class of execution failure. +type Code string + +const ( + // InvalidArgument reports invalid source, input, or output values. + InvalidArgument Code = "invalid_argument" + // FailedPrecondition reports a Python execution failure after validation. + FailedPrecondition Code = "failed_precondition" + // Canceled reports cancellation by the caller. + Canceled Code = "canceled" + // DeadlineExceeded reports an execution deadline that elapsed. + DeadlineExceeded Code = "deadline_exceeded" + // ResourceExhausted reports a configured capacity or resource limit. + ResourceExhausted Code = "resource_exhausted" + // Unavailable reports that no runtime worker can accept the request. + Unavailable Code = "unavailable" + // Internal reports an unclassified runner or worker failure. + Internal Code = "internal" +) diff --git a/go/cloud-query/internal/tools/python/internal/contract/validation.go b/go/cloud-query/internal/tools/python/internal/contract/validation.go new file mode 100644 index 0000000000..48b2de98cf --- /dev/null +++ b/go/cloud-query/internal/tools/python/internal/contract/validation.go @@ -0,0 +1,44 @@ +package contract + +import ( + "encoding/json" + "strings" +) + +// NormalizeInput returns raw after verifying it is a bounded JSON object. An +// empty input becomes an empty JSON object. +func NormalizeInput(raw string) (string, error) { + if strings.TrimSpace(raw) == "" { + return "{}", nil + } + if len(raw) > MaxInputBytes { + return "", Invalid("input exceeds the input limit", nil) + } + var object map[string]json.RawMessage + if err := json.Unmarshal([]byte(raw), &object); err != nil || object == nil { + return "", Invalid("input_json must be a JSON object", err) + } + return raw, nil +} + +// ValidateRun validates source and normalizes the JSON object supplied as input. +func ValidateRun(input RunInput) (RunInput, error) { + if strings.TrimSpace(input.Script) == "" { + return RunInput{}, Invalid("script is required", nil) + } + if len(input.Script) > MaxSourceBytes { + return RunInput{}, Invalid("script exceeds the source limit", nil) + } + normalized, err := NormalizeInput(input.InputJSON) + if err != nil { + return RunInput{}, err + } + input.InputJSON = normalized + return input, nil +} + +// IsJSONObject reports whether raw contains one JSON object and no other value. +func IsJSONObject(raw string) bool { + var object map[string]json.RawMessage + return json.Unmarshal([]byte(raw), &object) == nil && object != nil +} diff --git a/go/cloud-query/internal/tools/python/internal/pool/config.go b/go/cloud-query/internal/tools/python/internal/pool/config.go new file mode 100644 index 0000000000..249f17ee32 --- /dev/null +++ b/go/cloud-query/internal/tools/python/internal/pool/config.go @@ -0,0 +1,44 @@ +package pool + +import ( + "strings" + + "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/contract" +) + +// ProcessConfig configures the command used for an isolated worker process. +// Environment replaces the parent process environment. +type ProcessConfig struct { + Executable string + Arguments []string + Environment []string +} + +// Config configures a Runner's worker capacity, queue, recycling, and process. +type Config struct { + Workers int + QueueSize int + MaxSuccessfulRunsBeforeRecycle int + Process ProcessConfig +} + +func validate(config Config) error { + if config.Workers <= 0 { + return contract.Invalid("python runner workers must be greater than zero", nil) + } + if config.QueueSize < 0 { + return contract.Invalid("python runner queue size must not be negative", nil) + } + if config.MaxSuccessfulRunsBeforeRecycle <= 0 { + return contract.Invalid("python runner recycle limit must be greater than zero", nil) + } + if strings.TrimSpace(config.Process.Executable) == "" { + return contract.Invalid("python worker executable is required", nil) + } + for _, entry := range config.Process.Environment { + if !strings.Contains(entry, "=") { + return contract.Invalid("python worker environment entries must use NAME=value", nil) + } + } + return nil +} diff --git a/go/cloud-query/internal/tools/python/internal/pool/pool.go b/go/cloud-query/internal/tools/python/internal/pool/pool.go new file mode 100644 index 0000000000..3981fd0f39 --- /dev/null +++ b/go/cloud-query/internal/tools/python/internal/pool/pool.go @@ -0,0 +1,408 @@ +package pool + +import ( + "context" + "fmt" + "sync" + "sync/atomic" + "time" + + "k8s.io/klog/v2" + + "github.com/pluralsh/console/go/cloud-query/internal/log" + "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/contract" + "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/protocol" +) + +const wallTimeout = 65 * time.Second + +// Runner owns a bounded set of isolated worker processes. It replaces workers +// after failures and recycles them after the configured number of successes. +type Runner struct { + config Config + factory processFactory + jobs chan job + mu sync.Mutex + closed bool + processes map[process]struct{} + wg sync.WaitGroup + closeOnce sync.Once + done chan struct{} + requestIDs atomic.Uint64 + workerIDs atomic.Uint64 + retry func() <-chan time.Time +} + +type job struct { + requestID string + queuedAt time.Time + ctx context.Context + input contract.RunInput + response chan result +} + +type result struct { + output *contract.RunOutput + err error +} + +// New starts config.Workers health-checked worker processes. If startup fails, +// it shuts down workers already started before returning the error. +func New(ctx context.Context, config Config) (*Runner, error) { + return newRunner(ctx, config, newExecProcess) +} + +func newRunner(ctx context.Context, config Config, factory processFactory) (*Runner, error) { + if err := validate(config); err != nil { + return nil, err + } + + klog.V(log.LogLevelInfo).InfoS( + "initializing python worker pool", + "workers", config.Workers, + "queue_size", config.QueueSize, + "recycle_limit", config.MaxSuccessfulRunsBeforeRecycle, + ) + + runner := &Runner{ + config: config, + factory: factory, + jobs: make(chan job, config.QueueSize), + processes: map[process]struct{}{}, + done: make(chan struct{}), + retry: func() <-chan time.Time { + return time.After(time.Second) + }, + } + + for range config.Workers { + if err := runner.start(ctx); err != nil { + _ = runner.CloseContext(context.Background()) + return nil, contract.UnavailableError("python runtime is unavailable", err) + } + } + + klog.V(log.LogLevelInfo).InfoS("python worker pool is ready", "workers", config.Workers) + + return runner, nil +} + +func (r *Runner) start(ctx context.Context) error { + ctx, cancel := context.WithTimeout(ctx, wallTimeout) + defer cancel() + + // Worker IDs identify parent-side process lifecycles without exposing command details. + workerID := r.workerIDs.Add(1) + klog.V(log.LogLevelVerbose).InfoS("spawning python worker", "worker_id", workerID) + process, err := r.factory(r.config.Process) + if err != nil { + logWorkerFailure("failed to start python worker process", workerID, err) + return err + } + klog.V(log.LogLevelVerbose).InfoS("python worker process started", "worker_id", workerID) + + request := protocol.Request{ + Version: protocol.Version, + Kind: protocol.Health, + ID: fmt.Sprintf("%d", r.requestIDs.Add(1)), + } + + response, err := process.exchange(ctx, request) + if err != nil || response.Error != nil { + process.stop() + if err != nil { + logWorkerFailure("python worker health check failed", workerID, err) + return err + } + err = protocol.Error(response) + logWorkerFailure("python worker health check failed", workerID, err) + return err + } + klog.V(log.LogLevelVerbose).InfoS("python worker health check passed", "worker_id", workerID) + + r.mu.Lock() + if r.closed { + r.mu.Unlock() + process.stop() + return contract.UnavailableError("python runtime is unavailable", nil) + } + + r.processes[process] = struct{}{} + r.wg.Add(1) + r.mu.Unlock() + + go r.serve(process, workerID) + return nil +} + +// Run validates and enqueues input, then waits for its worker response or ctx. +// It rejects new work after shutdown begins or when the queue is full. +func (r *Runner) Run(ctx context.Context, input contract.RunInput) (*contract.RunOutput, error) { + input, err := contract.ValidateRun(input) + if err != nil { + return nil, err + } + + if err := ctx.Err(); err != nil { + return nil, contract.ContextError(err) + } + + job := job{ + requestID: fmt.Sprintf("%d", r.requestIDs.Add(1)), + queuedAt: time.Now(), + ctx: ctx, + input: input, + response: make(chan result, 1), + } + + r.mu.Lock() + if r.closed || len(r.processes) == 0 { + r.mu.Unlock() + return nil, contract.UnavailableError("python runtime is unavailable", nil) + } + + select { + case r.jobs <- job: + klog.V(log.LogLevelExtended).InfoS( + "python request queued", + "request_id", job.requestID, + "queued_jobs", len(r.jobs), + "queue_capacity", cap(r.jobs), + ) + r.mu.Unlock() + case <-ctx.Done(): + r.mu.Unlock() + return nil, contract.ContextError(ctx.Err()) + default: + r.mu.Unlock() + return nil, contract.Exhausted("python queue is full", nil) + } + + select { + case answer := <-job.response: + return answer.output, answer.err + case <-ctx.Done(): + return nil, contract.ContextError(ctx.Err()) + } +} + +func (r *Runner) serve(process process, workerID uint64) { + retirement := "shutdown" + var retirementErr error + defer func() { + process.stop() + + r.mu.Lock() + delete(r.processes, process) + closed := r.closed + r.mu.Unlock() + + if closed { + retirement = "shutdown" + } + + fields := []any{"worker_id", workerID, "reason", retirement} + if retirementErr != nil { + fields = append(fields, + "error_code", contract.CodeOf(retirementErr), + "error", contract.PublicMessage(retirementErr), + ) + } + + klog.V(log.LogLevelExtended).InfoS("retiring python worker", fields...) + if !closed { + r.replace() + } + + r.wg.Done() + }() + + successes := 0 + for job := range r.jobs { + klog.V(log.LogLevelExtended).InfoS( + "python request taken from queue", + "request_id", job.requestID, + "worker_id", workerID, + "queued_for", time.Since(job.queuedAt), + "queued_jobs", len(r.jobs), + ) + + if err := job.ctx.Err(); err != nil { + klog.V(log.LogLevelExtended).InfoS( + "python request canceled before execution", + "request_id", job.requestID, + "worker_id", workerID, + "error_code", contract.CodeOf(contract.ContextError(err)), + ) + job.response <- result{err: contract.ContextError(err)} + continue + } + + ctx, cancel := context.WithTimeout(job.ctx, wallTimeout) + request := protocol.Request{ + Version: protocol.Version, + Kind: protocol.Run, + ID: job.requestID, + Script: job.input.Script, + InputJSON: job.input.InputJSON, + } + + startedAt := time.Now() + klog.V(log.LogLevelExtended).InfoS( + "python request execution started", + "request_id", job.requestID, + "worker_id", workerID, + ) + + response, err := process.exchange(ctx, request) + cancel() + if err != nil { + logExecutionFinished(job.requestID, workerID, startedAt, err) + job.response <- result{err: err} + retirement = retirementReason(err) + retirementErr = err + return + } + if response.Error != nil { + err := protocol.Error(response) + logExecutionFinished(job.requestID, workerID, startedAt, err) + job.response <- result{err: err} + retirement = "remote_python_error" + retirementErr = err + return + } + + logExecutionFinished(job.requestID, workerID, startedAt, nil) + job.response <- result{output: &contract.RunOutput{ + ResultJSON: response.ResultJSON, + Stdout: response.Stdout, + }} + + successes++ + if successes >= r.config.MaxSuccessfulRunsBeforeRecycle { + retirement = "successful_recycle_limit" + return + } + } +} + +func logExecutionFinished(requestID string, workerID uint64, startedAt time.Time, err error) { + status := "success" + if err != nil { + status = "error" + } + + fields := []any{ + "request_id", requestID, + "worker_id", workerID, + "duration", time.Since(startedAt), + "status", status, + } + if err != nil { + fields = append(fields, "error_code", contract.CodeOf(err)) + } + + klog.V(log.LogLevelExtended).InfoS("python request execution finished", fields...) +} + +func (r *Runner) replace() { + for { + r.mu.Lock() + closed := r.closed + r.mu.Unlock() + + if closed { + return + } + klog.V(log.LogLevelExtended).InfoS("attempting python worker replacement") + if err := r.start(context.Background()); err == nil { + klog.V(log.LogLevelExtended).InfoS("python worker replacement succeeded") + return + } + + select { + case <-r.done: + return + case <-r.retry(): + } + } +} + +// CloseContext stops admitting jobs, waits for worker shutdown, and kills +// remaining processes if ctx expires. It is safe to call more than once. +func (r *Runner) CloseContext(ctx context.Context) error { + r.closeOnce.Do(func() { + r.mu.Lock() + r.closed = true + processCount := len(r.processes) + queuedJobs := len(r.jobs) + close(r.jobs) + r.mu.Unlock() + klog.V(log.LogLevelInfo).InfoS( + "starting python worker pool shutdown", + "processes", processCount, + "queued_jobs", queuedJobs, + ) + go func() { + r.wg.Wait() + klog.V(log.LogLevelInfo).InfoS("python worker pool shutdown completed") + close(r.done) + }() + }) + + select { + case <-r.done: + return nil + case <-ctx.Done(): + r.mu.Lock() + processes := make([]process, 0, len(r.processes)) + for process := range r.processes { + processes = append(processes, process) + } + r.mu.Unlock() + + for _, process := range processes { + process.stop() + } + + drainedJobs := 0 + for draining := true; draining; { + select { + case job, ok := <-r.jobs: + if !ok { + draining = false + continue + } + job.response <- result{err: contract.UnavailableError("python runtime is unavailable", nil)} + drainedJobs++ + default: + draining = false + } + } + klog.ErrorS( + nil, + "python worker pool shutdown deadline exceeded; forcing teardown", + "processes", len(processes), + "queued_jobs", drainedJobs, + ) + return ctx.Err() + } +} + +func retirementReason(err error) string { + if contract.CodeOf(err) == contract.Canceled { + return "cancellation" + } + return "execution_or_transport_failure" +} + +func logWorkerFailure(message string, workerID uint64, err error) { + fields := []any{ + "error_code", contract.CodeOf(err), + "error", contract.PublicMessage(err), + } + if workerID != 0 { + fields = append([]any{"worker_id", workerID}, fields...) + } + klog.ErrorS(nil, message, fields...) +} diff --git a/go/cloud-query/internal/tools/python/internal/pool/pool_test.go b/go/cloud-query/internal/tools/python/internal/pool/pool_test.go new file mode 100644 index 0000000000..869c3d1969 --- /dev/null +++ b/go/cloud-query/internal/tools/python/internal/pool/pool_test.go @@ -0,0 +1,290 @@ +package pool + +import ( + "context" + "errors" + "fmt" + "runtime" + "sync" + "testing" + "time" + + "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/contract" + "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/protocol" +) + +type fakeProcess struct { + id int + stopped bool + mu sync.Mutex + seen []string + block chan struct{} +} + +func (p *fakeProcess) exchange(ctx context.Context, request protocol.Request) (protocol.Response, error) { + p.mu.Lock() + p.seen = append(p.seen, request.ID) + p.mu.Unlock() + if request.Kind == protocol.Health { + return protocol.Response{Version: protocol.Version, Kind: request.Kind, ID: request.ID}, nil + } + if request.Script == "error" { + return protocol.Response{ + Version: protocol.Version, + Kind: request.Kind, + ID: request.ID, + Error: &protocol.WireError{ + Code: contract.FailedPrecondition, + PublicMessage: "python ValueError: bad value", + Detail: "private", + }, + }, nil + } + if request.Script == "block" { + select { + case <-p.block: + case <-ctx.Done(): + return protocol.Response{}, contract.ContextError(ctx.Err()) + } + } + return protocol.Response{ + Version: protocol.Version, + Kind: request.Kind, + ID: request.ID, + ResultJSON: fmt.Sprintf(`{"worker":%d}`, p.id), + }, nil +} + +func (p *fakeProcess) stop() { + p.mu.Lock() + p.stopped = true + p.mu.Unlock() +} + +type fakeFactory struct { + mu sync.Mutex + processes []*fakeProcess + failAt int +} + +type controlledProcess struct { + started chan struct{} + stopped chan struct{} + once sync.Once +} + +func (p *controlledProcess) exchange(ctx context.Context, request protocol.Request) (protocol.Response, error) { + if request.Kind == protocol.Health { + return protocol.Response{Version: protocol.Version, Kind: request.Kind, ID: request.ID}, nil + } + select { + case <-p.started: + default: + close(p.started) + } + select { + case <-p.stopped: + return protocol.Response{}, contract.InternalError(errors.New("process stopped")) + case <-ctx.Done(): + return protocol.Response{}, contract.ContextError(ctx.Err()) + } +} + +func (p *controlledProcess) stop() { + p.once.Do(func() { close(p.stopped) }) +} + +func (f *fakeFactory) count() int { + f.mu.Lock() + defer f.mu.Unlock() + + return len(f.processes) +} + +func (f *fakeFactory) new(ProcessConfig) (process, error) { + f.mu.Lock() + defer f.mu.Unlock() + if f.failAt > 0 && len(f.processes)+1 == f.failAt { + return nil, errors.New("start failed") + } + p := &fakeProcess{id: len(f.processes) + 1, block: make(chan struct{})} + f.processes = append(f.processes, p) + return p, nil +} + +func testConfig() Config { + return Config{ + Workers: 1, + QueueSize: 1, + MaxSuccessfulRunsBeforeRecycle: 2, + Process: ProcessConfig{ + Executable: "fake", + Arguments: []string{"python-worker"}, + Environment: []string{"TMPDIR=/tmp"}, + }, + } +} + +func TestWallTimeoutAllowsMontyExecutionLimit(t *testing.T) { + if wallTimeout != 65*time.Second { + t.Fatalf("wall timeout = %s, want 65s", wallTimeout) + } +} + +func TestPoolStartupRecycleAndRemoteFailure(t *testing.T) { + factory := &fakeFactory{} + runner, err := newRunner(context.Background(), testConfig(), factory.new) + if err != nil { + t.Fatal(err) + } + defer runner.CloseContext(context.Background()) + + first, err := runner.Run(context.Background(), contract.RunInput{Script: "ok"}) + if err != nil { + t.Fatal(err) + } + if _, err := runner.Run(context.Background(), contract.RunInput{Script: "ok"}); err != nil { + t.Fatal(err) + } + deadline := time.Now().Add(time.Second) + for factory.count() < 2 && time.Now().Before(deadline) { + time.Sleep(time.Millisecond) + } + third, err := runner.Run(context.Background(), contract.RunInput{Script: "ok"}) + if err != nil || third.ResultJSON == first.ResultJSON { + t.Fatalf("recycle: %#v %v", third, err) + } + if _, err := runner.Run(context.Background(), contract.RunInput{Script: "error"}); contract.CodeOf(err) != contract.FailedPrecondition { + t.Fatalf("remote error: %v", err) + } +} + +func TestPoolFailsEagerStartupAndRejectsAfterClose(t *testing.T) { + factory := &fakeFactory{failAt: 2} + config := Config{ + Workers: 2, + QueueSize: 1, + MaxSuccessfulRunsBeforeRecycle: 1, + Process: ProcessConfig{ + Executable: "fake", + Arguments: []string{"python-worker"}, + Environment: []string{"TMPDIR=/tmp"}, + }, + } + if _, err := newRunner(context.Background(), config, factory.new); contract.CodeOf(err) != contract.Unavailable { + t.Fatalf("startup: %v", err) + } + runner, err := newRunner(context.Background(), testConfig(), (&fakeFactory{}).new) + if err != nil { + t.Fatal(err) + } + if err := runner.CloseContext(context.Background()); err != nil { + t.Fatal(err) + } + if _, err := runner.Run(context.Background(), contract.RunInput{Script: "ok"}); contract.CodeOf(err) != contract.Unavailable { + t.Fatalf("admission: %v", err) + } +} + +func TestRunCancellationRetiresActiveProcess(t *testing.T) { + proc := &controlledProcess{started: make(chan struct{}), stopped: make(chan struct{})} + runner, err := newRunner(context.Background(), testConfig(), func(ProcessConfig) (process, error) { return proc, nil }) + if err != nil { + t.Fatal(err) + } + defer runner.CloseContext(context.Background()) + ctx, cancel := context.WithCancel(context.Background()) + result := make(chan error, 1) + go func() { _, err := runner.Run(ctx, contract.RunInput{Script: "block"}); result <- err }() + <-proc.started + cancel() + if err := <-result; contract.CodeOf(err) != contract.Canceled { + t.Fatalf("cancellation = %v", err) + } + select { + case <-proc.stopped: + case <-time.After(time.Second): + t.Fatal("active process was not stopped") + } +} + +func TestCloseDeadlineStopsActiveProcessAndDrainsQueue(t *testing.T) { + proc := &controlledProcess{started: make(chan struct{}), stopped: make(chan struct{})} + runner, err := newRunner(context.Background(), testConfig(), func(ProcessConfig) (process, error) { return proc, nil }) + if err != nil { + t.Fatal(err) + } + active := make(chan error, 1) + go func() { _, err := runner.Run(context.Background(), contract.RunInput{Script: "block"}); active <- err }() + <-proc.started + queued := make(chan result, 1) + runner.jobs <- job{ctx: context.Background(), input: contract.RunInput{Script: "queued"}, response: queued} + deadline, cancel := context.WithTimeout(context.Background(), time.Millisecond) + defer cancel() + if err := runner.CloseContext(deadline); !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("close = %v", err) + } + if answer := <-queued; contract.CodeOf(answer.err) != contract.Unavailable { + t.Fatalf("queued = %v", answer.err) + } + if err := <-active; contract.CodeOf(err) != contract.Internal { + t.Fatalf("active = %v", err) + } + if _, err := runner.Run(context.Background(), contract.RunInput{Script: "later"}); contract.CodeOf(err) != contract.Unavailable { + t.Fatalf("admission = %v", err) + } +} + +func TestReplacementRetriesUntilCapacityRecovers(t *testing.T) { + var mu sync.Mutex + attempts := 0 + processes := make(chan *fakeProcess, 4) + factory := func(ProcessConfig) (process, error) { + mu.Lock() + defer mu.Unlock() + attempts++ + if attempts > 1 && attempts < 4 { + return nil, errors.New("replacement failed") + } + + process := &fakeProcess{id: attempts, block: make(chan struct{})} + processes <- process + return process, nil + } + + runner, err := newRunner(context.Background(), testConfig(), factory) + if err != nil { + t.Fatal(err) + } + defer runner.CloseContext(context.Background()) + <-processes + runner.retry = func() <-chan time.Time { ready := make(chan time.Time); close(ready); return ready } + if _, err := runner.Run(context.Background(), contract.RunInput{Script: "error"}); contract.CodeOf(err) != contract.FailedPrecondition { + t.Fatalf("run = %v", err) + } + <-processes + waitForCapacity(t, runner) + if _, err := runner.Run(context.Background(), contract.RunInput{Script: "ok"}); err != nil { + t.Fatalf("recovered = %v", err) + } + mu.Lock() + got := attempts + mu.Unlock() + if got != 4 { + t.Fatalf("replacement attempts = %d, want 4", got) + } +} + +func waitForCapacity(t *testing.T, runner *Runner) { + t.Helper() + for range 100_000 { + runner.mu.Lock() + ready := len(runner.processes) > 0 + runner.mu.Unlock() + if ready { + return + } + runtime.Gosched() + } + t.Fatal("replacement did not restore capacity") +} diff --git a/go/cloud-query/internal/tools/python/internal/pool/process.go b/go/cloud-query/internal/tools/python/internal/pool/process.go new file mode 100644 index 0000000000..dc4bfb4bae --- /dev/null +++ b/go/cloud-query/internal/tools/python/internal/pool/process.go @@ -0,0 +1,102 @@ +package pool + +import ( + "context" + "fmt" + "io" + "os/exec" + "sync" + + "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/contract" + "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/protocol" +) + +type process interface { + exchange(context.Context, protocol.Request) (protocol.Response, error) + stop() +} + +type processFactory func(ProcessConfig) (process, error) + +type execProcess struct { + stdin io.WriteCloser + stdout io.ReadCloser + kill func() + codec protocol.Codec + mu sync.Mutex + once sync.Once +} + +func newExecProcess(config ProcessConfig) (process, error) { + cmd := exec.Command(config.Executable, config.Arguments...) + cmd.Env = append([]string(nil), config.Environment...) + + stdin, err := cmd.StdinPipe() + if err != nil { + return nil, fmt.Errorf("opening python worker stdin: %w", err) + } + + stdout, err := cmd.StdoutPipe() + if err != nil { + return nil, fmt.Errorf("opening python worker stdout: %w", err) + } + + if err = cmd.Start(); err != nil { + return nil, fmt.Errorf("starting python worker process: %w", err) + } + + p := &execProcess{ + stdin: stdin, + stdout: stdout, + codec: protocol.NewCodec(protocol.MaxFrameSize), + } + + p.kill = func() { + _ = stdin.Close() + _ = stdout.Close() + if cmd.Process != nil { + _ = cmd.Process.Kill() + } + _ = cmd.Wait() + } + + return p, nil +} + +func (p *execProcess) stop() { + p.once.Do(p.kill) +} + +func (p *execProcess) exchange(ctx context.Context, request protocol.Request) (protocol.Response, error) { + p.mu.Lock() + defer p.mu.Unlock() + + result := make(chan struct { + response protocol.Response + err error + }, 1) + go func() { + if err := p.codec.WriteRequest(p.stdin, request); err != nil { + result <- struct { + response protocol.Response + err error + }{err: err} + return + } + response, err := p.codec.ReadResponse(p.stdout, request) + result <- struct { + response protocol.Response + err error + }{response, err} + }() + select { + case result := <-result: + if result.err != nil { + return protocol.Response{}, contract.InternalError(result.err) + } + return result.response, nil + case <-ctx.Done(): + p.stop() + return protocol.Response{}, contract.ContextError(ctx.Err()) + } +} diff --git a/go/cloud-query/internal/tools/python/internal/protocol/codec.go b/go/cloud-query/internal/tools/python/internal/protocol/codec.go new file mode 100644 index 0000000000..3def9e254a --- /dev/null +++ b/go/cloud-query/internal/tools/python/internal/protocol/codec.go @@ -0,0 +1,244 @@ +package protocol + +import ( + "bytes" + "encoding/binary" + "encoding/json" + "errors" + "fmt" + "io" + + "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/contract" +) + +// Version identifies the private parent-to-worker protocol version. +const Version = 1 + +// MaxFrameSize bounds a framed protocol payload before the codec allocates it. +const MaxFrameSize = contract.MaxSourceBytes + + contract.MaxInputBytes + + contract.MaxResultBytes + + contract.MaxStdoutBytes + + contract.MaxDetailBytes + + 16<<10 +const ( + // Health requests a worker startup or liveness check. + Health = "health" + // Run requests one restricted Python execution. + Run = "run" +) + +// Request is a validated parent-to-worker protocol message. +type Request struct { + Version int `json:"version"` + Kind string `json:"kind"` + ID string `json:"id"` + Script string `json:"script,omitempty"` + InputJSON string `json:"input_json,omitempty"` +} + +// Response is a validated worker-to-parent protocol message. +type Response struct { + Version int `json:"version"` + Kind string `json:"kind"` + ID string `json:"id"` + ResultJSON string `json:"result_json,omitempty"` + Stdout string `json:"stdout,omitempty"` + Error *WireError `json:"error,omitempty"` +} + +// WireError carries a stable code, sanitized public summary, and bounded private +// diagnostic over the worker protocol. Detail must not reach untrusted callers. +type WireError struct { + Code contract.Code `json:"code"` + PublicMessage string `json:"public_message"` + Detail string `json:"detail"` +} + +// Codec reads and writes strict, bounded, length-prefixed protocol frames. +type Codec struct{ maxFrameSize int } + +// NewCodec creates a codec that rejects frames larger than maxFrameSize. +func NewCodec(maxFrameSize int) Codec { return Codec{maxFrameSize: maxFrameSize} } + +// ReadRequest reads and validates one request frame. +func (c Codec) ReadRequest(in io.Reader) (Request, error) { + var request Request + if err := c.read(in, &request); err != nil { + return Request{}, err + } + if err := c.validRequest(request); err != nil { + return Request{}, err + } + return request, nil +} + +// WriteRequest validates and writes one request frame. +func (c Codec) WriteRequest(out io.Writer, request Request) error { + if err := c.validRequest(request); err != nil { + return err + } + return c.write(out, request) +} + +// ReadResponse reads and validates a response matched to request. +func (c Codec) ReadResponse(in io.Reader, request Request) (Response, error) { + var response Response + if err := c.read(in, &response); err != nil { + return Response{}, err + } + if err := c.validResponse(request, response); err != nil { + return Response{}, err + } + return response, nil +} + +// WriteResponse validates and writes a response matched to request. +func (c Codec) WriteResponse(out io.Writer, request Request, response Response) error { + if err := c.validResponse(request, response); err != nil { + return err + } + return c.write(out, response) +} + +func (c Codec) read(in io.Reader, into any) error { + var header [4]byte + if _, err := io.ReadFull(in, header[:]); err != nil { + return err + } + + length := int(binary.BigEndian.Uint32(header[:])) + if length == 0 || length > c.maxFrameSize { + return errors.New("invalid python protocol frame") + } + + body := make([]byte, length) + if _, err := io.ReadFull(in, body); err != nil { + return err + } + + decoder := json.NewDecoder(bytes.NewReader(body)) + decoder.DisallowUnknownFields() + if err := decoder.Decode(into); err != nil { + return errors.New("invalid python protocol payload") + } + if err := decoder.Decode(&struct{}{}); !errors.Is(err, io.EOF) { + return errors.New("invalid python protocol payload") + } + return nil +} + +func (c Codec) write(out io.Writer, value any) error { + body, err := json.Marshal(value) + if err != nil || len(body) == 0 || len(body) > c.maxFrameSize { + return errors.New("invalid python protocol frame") + } + + frame := make([]byte, 4+len(body)) + binary.BigEndian.PutUint32(frame, uint32(len(body))) + copy(frame[4:], body) + + for len(frame) > 0 { + n, err := out.Write(frame) + if err != nil { + return err + } + if n <= 0 { + return io.ErrShortWrite + } + frame = frame[n:] + } + return nil +} + +func (c Codec) validRequest(request Request) error { + if request.Version != Version || request.ID == "" { + return errors.New("invalid python protocol request") + } + + switch request.Kind { + case Health: + if request.Script != "" || request.InputJSON != "" { + return errors.New("invalid python protocol request") + } + case Run: + if request.Script == "" || + len(request.Script) > contract.MaxSourceBytes || + len(request.InputJSON) > contract.MaxInputBytes { + return errors.New("invalid python protocol request") + } + if _, err := contract.NormalizeInput(request.InputJSON); err != nil { + return errors.New("invalid python protocol request") + } + default: + return errors.New("invalid python protocol request") + } + return nil +} + +func (c Codec) validResponse(request Request, response Response) error { + if response.Version != Version || + response.Kind != request.Kind || + response.ID != request.ID { + return errors.New("invalid python protocol response") + } + + if response.Error != nil { + if !c.known(response.Error.Code) || + response.Error.PublicMessage == "" || + len(response.Error.PublicMessage) > contract.MaxPublicMessageBytes || + response.Error.Detail == "" || + len(response.Error.Detail) > contract.MaxDetailBytes || + response.ResultJSON != "" || + response.Stdout != "" { + return errors.New("invalid python protocol response") + } + return nil + } + + if request.Kind == Health { + if response.ResultJSON != "" || response.Stdout != "" { + return errors.New("invalid python protocol response") + } + return nil + } + + if response.ResultJSON == "" || + len(response.ResultJSON) > contract.MaxResultBytes || + len(response.Stdout) > contract.MaxStdoutBytes || + !contract.IsJSONObject(response.ResultJSON) { + return errors.New("invalid python protocol response") + } + return nil +} + +func (Codec) known(code contract.Code) bool { + switch code { + case contract.InvalidArgument, + contract.FailedPrecondition, + contract.Canceled, + contract.DeadlineExceeded, + contract.ResourceExhausted, + contract.Unavailable, + contract.Internal: + return true + } + return false +} + +// Error converts a worker error response into an error with a safe public +// summary and private bounded diagnostic. +func Error(response Response) error { + if response.Error == nil { + return nil + } + if !(Codec{}).known(response.Error.Code) { + return contract.InternalError(fmt.Errorf("unknown python worker error code")) + } + + return contract.New( + response.Error.Code, + response.Error.PublicMessage, + errors.New(response.Error.Detail), + ) +} diff --git a/go/cloud-query/internal/tools/python/internal/protocol/codec_test.go b/go/cloud-query/internal/tools/python/internal/protocol/codec_test.go new file mode 100644 index 0000000000..c5f69f4c43 --- /dev/null +++ b/go/cloud-query/internal/tools/python/internal/protocol/codec_test.go @@ -0,0 +1,150 @@ +package protocol + +import ( + "bytes" + "encoding/binary" + "encoding/json" + "io" + "strings" + "testing" + + "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/contract" +) + +func TestCodecRejectsMalformedAndInvalidResponses(t *testing.T) { + codec := NewCodec(128) + invalidBodies := [][]byte{ + nil, + []byte(`{"version":1,"kind":"health","id":"x","extra":true}`), + []byte(`{"version":2,"kind":"health","id":"x"}`), + } + for _, body := range invalidBodies { + var input bytes.Buffer + var header [4]byte + binary.BigEndian.PutUint32(header[:], uint32(len(body))) + input.Write(header[:]) + input.Write(body) + if _, err := codec.ReadRequest(&input); err == nil { + t.Fatal("accepted invalid request") + } + } + request := Request{Version: Version, Kind: Run, ID: "x", Script: "output={}", InputJSON: "{}"} + response := Response{ + Version: Version, + Kind: Run, + ID: "x", + Error: &WireError{ + Code: contract.Internal, + PublicMessage: "python runtime failed", + Detail: "private", + }, + } + if err := codec.WriteResponse(&bytes.Buffer{}, request, response); err != nil { + t.Fatal(err) + } + if Public := Error(response); contract.PublicMessage(Public) != "python runtime failed" { + t.Fatal("private detail leaked") + } +} + +func TestCodecFrameAndSemanticInvariants(t *testing.T) { + codec := NewCodec(MaxFrameSize) + request := Request{Version: Version, Kind: Run, ID: "request", Script: "output={}", InputJSON: "{}"} + for _, body := range [][]byte{ + []byte(`{"version":1,"kind":"unknown","id":"request"}`), + []byte(`{"version":1,"kind":"run","id":"request","script":"","input_json":"{}"}`), + []byte(`{"version":1,"kind":"health","id":"request","script":"x"}`), + []byte(`{"version":1,"kind":"run","id":"request","script":"output={}","input_json":"{}"} {}`), + } { + if _, err := codec.ReadRequest(frame(body)); err == nil { + t.Fatalf("accepted request %s", body) + } + } + for _, response := range []Response{ + {Version: 2, Kind: Run, ID: request.ID, ResultJSON: "{}"}, + {Version: Version, Kind: Health, ID: request.ID, ResultJSON: "{}"}, + {Version: Version, Kind: Run, ID: "other", ResultJSON: "{}"}, + { + Version: Version, + Kind: Run, + ID: request.ID, + ResultJSON: "{}", + Error: &WireError{ + Code: contract.Internal, + PublicMessage: "python runtime failed", + Detail: "x", + }, + }, + {Version: Version, Kind: Run, ID: request.ID, Error: &WireError{Code: contract.Internal}}, + { + Version: Version, + Kind: Run, + ID: request.ID, + Error: &WireError{ + Code: contract.Internal, + PublicMessage: "python runtime failed", + Detail: strings.Repeat("x", contract.MaxDetailBytes+1), + }, + }, + {Version: Version, Kind: Run, ID: request.ID, ResultJSON: "[]"}, + {Version: Version, Kind: Run, ID: request.ID, ResultJSON: strings.Repeat("x", contract.MaxResultBytes+1)}, + } { + var encoded bytes.Buffer + body, _ := jsonMarshal(response) + encoded.Write(frameBytes(body)) + if _, err := codec.ReadResponse(&encoded, request); err == nil { + t.Fatalf("accepted response %#v", response) + } + } + for _, input := range []io.Reader{bytes.NewReader([]byte{0, 0, 0, 0}), bytes.NewReader([]byte{0, 0, 0, 1, '{'}), frame([]byte(`{`))} { + if _, err := codec.ReadRequest(input); err == nil { + t.Fatal("accepted malformed frame") + } + } + oversized := make([]byte, MaxFrameSize+1) + binary.BigEndian.PutUint32(oversized[:4], uint32(MaxFrameSize+1)) + if _, err := codec.ReadRequest(bytes.NewReader(oversized[:4])); err == nil { + t.Fatal("accepted oversized frame") + } +} + +type shortWriter struct{ writes int } + +func (w *shortWriter) Write(p []byte) (int, error) { + w.writes++ + if w.writes == 1 { + return 1, nil + } + return 0, nil +} + +func TestCodecHandlesShortWritesAndPrivateRemoteDetail(t *testing.T) { + codec := NewCodec(MaxFrameSize) + request := Request{Version: Version, Kind: Health, ID: "x"} + if err := codec.WriteRequest(&shortWriter{}, request); err != io.ErrShortWrite { + t.Fatalf("short write: %v", err) + } + err := Error(Response{Error: &WireError{ + Code: contract.Internal, + PublicMessage: "python runtime failed", + Detail: "secret path /tmp/private", + }}) + if strings.Contains(contract.PublicMessage(err), "secret") { + t.Fatal("remote detail leaked") + } +} + +func frame(body []byte) io.Reader { return bytes.NewReader(frameBytes(body)) } + +func frameBytes(body []byte) []byte { + var header [4]byte + binary.BigEndian.PutUint32(header[:], uint32(len(body))) + return append(header[:], body...) +} + +func jsonMarshal(value any) ([]byte, error) { + var buffer bytes.Buffer + encoder := json.NewEncoder(&buffer) + err := encoder.Encode(value) + return bytes.TrimSpace(buffer.Bytes()), err +} diff --git a/go/cloud-query/internal/tools/python/internal/worker/os.go b/go/cloud-query/internal/tools/python/internal/worker/os.go new file mode 100644 index 0000000000..0d0b882b43 --- /dev/null +++ b/go/cloud-query/internal/tools/python/internal/worker/os.go @@ -0,0 +1,86 @@ +package worker + +import ( + "context" + "time" + + monty "github.com/ewhauser/gomonty" +) + +const ( + osDateToday monty.OSFunction = "date.today" + osDateTimeNow monty.OSFunction = "datetime.now" +) + +// handleOS exposes the only host callbacks available to sandboxed Python. +// Clock values are always derived from UTC; filesystem, environment, and all +// other OS callbacks remain unavailable. +func (m *montyRuntime) handleOS(_ context.Context, call monty.OSCall) (monty.Result, error) { + switch call.Function { + case osDateToday: + if len(call.Args) != 0 || len(call.Kwargs) != 0 { + return m.invalidClockArguments(call.Function) + } + return monty.Return(monty.DateValue(m.dateOf(m.utcNow()))), nil + case osDateTimeNow: + return m.dateTimeNow(call) + default: + message := "OS function " + string(call.Function) + " is not available" + return monty.Raise(monty.Exception{Type: "NotImplementedError", Arg: &message}), nil + } +} + +func (m *montyRuntime) dateTimeNow(call monty.OSCall) (monty.Result, error) { + if len(call.Kwargs) != 0 || len(call.Args) > 1 { + return m.invalidClockArguments(call.Function) + } + + now := m.dateTimeOf(m.utcNow()) + if len(call.Args) == 0 || call.Args[0].Raw() == nil { + return monty.Return(monty.DateTimeValue(now)), nil + } + + timezone, ok := call.Args[0].TimeZone() + if !ok || timezone.OffsetSeconds != 0 { + message := "datetime.now only supports UTC" + return monty.Raise(monty.Exception{Type: "ValueError", Arg: &message}), nil + } + + zeroOffset := int32(0) + utcName := "UTC" + now.OffsetSeconds = &zeroOffset + now.TimezoneName = &utcName + return monty.Return(monty.DateTimeValue(now)), nil +} + +func (m *montyRuntime) utcNow() time.Time { + if m.now == nil { + return time.Now().UTC() + } + return m.now().UTC() +} + +func (m *montyRuntime) dateOf(value time.Time) monty.Date { + return monty.Date{ + Year: int32(value.Year()), + Month: uint8(value.Month()), + Day: uint8(value.Day()), + } +} + +func (m *montyRuntime) dateTimeOf(value time.Time) monty.DateTime { + return monty.DateTime{ + Year: int32(value.Year()), + Month: uint8(value.Month()), + Day: uint8(value.Day()), + Hour: uint8(value.Hour()), + Minute: uint8(value.Minute()), + Second: uint8(value.Second()), + Microsecond: uint32(time.Duration(value.Nanosecond()) / time.Microsecond), + } +} + +func (m *montyRuntime) invalidClockArguments(function monty.OSFunction) (monty.Result, error) { + message := string(function) + " received unsupported arguments" + return monty.Raise(monty.Exception{Type: "TypeError", Arg: &message}), nil +} diff --git a/go/cloud-query/internal/tools/python/internal/worker/runtime.go b/go/cloud-query/internal/tools/python/internal/worker/runtime.go new file mode 100644 index 0000000000..e04bc17b8a --- /dev/null +++ b/go/cloud-query/internal/tools/python/internal/worker/runtime.go @@ -0,0 +1,165 @@ +package worker + +import ( + "context" + "errors" + "strconv" + "strings" + "time" + + monty "github.com/ewhauser/gomonty" + + "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/contract" +) + +const executionTimeout = 60 * time.Second +const maxMemoryBytes = 100 << 20 +const maxRecursionDepth = 100 + +type montyRuntime struct { + now func() time.Time +} + +func newMontyRuntime() *montyRuntime { return &montyRuntime{now: time.Now} } + +func (m *montyRuntime) Health() error { + _, err := monty.NewRepl(monty.ReplOptions{ + ScriptName: "workbench.py", + Limits: m.limits(), + }) + if err != nil { + return contract.InternalError(err) + } + return nil +} + +func (m *montyRuntime) Run(ctx context.Context, script, inputJSON string) (*contract.RunOutput, error) { + input, err := contract.ValidateRun(contract.RunInput{Script: script, InputJSON: inputJSON}) + if err != nil { + return nil, err + } + + repl, err := monty.NewRepl(monty.ReplOptions{ + ScriptName: "workbench.py", + Limits: m.limits(), + }) + if err != nil { + return nil, contract.InternalError(err) + } + + ctx, cancel := context.WithTimeout(ctx, executionTimeout) + defer cancel() + + initialization := "import json as __workbench_json\n" + + "input = __workbench_json.loads(" + strconv.Quote(input.InputJSON) + ")\n" + + "output = {}" + + if _, err = repl.FeedRun(ctx, initialization, m.feedOptions()); err != nil { + return nil, m.mapError(err) + } + + var stdout strings.Builder + tooLarge := false + printFn := func(stream, text string) { + if stream != "stdout" { + return + } + if stdout.Len()+len(text) > contract.MaxStdoutBytes { + tooLarge = true + return + } + stdout.WriteString(text) + } + + options := m.feedOptions() + options.Print = printFn + if _, err = repl.FeedRun(ctx, input.Script, options); err != nil { + return nil, m.mapError(err) + } + + if tooLarge { + return nil, contract.Exhausted("python stdout exceeds the stdout limit", nil) + } + + value, err := repl.FeedRun(ctx, "__workbench_json.dumps(output)", m.feedOptions()) + if err != nil { + return nil, m.mapError(err) + } + + result, ok := value.Raw().(string) + if !ok || !contract.IsJSONObject(result) { + return nil, contract.Invalid("output must be a JSON object", nil) + } + + if len(result) > contract.MaxResultBytes { + return nil, contract.Exhausted("python result exceeds the result limit", nil) + } + + return &contract.RunOutput{ + ResultJSON: result, + Stdout: stdout.String(), + }, nil +} + +func (m *montyRuntime) feedOptions() monty.FeedOptions { + return monty.FeedOptions{OS: m.handleOS} +} + +func (m *montyRuntime) limits() *monty.ResourceLimits { + return &monty.ResourceLimits{ + MaxDuration: executionTimeout, + MaxMemory: maxMemoryBytes, + MaxRecursionDepth: maxRecursionDepth, + } +} + +func (m *montyRuntime) mapError(err error) error { + if errors.Is(err, context.DeadlineExceeded) { + return contract.Deadline(err) + } + + var syntax *monty.SyntaxError + if errors.As(err, &syntax) { + return contract.Invalid("python code is invalid", err) + } + + var runtime *monty.RuntimeError + if !errors.As(err, &runtime) { + return contract.InternalError(err) + } + + kind, detail, ok := strings.Cut(strings.TrimSpace(runtime.Error()), ":") + if !ok { + kind, detail = "Exception", runtime.Error() + } + + switch kind { + case "SyntaxError", "TypeError": + return contract.Invalid(m.toErrorSummary(kind, detail), err) + case "MemoryError", "RecursionError": + return contract.Exhausted("python resource limit exceeded", err) + case "TimeoutError": + return contract.Deadline(err) + case "ResourceError": + if strings.Contains(strings.ToLower(detail), "time") || strings.Contains(strings.ToLower(detail), "duration") { + return contract.Deadline(err) + } + return contract.Exhausted("python resource limit exceeded", err) + default: + return contract.Failed(m.toErrorSummary(kind, detail), err) + } +} + +func (m *montyRuntime) toErrorSummary(kind, detail string) string { + kind = strings.TrimSpace(kind) + detail = strings.Join(strings.Fields(detail), " ") + + if kind == "" { + kind = "Exception" + } + + if detail == "" { + return "python " + kind + } + return "python " + kind + ": " + detail +} diff --git a/go/cloud-query/internal/tools/python/internal/worker/server.go b/go/cloud-query/internal/tools/python/internal/worker/server.go new file mode 100644 index 0000000000..050bc169fc --- /dev/null +++ b/go/cloud-query/internal/tools/python/internal/worker/server.go @@ -0,0 +1,76 @@ +package worker + +import ( + "context" + "errors" + "io" + + "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/contract" + "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/protocol" +) + +type runtime interface { + Health() error + Run(context.Context, string, string) (*contract.RunOutput, error) +} + +// Server processes private protocol requests using one restricted runtime. +type Server struct { + runtime runtime + codec protocol.Codec +} + +// NewServer creates a server backed by the restricted Monty runtime. +func NewServer() *Server { + return newServer(newMontyRuntime()) +} + +func newServer(runtime runtime) *Server { + return &Server{runtime: runtime, codec: protocol.NewCodec(protocol.MaxFrameSize)} +} + +// Run processes protocol requests until input reaches EOF or a protocol or I/O +// failure occurs. It returns only private diagnostics to its caller. +func (s *Server) Run(in io.Reader, out io.Writer) error { + for { + request, err := s.codec.ReadRequest(in) + if errors.Is(err, io.EOF) { + return nil + } + if err != nil { + return contract.InternalError(err) + } + + response := protocol.Response{ + Version: protocol.Version, + Kind: request.Kind, + ID: request.ID, + } + + switch request.Kind { + case protocol.Health: + err = s.runtime.Health() + case protocol.Run: + var result *contract.RunOutput + result, err = s.runtime.Run(context.Background(), request.Script, request.InputJSON) + if err == nil { + response.ResultJSON = result.ResultJSON + response.Stdout = result.Stdout + } + default: + err = contract.InternalError(errors.New("unreachable protocol kind")) + } + + if err != nil { + response.Error = &protocol.WireError{ + Code: contract.CodeOf(err), + PublicMessage: contract.PublicMessage(err), + Detail: contract.Detail(err), + } + } + + if err := s.codec.WriteResponse(out, request, response); err != nil { + return contract.InternalError(err) + } + } +} diff --git a/go/cloud-query/internal/tools/python/internal/worker/server_test.go b/go/cloud-query/internal/tools/python/internal/worker/server_test.go new file mode 100644 index 0000000000..b80b9f0826 --- /dev/null +++ b/go/cloud-query/internal/tools/python/internal/worker/server_test.go @@ -0,0 +1,141 @@ +package worker + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "strings" + "testing" + "time" + + "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/contract" + "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/protocol" +) + +type fakeRuntime struct { + health error + output *contract.RunOutput + run error + calls int +} + +func (f *fakeRuntime) Health() error { return f.health } + +func (f *fakeRuntime) Run(_ context.Context, _, _ string) (*contract.RunOutput, error) { + f.calls++ + return f.output, f.run +} + +func TestServerDispatchesAndKeepsErrorDetailPrivate(t *testing.T) { + codec := protocol.NewCodec(protocol.MaxFrameSize) + runtime := &fakeRuntime{output: &contract.RunOutput{ResultJSON: "{}"}} + var in, out bytes.Buffer + request := protocol.Request{Version: protocol.Version, Kind: protocol.Run, ID: "1", Script: "output={}", InputJSON: "{}"} + if err := codec.WriteRequest(&in, request); err != nil { + t.Fatal(err) + } + if err := newServer(runtime).Run(&in, &out); err != nil { + t.Fatal(err) + } + response, err := codec.ReadResponse(&out, request) + if err != nil || response.ResultJSON != "{}" || runtime.calls != 1 { + t.Fatalf("response=%#v err=%v", response, err) + } + runtime.run = contract.Failed("python execution failed", errors.New("private host detail")) + in.Reset() + out.Reset() + _ = codec.WriteRequest(&in, request) + if err := newServer(runtime).Run(&in, &out); err != nil { + t.Fatal(err) + } + response, err = codec.ReadResponse(&out, request) + if err != nil || response.Error == nil || contract.PublicMessage(protocol.Error(response)) != "python execution failed" { + t.Fatalf("error response=%#v err=%v", response, err) + } +} + +func TestMontyRuntimeFreshStateAndStdout(t *testing.T) { + runtime := newMontyRuntime() + first, err := runtime.Run(context.Background(), "secret = 42\noutput = {'sum': input['value'] + 1}\nprint('ok')", `{"value":1}`) + if err != nil || first.ResultJSON != "{\"sum\": 2}" || first.Stdout != "ok\n" { + t.Fatalf("first=%#v err=%v", first, err) + } + if _, err := runtime.Run(context.Background(), "output = {'secret': secret}", "{}"); contract.CodeOf(err) != contract.FailedPrecondition { + t.Fatalf("state leaked: %v", err) + } +} + +func TestMontyRuntimeLimitsMatchWorkbenchDefaults(t *testing.T) { + limits := newMontyRuntime().limits() + if limits.MaxDuration != 60*time.Second || limits.MaxMemory != 100<<20 || limits.MaxRecursionDepth != 100 { + t.Fatalf("limits = %#v", limits) + } +} + +func TestMontyRuntimeUsesUTCClockOnly(t *testing.T) { + runtime := &montyRuntime{ + now: func() time.Time { + return time.Date(2026, time.April, 5, 6, 7, 8, 123456000, time.FixedZone("CEST", 2*60*60)) + }, + } + + output, err := runtime.Run(context.Background(), ` +from datetime import date, datetime +now = datetime.now() +today = date.today() +output = {"now": now.isoformat(), "today": today.isoformat()} +`, "{}") + if err != nil { + t.Fatal(err) + } + + var result map[string]string + if err := json.Unmarshal([]byte(output.ResultJSON), &result); err != nil { + t.Fatal(err) + } + if want := map[string]string{ + "now": "2026-04-05T04:07:08.123456", + "today": "2026-04-05", + }; !equalStringMaps(result, want) { + t.Fatalf("result = %#v, want %#v", result, want) + } + + for name, script := range map[string]string{ + "environment": "import os\noutput = {'value': os.getenv('SECRET')}", + "filesystem": "from pathlib import Path\noutput = {'value': Path('/tmp/test').exists()}", + "non-UTC clock": ` +from datetime import datetime, timedelta, timezone +output = {"value": datetime.now(timezone(timedelta(hours=1))).isoformat()} +`, + } { + t.Run(name, func(t *testing.T) { + _, err := runtime.Run(context.Background(), script, "{}") + if contract.CodeOf(err) != contract.FailedPrecondition { + t.Fatalf("err = %v", err) + } + + message := contract.PublicMessage(err) + if name == "non-UTC clock" && !strings.Contains(message, "only supports UTC") { + t.Fatalf("err = %v", err) + } + if name != "non-UTC clock" && !strings.Contains(message, "NotImplementedError") { + t.Fatalf("err = %v", err) + } + }) + } +} + +func equalStringMaps(got, want map[string]string) bool { + if len(got) != len(want) { + return false + } + + for key, wantValue := range want { + if got[key] != wantValue { + return false + } + } + + return true +} diff --git a/go/cloud-query/internal/tools/python/python_test.go b/go/cloud-query/internal/tools/python/python_test.go new file mode 100644 index 0000000000..88254808a9 --- /dev/null +++ b/go/cloud-query/internal/tools/python/python_test.go @@ -0,0 +1,54 @@ +package python + +import ( + "context" + "os" + "strings" + "testing" +) + +func TestMain(m *testing.M) { + if len(os.Args) == 2 && os.Args[1] == "python-worker" { + if err := NewWorker().Run(os.Stdin, os.Stdout); err != nil { + os.Exit(1) + } + os.Exit(0) + } + os.Exit(m.Run()) +} + +func TestRunnerSelfSpawnsWorker(t *testing.T) { + runner, err := NewRunner(context.Background(), RunnerConfig{Workers: 1, QueueSize: 1, MaxSuccessfulRunsBeforeRecycle: 1}) + if err != nil { + t.Fatal(err) + } + defer runner.CloseContext(context.Background()) + output, err := runner.Run(context.Background(), RunInput{Script: "output = {'value': 42}"}) + if err != nil || output.ResultJSON != "{\"value\": 42}" { + t.Fatalf("run = %#v, %v", output, err) + } +} + +func TestRunnerPreservesSandboxErrorMessage(t *testing.T) { + runner, err := NewRunner(context.Background(), RunnerConfig{ + Workers: 1, + QueueSize: 1, + MaxSuccessfulRunsBeforeRecycle: 1, + }) + if err != nil { + t.Fatal(err) + } + defer runner.CloseContext(context.Background()) + + _, err = runner.Run(context.Background(), RunInput{ + Script: "import socket\nsocket.create_connection(('example.com', 443))\noutput = {'reachable': True}", + }) + if err == nil { + t.Fatal("sandboxed socket import unexpectedly succeeded") + } + + message := PublicMessage(err) + if message == "python execution failed" || !strings.Contains(strings.ToLower(message), "socket") { + t.Fatalf("sandbox error message = %q", message) + } +} diff --git a/go/cloud-query/internal/tools/python/runner.go b/go/cloud-query/internal/tools/python/runner.go new file mode 100644 index 0000000000..6c2c8e6a0c --- /dev/null +++ b/go/cloud-query/internal/tools/python/runner.go @@ -0,0 +1,34 @@ +package python + +import ( + "context" + + "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/pool" +) + +// Runner accepts Python executions and manages its worker processes. +type Runner interface { + // Run validates and executes input, respecting ctx while queued and running. + Run(context.Context, RunInput) (*RunOutput, error) + // CloseContext stops admission and waits for workers until ctx expires. + CloseContext(context.Context) error +} + +// NewRunner starts a configured set of isolated Python workers. +func NewRunner(ctx context.Context, config RunnerConfig) (Runner, error) { + config, err := resolveRunnerConfig(config) + if err != nil { + return nil, err + } + + return pool.New(ctx, pool.Config{ + Workers: config.Workers, + QueueSize: config.QueueSize, + MaxSuccessfulRunsBeforeRecycle: config.MaxSuccessfulRunsBeforeRecycle, + Process: pool.ProcessConfig{ + Executable: config.WorkerProcess.Executable, + Arguments: config.WorkerProcess.Arguments, + Environment: config.WorkerProcess.Environment, + }, + }) +} diff --git a/go/cloud-query/internal/tools/python/types.go b/go/cloud-query/internal/tools/python/types.go new file mode 100644 index 0000000000..deac542ecc --- /dev/null +++ b/go/cloud-query/internal/tools/python/types.go @@ -0,0 +1,29 @@ +package python + +import "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/contract" + +// RunInput contains the source and JSON object supplied to a Python execution. +type RunInput = contract.RunInput + +// RunOutput contains the JSON object assigned to output and captured stdout. +type RunOutput = contract.RunOutput + +// Code identifies a stable category of Python runner failure. +type Code = contract.Code + +const ( + // InvalidArgument reports invalid source, input, or output values. + InvalidArgument = contract.InvalidArgument + // FailedPrecondition reports an execution failure after input validation. + FailedPrecondition = contract.FailedPrecondition + // Canceled reports cancellation by the caller. + Canceled = contract.Canceled + // DeadlineExceeded reports an execution deadline that elapsed. + DeadlineExceeded = contract.DeadlineExceeded + // ResourceExhausted reports a configured capacity or resource limit. + ResourceExhausted = contract.ResourceExhausted + // Unavailable reports that no Python runtime worker can accept the request. + Unavailable = contract.Unavailable + // Internal reports an unclassified runner or worker failure. + Internal = contract.Internal +) diff --git a/go/cloud-query/internal/tools/python/worker.go b/go/cloud-query/internal/tools/python/worker.go new file mode 100644 index 0000000000..816ca14153 --- /dev/null +++ b/go/cloud-query/internal/tools/python/worker.go @@ -0,0 +1,16 @@ +package python + +import ( + "io" + + internalworker "github.com/pluralsh/console/go/cloud-query/internal/tools/python/internal/worker" +) + +// Worker serves the private parent-to-worker protocol on a pair of streams. +type Worker interface { + // Run processes requests until the input reaches EOF or an I/O error occurs. + Run(io.Reader, io.Writer) error +} + +// NewWorker constructs the worker implementation used by the python-worker command. +func NewWorker() Worker { return internalworker.NewServer() } diff --git a/go/controller/go.mod b/go/controller/go.mod index 705b793031..93e10c71b1 100644 --- a/go/controller/go.mod +++ b/go/controller/go.mod @@ -1,6 +1,6 @@ module github.com/pluralsh/console/go/controller -go 1.26.5 +go 1.26.6 replace ( github.com/pluralsh/console/go/client => ../client diff --git a/go/datastore/go.mod b/go/datastore/go.mod index 64990a3532..9fcbc54bb6 100644 --- a/go/datastore/go.mod +++ b/go/datastore/go.mod @@ -1,6 +1,6 @@ module github.com/pluralsh/console/go/datastore -go 1.26.5 +go 1.26.6 replace ( github.com/pluralsh/console/go/client => ../client diff --git a/go/demo/flaky-service/go.mod b/go/demo/flaky-service/go.mod index 72aabb4930..5d4349cd47 100644 --- a/go/demo/flaky-service/go.mod +++ b/go/demo/flaky-service/go.mod @@ -1,6 +1,6 @@ module github.com/pluralsh/console/go/demo/flaky-service -go 1.26.5 +go 1.26.6 require github.com/prometheus/client_golang v1.23.2 diff --git a/go/deployment-operator/dockerfiles/agent-harness/base.Dockerfile b/go/deployment-operator/dockerfiles/agent-harness/base.Dockerfile index 142e4412c2..3c336c4edf 100644 --- a/go/deployment-operator/dockerfiles/agent-harness/base.Dockerfile +++ b/go/deployment-operator/dockerfiles/agent-harness/base.Dockerfile @@ -219,4 +219,4 @@ WORKDIR /plural HEALTHCHECK --interval=60s --timeout=10s --start-period=30s --retries=5 \ CMD kill -0 1 || exit 1 -ENTRYPOINT ["/entrypoint.sh", "/agent-harness", "--working-dir=/plural"] \ No newline at end of file +ENTRYPOINT ["/entrypoint.sh", "/agent-harness", "--working-dir=/plural"] diff --git a/go/deployment-operator/dockerfiles/sentinel-harness/base.Dockerfile b/go/deployment-operator/dockerfiles/sentinel-harness/base.Dockerfile index a321a613e3..965b12a44d 100644 --- a/go/deployment-operator/dockerfiles/sentinel-harness/base.Dockerfile +++ b/go/deployment-operator/dockerfiles/sentinel-harness/base.Dockerfile @@ -55,4 +55,4 @@ WORKDIR /plural USER 65532:65532 HEALTHCHECK --interval=60s --timeout=10s --start-period=30s --retries=5 \ - CMD kill -0 1 || exit 1 \ No newline at end of file + CMD kill -0 1 || exit 1 diff --git a/go/deployment-operator/dockerfiles/sentinel-harness/terratest.Dockerfile b/go/deployment-operator/dockerfiles/sentinel-harness/terratest.Dockerfile index 398c1ba0b9..039063e78c 100644 --- a/go/deployment-operator/dockerfiles/sentinel-harness/terratest.Dockerfile +++ b/go/deployment-operator/dockerfiles/sentinel-harness/terratest.Dockerfile @@ -16,7 +16,7 @@ ENV CGO_ENABLED=0 \ WORKDIR /sentinel/terratest # Copy test files -COPY deployment-operator/terratest ./ +COPY --chown=65532:65532 deployment-operator/terratest ./ RUN go mod download diff --git a/go/deployment-operator/go.mod b/go/deployment-operator/go.mod index b9d2ece481..8301dd1173 100644 --- a/go/deployment-operator/go.mod +++ b/go/deployment-operator/go.mod @@ -1,6 +1,6 @@ module github.com/pluralsh/console/go/deployment-operator -go 1.26.5 +go 1.26.6 replace ( github.com/containerd/containerd => github.com/containerd/containerd v1.7.33 diff --git a/go/deployment-operator/terratest/go.mod b/go/deployment-operator/terratest/go.mod index b66fa82a15..1e0ae83af7 100644 --- a/go/deployment-operator/terratest/go.mod +++ b/go/deployment-operator/terratest/go.mod @@ -1,6 +1,6 @@ module github.com/pluralsh/console/go/deployment-operator/terratest -go 1.26.5 +go 1.26.6 require ( github.com/gruntwork-io/terratest v1.0.1 @@ -181,7 +181,7 @@ require ( golang.org/x/text v0.39.0 // indirect golang.org/x/time v0.15.0 // indirect golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da // indirect - google.golang.org/genproto/googleapis/rpc v0.0.0-20260406210006-6f92a3bedf2d // indirect + google.golang.org/genproto/googleapis/rpc v0.0.0-20260414002931-afd174a4e478 // indirect google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af // indirect gopkg.in/evanphx/json-patch.v4 v4.13.0 // indirect gopkg.in/inf.v0 v0.9.1 // indirect diff --git a/go/deployment-operator/terratest/go.sum b/go/deployment-operator/terratest/go.sum index d8958d7f02..f314779933 100644 --- a/go/deployment-operator/terratest/go.sum +++ b/go/deployment-operator/terratest/go.sum @@ -513,8 +513,7 @@ golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da h1:noIWHXmPHxILtqtCOPIhS golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da/go.mod h1:NDW/Ps6MPRej6fsCIbMTohpP40sJ/P/vI1MoTEGwX90= google.golang.org/appengine v1.6.8 h1:IhEN5q69dyKagZPYMSdIjS2HqprW324FRQZJcGqPAsM= google.golang.org/appengine v1.6.8/go.mod h1:1jJ3jBArFh5pcgW8gCtRJnepW8FzD1V44FJffLiz/Ds= -google.golang.org/genproto/googleapis/rpc v0.0.0-20260406210006-6f92a3bedf2d h1:wT2n40TBqFY6wiwazVK9/iTWbsQrgk5ZfCSVFLO9LQA= -google.golang.org/genproto/googleapis/rpc v0.0.0-20260406210006-6f92a3bedf2d/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8= +google.golang.org/genproto/googleapis/rpc v0.0.0-20260414002931-afd174a4e478 h1:RmoJA1ujG+/lRGNfUnOMfhCy5EipVMyvUE+KNbPbTlw= google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af h1:+5/Sw3GsDNlEmu7TfklWKPdQ0Ykja5VEmq2i817+jbI= google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= diff --git a/go/go.work b/go/go.work index 1e16b66388..078d88fe76 100644 --- a/go/go.work +++ b/go/go.work @@ -1,4 +1,4 @@ -go 1.26.5 +go 1.26.6 use ( ./ai-proxy diff --git a/go/helm-test/go.mod b/go/helm-test/go.mod index c9f0fb052d..79b7c1e87b 100644 --- a/go/helm-test/go.mod +++ b/go/helm-test/go.mod @@ -1,6 +1,6 @@ module github.com/pluralsh/console/go/helm-test -go 1.26.5 +go 1.26.6 require ( github.com/onsi/ginkgo/v2 v2.28.1 diff --git a/go/helmchartutils/go.mod b/go/helmchartutils/go.mod index 45bcdf3118..a14fc5f943 100644 --- a/go/helmchartutils/go.mod +++ b/go/helmchartutils/go.mod @@ -1,3 +1,3 @@ module github.com/pluralsh/console/go/helmchartutils -go 1.26.5 +go 1.26.6 diff --git a/go/kubernetes-agent/api/go.mod b/go/kubernetes-agent/api/go.mod index 544696208b..f22f41b579 100644 --- a/go/kubernetes-agent/api/go.mod +++ b/go/kubernetes-agent/api/go.mod @@ -1,6 +1,6 @@ module github.com/pluralsh/console/go/kubernetes-agent/api -go 1.26.5 +go 1.26.6 require ( github.com/Yiling-J/theine-go v0.6.0 diff --git a/go/kubernetes-agent/kas/go.mod b/go/kubernetes-agent/kas/go.mod index 7545255a5e..3760001404 100644 --- a/go/kubernetes-agent/kas/go.mod +++ b/go/kubernetes-agent/kas/go.mod @@ -1,6 +1,6 @@ module github.com/pluralsh/console/go/kubernetes-agent -go 1.26.5 +go 1.26.6 replace ( github.com/pluralsh/console/go/client => ../../client diff --git a/go/nexus/go.mod b/go/nexus/go.mod index d012eefe83..b478ab520a 100644 --- a/go/nexus/go.mod +++ b/go/nexus/go.mod @@ -1,6 +1,6 @@ module github.com/pluralsh/console/go/nexus -go 1.26.5 +go 1.26.6 require ( github.com/bytedance/sonic v1.15.0 diff --git a/go/observability-proxy/go.mod b/go/observability-proxy/go.mod index e01023e260..74d18cc0d2 100644 --- a/go/observability-proxy/go.mod +++ b/go/observability-proxy/go.mod @@ -1,6 +1,6 @@ module github.com/pluralsh/console/go/observability-proxy -go 1.26.5 +go 1.26.6 require ( golang.org/x/sync v0.22.0 diff --git a/go/oci-auth/go.mod b/go/oci-auth/go.mod index a535253269..ba29efd2ae 100644 --- a/go/oci-auth/go.mod +++ b/go/oci-auth/go.mod @@ -1,6 +1,6 @@ module github.com/pluralsh/console/go/oci-auth -go 1.26.5 +go 1.26.6 require ( github.com/Azure/azure-sdk-for-go/sdk/azcore v1.21.1 diff --git a/go/polly/go.mod b/go/polly/go.mod index a2d42c712f..105f1272b7 100644 --- a/go/polly/go.mod +++ b/go/polly/go.mod @@ -1,6 +1,6 @@ module github.com/pluralsh/console/go/polly -go 1.26.5 +go 1.26.6 require ( dario.cat/mergo v1.0.2 diff --git a/go/sign/go.mod b/go/sign/go.mod index 33c828bac5..b569caafa7 100644 --- a/go/sign/go.mod +++ b/go/sign/go.mod @@ -1,6 +1,6 @@ module github.com/pluralsh/console/go/sign -go 1.26.5 +go 1.26.6 require github.com/aws/aws-lambda-go v1.54.0 diff --git a/go/tools/go.mod b/go/tools/go.mod index 976bf151d7..09f8448af9 100644 --- a/go/tools/go.mod +++ b/go/tools/go.mod @@ -1,6 +1,6 @@ module github.com/pluralsh/console/go/tools -go 1.26.5 +go 1.26.6 require ( github.com/99designs/gqlgen v0.17.78 diff --git a/lib/cloud_query/toolquery.pb.ex b/lib/cloud_query/toolquery.pb.ex index 7cf14c6a84..1c96cf6319 100644 --- a/lib/cloud_query/toolquery.pb.ex +++ b/lib/cloud_query/toolquery.pb.ex @@ -639,6 +639,30 @@ defmodule Toolquery.RunLuaOutput do field :result_json, 1, type: :string, json_name: "resultJson" end +defmodule Toolquery.RunPythonInput do + @moduledoc false + + use Protobuf, + full_name: "toolquery.RunPythonInput", + protoc_gen_elixir_version: "0.16.0", + syntax: :proto3 + + field :script, 1, type: :string + field :input_json, 2, type: :string, json_name: "inputJson" +end + +defmodule Toolquery.RunPythonOutput do + @moduledoc false + + use Protobuf, + full_name: "toolquery.RunPythonOutput", + protoc_gen_elixir_version: "0.16.0", + syntax: :proto3 + + field :result_json, 1, type: :string, json_name: "resultJson" + field :stdout, 2, type: :string +end + defmodule Toolquery.ToolQuery.Service do @moduledoc false @@ -657,6 +681,8 @@ defmodule Toolquery.ToolQuery.Service do rpc :InvokeLambda, Toolquery.InvokeLambdaInput, Toolquery.InvokeLambdaOutput rpc :RunLua, Toolquery.RunLuaInput, Toolquery.RunLuaOutput + + rpc :RunPython, Toolquery.RunPythonInput, Toolquery.RunPythonOutput end defmodule Toolquery.ToolQuery.Stub do diff --git a/lib/console/ai/tools/workbench/python.ex b/lib/console/ai/tools/workbench/python.ex new file mode 100644 index 0000000000..96c7ff3af8 --- /dev/null +++ b/lib/console/ai/tools/workbench/python.ex @@ -0,0 +1,39 @@ +defmodule Console.AI.Tools.Workbench.Python do + use Console.AI.Tools.Workbench.Base + alias CloudQuery.Client + alias Console.AI.Tools.Workbench.Output + alias Toolquery.ToolQuery.Stub + alias Toolquery.{RunPythonInput, RunPythonOutput} + + embedded_schema do + field :explanation, :string + field :code, :string + field :input, :map + end + + @json_schema Console.priv_file!("tools/workbench/python.json") |> Jason.decode!() + + def name(), do: "python_sandbox" + + def description(), + do: + "Execute a sandboxed Monty Python script for precise computation. It supports a limited Python subset only and cannot access the filesystem, environment, network, subprocesses, pip, packages, or host tools." + + def json_schema(), do: @json_schema + + def changeset(model, attrs) do + model + |> cast(attrs, [:explanation, :code, :input]) + |> validate_required([:explanation, :code]) + end + + def implement(%__MODULE__{code: code, input: input}) do + with {:ok, client} <- Client.connect(), + {:ok, input_json} <- Jason.encode(input || %{}), + request = %RunPythonInput{script: code, input_json: input_json}, + {:ok, %RunPythonOutput{result_json: result_json, stdout: stdout}} <- Stub.run_python(client, request, Client.cloud_query_rpc_opts()), + {:ok, result} <- Jason.decode(result_json) do + Output.json(%{result: result, stdout: stdout}) + end + end +end diff --git a/lib/console/ai/workbench/engine.ex b/lib/console/ai/workbench/engine.ex index 9e6603a372..eb46734da6 100644 --- a/lib/console/ai/workbench/engine.ex +++ b/lib/console/ai/workbench/engine.ex @@ -30,6 +30,7 @@ defmodule Console.AI.Workbench.Engine do } alias Console.AI.Tools.Workbench.{ Codemode, + Python, Complete, Subagents, Subagent, @@ -335,6 +336,7 @@ defmodule Console.AI.Workbench.Engine do %Subagent{subagents: subagents}, %FetchNotes{job: job}, %Codemode{tools: []}, + Python, Notes, Complete, ] ++ type_tools(job) diff --git a/priv/prompts/workbench/job.md.eex b/priv/prompts/workbench/job.md.eex index 207f6e3d32..d695ae767a 100644 --- a/priv/prompts/workbench/job.md.eex +++ b/priv/prompts/workbench/job.md.eex @@ -164,7 +164,7 @@ do not run them simultaneously, since they won't be able to communicate with eac In general, you can call any of the subagents and tools as many times as you need to accomplish your task. -For any tool you delegate to, you should have them give as close to exact raw data as possible, rather than computing on it. When you need an exact, deterministic computation, always use the `workbench_lua` tool with the appropriate Lua code to get the result for you. +For any tool you delegate to, you should have them give as close to exact raw data as possible, rather than computing on it. When you need an exact, deterministic computation, always use `python_sandbox` tool with an appropriate supported script to get the result for you. For instance, @@ -178,6 +178,14 @@ output.total = total (you must always set output to see the returned values). +For Python, `input` is the optional JSON object supplied to the tool and `output` must remain a JSON-serializable dictionary: + +```python +output = {"total": sum(input["values"])} +``` + +Monty is a limited Python subset, not CPython; use Lua when its supported libraries are a better fit, and do not expect files, environment variables, network access, subprocesses, pip, third-party packages, or host tools from Python. + That said, you *should* avoid repeatedly calling the notes tool unless you're meaningfully updating your working theory, plan, or conclusion. This will spam users and is unnneeded. In general, you should be focused on always making progress, don't repeatedly call tools that get no where important. diff --git a/priv/tools/workbench/python.json b/priv/tools/workbench/python.json new file mode 100644 index 0000000000..187136a61a --- /dev/null +++ b/priv/tools/workbench/python.json @@ -0,0 +1,22 @@ +{ + "type": "object", + "description": "Execute a sandboxed Monty Python script for precise computation. Monty supports a limited Python subset for expressions, variables, collections, control flow, functions, and selected built-ins. The JSON object `input` is available to the script as a global dictionary; write structured results to the global `output` dictionary and use print() for diagnostic output. Filesystem, environment, network, subprocesses, pip, packages, and host tools are unavailable.", + "properties": { + "code": { + "type": "string", + "description": "Monty-compatible Python source code. Read the optional global input dictionary, assign JSON-serializable fields on the global output dictionary, and use print() for diagnostics." + }, + "explanation": { + "type": "string", + "description": "Explanation of why you're running this script, such as 'I need to calculate the exact p95 latency from these samples.'." + }, + "input": { + "type": "object", + "description": "Optional JSON object exposed to the script as the global input dictionary." + } + }, + "required": [ + "code", + "explanation" + ] +} diff --git a/test/console/ai/tools/workbench/python_test.exs b/test/console/ai/tools/workbench/python_test.exs new file mode 100644 index 0000000000..d04955aefc --- /dev/null +++ b/test/console/ai/tools/workbench/python_test.exs @@ -0,0 +1,69 @@ +defmodule Console.AI.Tools.Workbench.PythonTest do + use Console.DataCase, async: false + use Mimic + + alias CloudQuery.Client + alias Console.AI.Tools.Workbench.{Output, Python} + alias Toolquery.{RunPythonOutput} + alias Toolquery.ToolQuery.Stub + + setup :set_mimic_global + + test "exposes the Python tool contract" do + assert Python.name() == "python_sandbox" + assert %{"input" => %{"type" => "object"}} = Python.json_schema()["properties"] + assert Python.json_schema()["required"] == ["code", "explanation"] + end + + test "requires an explanation and Python code" do + changeset = Python.changeset(%Python{}, %{}) + + refute changeset.valid? + assert %{code: ["can't be blank"], explanation: ["can't be blank"]} = errors_on(changeset) + end + + test "runs Monty Python with JSON input and returns serialized output and stdout" do + expect(Client, :connect, fn -> {:ok, :channel} end) + + expect(Stub, :run_python, fn :channel, request, opts -> + assert request.script == "output['total'] = input['first'] + input['second']" + assert request.input_json == ~s({"first":20,"second":22}) + assert opts == Client.cloud_query_rpc_opts() + + {:ok, %RunPythonOutput{result_json: ~s({"total":42}), stdout: "calculated total\n"}} + end) + + assert {:ok, output} = + Python.implement(%Python{ + code: "output['total'] = input['first'] + input['second']", + input: %{"first" => 20, "second" => 22} + }) + + assert Jason.decode!(output) == %{"result" => %{"total" => 42}, "stdout" => "calculated total\n"} + end + + test "defaults omitted input to an empty JSON object" do + expect(Client, :connect, fn -> {:ok, :channel} end) + + expect(Stub, :run_python, fn :channel, request, _opts -> + assert request.input_json == "{}" + {:ok, %RunPythonOutput{result_json: "{}", stdout: ""}} + end) + + assert {:ok, output} = Python.implement(%Python{code: "output = {}"}) + assert Jason.decode!(output) == %{"result" => %{}, "stdout" => ""} + end + + test "bounds oversized Python output" do + expect(Client, :connect, fn -> {:ok, :channel} end) + + expect(Stub, :run_python, fn :channel, _request, _opts -> + result = Jason.encode!(%{"value" => String.duplicate("x", Output.max_bytes())}) + {:ok, %RunPythonOutput{result_json: result, stdout: ""}} + end) + + assert {:ok, output} = Python.implement(%Python{code: "output = {}"}) + assert byte_size(output) == Output.max_bytes() + assert output =~ "output truncated at 50 KiB" + end +end