Skip to content

Commit

Permalink
pkg/cwhub: download data assets to temporary files to avoid partial f…
Browse files Browse the repository at this point in the history
…etch (#2879)
  • Loading branch information
mmetc authored Mar 8, 2024
1 parent 1eab943 commit 6c5e8af
Show file tree
Hide file tree
Showing 5 changed files with 33 additions and 17 deletions.
22 changes: 18 additions & 4 deletions pkg/cwhub/dataset.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import (
"io"
"net/http"
"os"
"path/filepath"
"time"

"github.com/sirupsen/logrus"
Expand All @@ -31,19 +32,32 @@ func downloadFile(url string, destPath string) error {
return fmt.Errorf("bad http code %d for %s", resp.StatusCode, url)
}

file, err := os.Create(destPath)
tmpFile, err := os.CreateTemp(filepath.Dir(destPath), filepath.Base(destPath)+".*.tmp")
if err != nil {
return err
}
defer file.Close()

tmpFileName := tmpFile.Name()
defer func() {
tmpFile.Close()
os.Remove(tmpFileName)
}()

// avoid reading the whole file in memory
_, err = io.Copy(file, resp.Body)
_, err = io.Copy(tmpFile, resp.Body)
if err != nil {
return err
}

if err = file.Sync(); err != nil {
if err = tmpFile.Sync(); err != nil {
return err
}

if err = tmpFile.Close(); err != nil {
return err
}

if err = os.Rename(tmpFileName, destPath); err != nil {
return err
}

Expand Down
8 changes: 4 additions & 4 deletions pkg/cwhub/dataset_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ func TestDownloadFile(t *testing.T) {
httpmock.Activate()
defer httpmock.DeactivateAndReset()

//OK
// OK
httpmock.RegisterResponder(
"GET",
"https://example.com/xx",
Expand All @@ -36,15 +36,15 @@ func TestDownloadFile(t *testing.T) {
assert.Equal(t, "example content oneoneone", string(content))
require.NoError(t, err)

//bad uri
// bad uri
err = downloadFile("https://zz.com", examplePath)
require.Error(t, err)

//404
// 404
err = downloadFile("https://example.com/x", examplePath)
require.Error(t, err)

//bad target
// bad target
err = downloadFile("https://example.com/xx", "")
require.Error(t, err)
}
2 changes: 1 addition & 1 deletion pkg/cwhub/errors.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ import (
)

var (
// ErrNilRemoteHub is returned when the remote hub configuration is not provided to the NewHub constructor.
// ErrNilRemoteHub is returned when trying to download with a local-only configuration.
ErrNilRemoteHub = errors.New("remote hub configuration is not provided. Please report this issue to the developers")
)

Expand Down
3 changes: 2 additions & 1 deletion pkg/cwhub/hub.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package cwhub
import (
"bytes"
"encoding/json"
"errors"
"fmt"
"io"
"os"
Expand Down Expand Up @@ -34,7 +35,7 @@ func (h *Hub) GetDataDir() string {
// All download operations (including updateIndex) return ErrNilRemoteHub if the remote configuration is not set.
func NewHub(local *csconfig.LocalHubCfg, remote *RemoteHubCfg, updateIndex bool, logger *logrus.Logger) (*Hub, error) {
if local == nil {
return nil, fmt.Errorf("no hub configuration found")
return nil, errors.New("no hub configuration found")
}

if logger == nil {
Expand Down
15 changes: 8 additions & 7 deletions pkg/cwhub/sync.go
Original file line number Diff line number Diff line change
Expand Up @@ -77,9 +77,9 @@ func (h *Hub) getItemFileInfo(path string, logger *logrus.Logger) (*itemFileInfo
if strings.HasPrefix(path, hubDir) {
logger.Tracef("in hub dir")

//.../hub/parsers/s00-raw/crowdsec/skip-pretag.yaml
//.../hub/scenarios/crowdsec/ssh_bf.yaml
//.../hub/profiles/crowdsec/linux.yaml
// .../hub/parsers/s00-raw/crowdsec/skip-pretag.yaml
// .../hub/scenarios/crowdsec/ssh_bf.yaml
// .../hub/profiles/crowdsec/linux.yaml
if len(subs) < 4 {
return nil, fmt.Errorf("path is too short: %s (%d)", path, len(subs))
}
Expand All @@ -93,13 +93,14 @@ func (h *Hub) getItemFileInfo(path string, logger *logrus.Logger) (*itemFileInfo
}
} else if strings.HasPrefix(path, installDir) { // we're in install /etc/crowdsec/<type>/...
logger.Tracef("in install dir")

if len(subs) < 3 {
return nil, fmt.Errorf("path is too short: %s (%d)", path, len(subs))
}
///.../config/parser/stage/file.yaml
///.../config/postoverflow/stage/file.yaml
///.../config/scenarios/scenar.yaml
///.../config/collections/linux.yaml //file is empty
// .../config/parser/stage/file.yaml
// .../config/postoverflow/stage/file.yaml
// .../config/scenarios/scenar.yaml
// .../config/collections/linux.yaml //file is empty
ret = &itemFileInfo{
inhub: false,
fname: subs[len(subs)-1],
Expand Down

0 comments on commit 6c5e8af

Please sign in to comment.