From fd25a83f7ca8f2a99cc3d909b69ce8159971e127 Mon Sep 17 00:00:00 2001 From: Mike Cohen Date: Wed, 16 Aug 2023 15:38:48 +1000 Subject: [PATCH 1/2] Bugfix: Allow serve url to be set without materializing Some servers have locked down egress and can not materialize the tool themselves (to calculate the tool hash). This means that the user is not able to set an alternative serve url from these servers without triggering a materialize (which fails) This PR does not require the server to materialize the tool when updating the serve url. This allows a user to manually upload a tool (thereby calculating its hash) then change the serve url, preserving the hash - all without the server fetching the tool from the internet. --- actions/vql_test.go | 42 ++++---- api/tools.go | 2 +- api/upload.go | 6 +- docs/references/vql.yaml | 8 +- file_store/test_utils/retry.go | 99 +++++++++++++++++++ .../notebooks/notebook-cell-renderer.jsx | 29 +++++- .../src/components/tools/tool-viewer.jsx | 11 ++- services/inventory/inventory.go | 6 ++ vtesting/helpers.go | 6 +- 9 files changed, 175 insertions(+), 34 deletions(-) create mode 100644 file_store/test_utils/retry.go diff --git a/actions/vql_test.go b/actions/vql_test.go index bc60a026569..7bfc0372720 100644 --- a/actions/vql_test.go +++ b/actions/vql_test.go @@ -130,27 +130,31 @@ func (self *ClientVQLTestSuite) TestMaxRows() { } func (self *ClientVQLTestSuite) TestMaxWait() { - resp := responder.TestResponderWithFlowId(self.ConfigObj, "TestMaxWait") + assert.True(self.T(), test_utils.Retry(self.T(), 5, time.Millisecond*500, + func(r *test_utils.R) { + resp := responder.TestResponderWithFlowId(self.ConfigObj, "TestMaxWait") + + actions.VQLClientAction{}.StartQuery(self.ConfigObj, self.Sm.Ctx, resp, + &actions_proto.VQLCollectorArgs{ + MaxRow: 1000, + MaxWait: 1, + Query: []*actions_proto.VQLRequest{ + { + Name: "Query", + VQL: "SELECT sleep(ms=400) FROM range(end=4)", + }, + }, + }) - actions.VQLClientAction{}.StartQuery(self.ConfigObj, self.Sm.Ctx, resp, - &actions_proto.VQLCollectorArgs{ - MaxRow: 1000, - MaxWait: 1, - Query: []*actions_proto.VQLRequest{ - { - Name: "Query", - VQL: "SELECT sleep(ms=400) FROM range(end=4)", - }, - }, - }) + var responses []*crypto_proto.VeloMessage - var responses []*crypto_proto.VeloMessage - vtesting.WaitUntil(5*time.Second, self.T(), func() bool { - responses = resp.Drain.Messages() - payloads := getResponsePacketCounts(responses) - // Message will be split into 2 packets 2 messages in each - return len(payloads) == 2 && payloads[0] == 2 && payloads[1] == 2 - }) + vtesting.WaitUntil(5*time.Second, r, func() bool { + responses = resp.Drain.Messages() + payloads := getResponsePacketCounts(responses) + // Message will be split into 2 packets 2 messages in each + return len(payloads) == 2 && payloads[0] == 2 && payloads[1] == 2 + }) + })) } func getLogs(responses []*crypto_proto.VeloMessage) string { diff --git a/api/tools.go b/api/tools.go index 9d97bee6f50..1aeda220476 100644 --- a/api/tools.go +++ b/api/tools.go @@ -63,7 +63,7 @@ func (self *ApiServer) SetToolInfo(ctx context.Context, return nil, Status(self.verbose, err) } - // Clear internally managed tools the user should not be allowed + // Clear internally managed fields the user should not be allowed // to set. in.Versions = nil in.ServeUrl = "" diff --git a/api/upload.go b/api/upload.go index d1957367550..d9a6c8e8942 100644 --- a/api/upload.go +++ b/api/upload.go @@ -209,8 +209,10 @@ func formUploadHandler() http.Handler { org_config_obj, form_desc.Filename) pathspec, file_store_factory, err := path_manager.Path() - returnError(w, 403, fmt.Sprintf("Error: %v", err)) - return + if err != nil { + returnError(w, 403, fmt.Sprintf("Error: %v", err)) + return + } form_desc.Url = path_manager.URL() diff --git a/docs/references/vql.yaml b/docs/references/vql.yaml index 07fe7726158..e324fdb924d 100644 --- a/docs/references/vql.yaml +++ b/docs/references/vql.yaml @@ -5421,7 +5421,7 @@ type: string description: SQL Connection String - name: file - type: string + type: accessors.OSPath description: Required if using sqlite driver - name: accessor type: string @@ -5439,7 +5439,7 @@ type: Plugin args: - name: file - type: string + type: accessors.OSPath required: true - name: accessor type: string @@ -6560,7 +6560,7 @@ type: Plugin args: - name: filename - type: string + type: accessors.OSPath description: CSV files to open required: true - name: accessor @@ -6578,7 +6578,7 @@ type: Plugin args: - name: filename - type: string + type: accessors.OSPath description: CSV files to open required: true - name: accessor diff --git a/file_store/test_utils/retry.go b/file_store/test_utils/retry.go new file mode 100644 index 00000000000..c580938ef6b --- /dev/null +++ b/file_store/test_utils/retry.go @@ -0,0 +1,99 @@ +// Copyright 2019 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// This file from +// https://github.com/googlecloudplatform/golang-samples/blob/5465a7a4f5b2/internal/testutil/retry.go + +package test_utils + +import ( + "bytes" + "fmt" + "path/filepath" + "runtime" + "strconv" + "testing" + "time" +) + +// Retry runs function f for up to maxAttempts times until f returns successfully, and reports whether f was run successfully. +// It will sleep for the given period between invocations of f. +// Use the provided *testutil.R instead of a *testing.T from the function. +func Retry(t *testing.T, maxAttempts int, sleep time.Duration, f func(r *R)) bool { + for attempt := 1; attempt <= maxAttempts; attempt++ { + r := &R{Attempt: attempt, log: &bytes.Buffer{}} + + f(r) + + if !r.failed { + if r.log.Len() != 0 { + t.Logf("Success after %d attempts:%s", attempt, r.log.String()) + } + return true + } + + if attempt == maxAttempts { + t.Logf("FAILED after %d attempts:%s", attempt, r.log.String()) + t.Fail() + } + + time.Sleep(sleep) + } + return false +} + +// R is passed to each run of a flaky test run, manages state and accumulates log statements. +type R struct { + // The number of current attempt. + Attempt int + + failed bool + log *bytes.Buffer +} + +// Fail marks the run as failed, and will retry once the function returns. +func (r *R) Fail() { + r.failed = true +} + +// Errorf is equivalent to Logf followed by Fail. +func (r *R) Errorf(s string, v ...interface{}) { + r.logf(s, v...) + r.Fail() +} + +func (r *R) Fatalf(s string, v ...interface{}) { + r.logf(s, v...) + r.Fail() +} + +// Logf formats its arguments and records it in the error log. +// The text is only printed for the final unsuccessful run or the first successful run. +func (r *R) Logf(s string, v ...interface{}) { + r.logf(s, v...) +} + +func (r *R) logf(s string, v ...interface{}) { + fmt.Fprint(r.log, "\n") + fmt.Fprint(r.log, lineNumber()) + fmt.Fprintf(r.log, s, v...) +} + +func lineNumber() string { + _, file, line, ok := runtime.Caller(3) // logf, public func, user function + if !ok { + return "" + } + return filepath.Base(file) + ":" + strconv.Itoa(line) + ": " +} diff --git a/gui/velociraptor/src/components/notebooks/notebook-cell-renderer.jsx b/gui/velociraptor/src/components/notebooks/notebook-cell-renderer.jsx index ed53fca2be5..b8672266bb5 100644 --- a/gui/velociraptor/src/components/notebooks/notebook-cell-renderer.jsx +++ b/gui/velociraptor/src/components/notebooks/notebook-cell-renderer.jsx @@ -176,11 +176,13 @@ export default class NotebookCellRenderer extends React.Component { componentDidMount() { this.source = CancelToken.source(); + this.update_source = CancelToken.source(); this.fetchCellContents(); } componentWillUnmount() { this.source.cancel(); + this.update_source.cancel(); } componentDidUpdate = (prevProps, prevState, rootNode) => { @@ -288,10 +290,17 @@ export default class NotebookCellRenderer extends React.Component { recalculate = () => { let cell = this.state.cell; - cell.output = "Loading"; + cell.output = T("Loading"); cell.timestamp = 0; + cell.calculating = true; + cell.messages = []; + this.setState({cell: cell}); + // Reset any inflight calls. + this.update_source.cancel(); + this.update_source = CancelToken.source(); + api.post('v1/UpdateNotebookCell', { notebook_id: this.props.notebook_id, cell_id: this.state.cell.cell_id, @@ -299,8 +308,9 @@ export default class NotebookCellRenderer extends React.Component { env: this.state.cell.env, currently_editing: false, input: this.state.cell.input, - }, this.source.token).then( (response) => { + }, this.update_source.token).then( (response) => { if (response.cancel) { + this.setState({currently_editing: false}); return; } @@ -310,6 +320,7 @@ export default class NotebookCellRenderer extends React.Component { api.error(response.data.error); keep_editing = true; } + this.setState({cell: response.data, currently_editing: keep_editing}); }); @@ -334,9 +345,16 @@ export default class NotebookCellRenderer extends React.Component { saveCell = () => { let cell = this.state.cell; cell.input = this.state.ace.getValue(); - cell.output = "Loading"; + cell.output = T("Loading"); cell.timestamp = 0; - this.setState({cell: cell}); + cell.calculating = true; + cell.messages = []; + + this.setState({cell: cell, currently_editing: false}); + + // Reset any inflight calls. + this.update_source.cancel(); + this.update_source = CancelToken.source(); api.post('v1/UpdateNotebookCell', { notebook_id: this.props.notebook_id, @@ -345,8 +363,9 @@ export default class NotebookCellRenderer extends React.Component { env: this.state.cell.env, currently_editing: false, input: cell.input, - }, this.source.token).then( (response) => { + }, this.update_source.token).then( (response) => { if (response.cancel) { + this.setState({currently_editing: false}); return; } diff --git a/gui/velociraptor/src/components/tools/tool-viewer.jsx b/gui/velociraptor/src/components/tools/tool-viewer.jsx index 17b606e7043..ebd4627d49d 100644 --- a/gui/velociraptor/src/components/tools/tool-viewer.jsx +++ b/gui/velociraptor/src/components/tools/tool-viewer.jsx @@ -128,11 +128,18 @@ export default class ToolViewer extends React.Component { setServeUrl = url=>{ let tool = Object.assign({}, this.state.tool); tool.url = this.state.remote_url; - tool.hash = ""; + + // Preserve the hash if we already know it. This avoids us + // having to refetch the file all the time. + // tool.hash = ""; tool.filename = ""; tool.github_project = ""; tool.serve_locally = false; - tool.materialize = true; + + // Do not force a materialize - it is possible that the server + // has no egress access and it is not possible to materialize + // the tool from the server. + tool.materialize = false; this.setToolInfo(tool); }; diff --git a/services/inventory/inventory.go b/services/inventory/inventory.go index 62d29a2a244..f4c30636ff1 100644 --- a/services/inventory/inventory.go +++ b/services/inventory/inventory.go @@ -483,6 +483,8 @@ func (self *InventoryService) AddTool( tool := proto.Clone(tool_request).(*artifacts_proto.Tool) tool.FilestorePath = paths.ObfuscateName(config_obj, tool.Name) + // No client config so we dont know any server urls - therefore we + // can not serve locally at all. if tool.ServeLocally && config_obj.Client == nil { tool.ServeLocally = false } @@ -492,6 +494,10 @@ func (self *InventoryService) AddTool( return errors.New("No server URLs configured!") } tool.ServeUrl = config_obj.Client.ServerUrls[0] + "public/" + tool.FilestorePath + } else { + // If we dont serve the tool, the clients will directly get + // the tool from its upstream URL. + tool.ServeUrl = tool.Url } // Set the filename to something sensible so it is always valid. diff --git a/vtesting/helpers.go b/vtesting/helpers.go index 0a5441fc7d2..41e58363ed6 100644 --- a/vtesting/helpers.go +++ b/vtesting/helpers.go @@ -31,6 +31,10 @@ import ( "www.velocidex.com/golang/vfilter/types" ) +type TestFailer interface { + Fatalf(format string, args ...interface{}) +} + func ReadFile(t *testing.T, filename string) []byte { result, err := ioutil.ReadFile(filename) if err != nil { @@ -62,7 +66,7 @@ func ContainsString(expected string, watched []string) bool { return false } -func WaitUntil(deadline time.Duration, t *testing.T, cb func() bool) { +func WaitUntil(deadline time.Duration, t TestFailer, cb func() bool) { end_time := time.Now().Add(deadline) for end_time.After(time.Now()) { From bc42c7f12cfaae17e85176aad80b0c44eba68f77 Mon Sep 17 00:00:00 2001 From: Mike Cohen Date: Wed, 16 Aug 2023 15:57:19 +1000 Subject: [PATCH 2/2] Fix tests --- artifacts/testdata/server/testcases/binary_blobs.out.yaml | 2 +- artifacts/testdata/server/testcases/tools.out.yaml | 1 + services/sanity/fixtures/TestUpgradeTools.golden | 2 ++ 3 files changed, 4 insertions(+), 1 deletion(-) diff --git a/artifacts/testdata/server/testcases/binary_blobs.out.yaml b/artifacts/testdata/server/testcases/binary_blobs.out.yaml index 6fb264daf4d..ecef6ec7479 100644 --- a/artifacts/testdata/server/testcases/binary_blobs.out.yaml +++ b/artifacts/testdata/server/testcases/binary_blobs.out.yaml @@ -37,7 +37,7 @@ SELECT * FROM switch( b={SELECT Complete FROM execve(argv=["rm", "-f", "/tmp/aut "materialize": false, "artifact": "", "filestore_path": "351b4f6d59a4266cc7a2eab9cedf959eb6a4c924746044e6edeabdd1a477643e", - "serve_url": "", + "serve_url": "https://storage.googleapis.com/go.velocidex.com/winpmem_v3.3.rc3.exe", "serve_path": "", "filename": "winpmem_v3.3.rc3.exe", "hash": "", diff --git a/artifacts/testdata/server/testcases/tools.out.yaml b/artifacts/testdata/server/testcases/tools.out.yaml index 028f141925d..2cc0416ea8f 100644 --- a/artifacts/testdata/server/testcases/tools.out.yaml +++ b/artifacts/testdata/server/testcases/tools.out.yaml @@ -31,6 +31,7 @@ SELECT inventory_add(tool='Autorun_amd64', url='https://storage.googleapis.com/g "url": "https://storage.googleapis.com/go.velocidex.com/autorunsc.exe", "admin_override": true, "filestore_path": "affbb298885d99a7569ef9ec463bbfbf9db112a11cec65cf78ec8657ab78950a", + "serve_url": "https://storage.googleapis.com/go.velocidex.com/autorunsc.exe", "filename": "autorunsc_x64.exe", "hash": "083d7eee4ed40a3e5a35675503b0b6be0cb627b4cb1009d185a558a805f64153", "versions": [ diff --git a/services/sanity/fixtures/TestUpgradeTools.golden b/services/sanity/fixtures/TestUpgradeTools.golden index 6a7768a4f61..715cc34c6c2 100644 --- a/services/sanity/fixtures/TestUpgradeTools.golden +++ b/services/sanity/fixtures/TestUpgradeTools.golden @@ -6,6 +6,7 @@ "filename": "www.company.com", "filestore_path": "95b097fd386adb06d90ff35d45f18fb507675b076522e7993b5be4a695f4247a", "name": "Tool1", + "serve_url": "https://www.company.com", "url": "https://www.company.com" }, { @@ -13,6 +14,7 @@ "filename": "www.example2.com", "filestore_path": "9a0dbc5804c7e06fc2e371a936bbe1d90b206aa9b8fbad7601f4a7f70fb596e1", "name": "Tool2", + "serve_url": "https://www.example2.com/", "url": "https://www.example2.com/" } ],