Skip to content

Commit

Permalink
fix 0xrawsec#122: Use context.Context in forwarder.go
Browse files Browse the repository at this point in the history
  • Loading branch information
@qjerome
qjerome committed Jul 25, 2022
1 parent daf1602 commit 28e1e43
Show file tree
Hide file tree
Showing 8 changed files with 274 additions and 248 deletions.
2 changes: 1 addition & 1 deletion .github/coverage/badge.svg
View file
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
100 changes: 50 additions & 50 deletions .github/coverage/coverage.txt
View file
Original file line number Diff line number Diff line change
@@ -1,12 +1,12 @@
ok github.com/0xrawsec/whids/api 167.841s coverage: 66.8% of statements
ok github.com/0xrawsec/whids/event 43.388s coverage: 75.3% of statements
ok github.com/0xrawsec/whids/hids 39.421s coverage: 52.5% of statements
ok github.com/0xrawsec/whids/hids/sysinfo 0.457s coverage: 95.2% of statements
ok github.com/0xrawsec/whids/ioc 26.004s coverage: 73.3% of statements
ok github.com/0xrawsec/whids/logger 43.282s coverage: 76.7% of statements
ok github.com/0xrawsec/whids/sysmon 6.222s coverage: 83.1% of statements
ok github.com/0xrawsec/whids/utils 11.752s coverage: 18.1% of statements
ok github.com/0xrawsec/whids/utils/command 0.541s coverage: 100.0% of statements
ok github.com/0xrawsec/whids/api 221.707s coverage: 66.8% of statements
ok github.com/0xrawsec/whids/event 83.783s coverage: 75.3% of statements
ok github.com/0xrawsec/whids/hids 72.862s coverage: 49.0% of statements
ok github.com/0xrawsec/whids/hids/sysinfo 0.722s coverage: 95.2% of statements
ok github.com/0xrawsec/whids/ioc 50.828s coverage: 73.3% of statements
ok github.com/0xrawsec/whids/logger 73.084s coverage: 76.7% of statements
ok github.com/0xrawsec/whids/sysmon 18.651s coverage: 83.1% of statements
ok github.com/0xrawsec/whids/utils 22.620s coverage: 18.1% of statements
ok github.com/0xrawsec/whids/utils/command 1.303s coverage: 100.0% of statements
github.com/0xrawsec/whids/api/api_client.go:36: ValidateRespStatus 75.0%
github.com/0xrawsec/whids/api/api_client.go:61: ManagerIP 60.0%
github.com/0xrawsec/whids/api/api_client.go:73: DialContext 0.0%
Expand Down Expand Up @@ -53,20 +53,20 @@ github.com/0xrawsec/whids/api/endpoint.go:30: NewEndpoint 100.0%
github.com/0xrawsec/whids/api/endpoint.go:37: Validate 66.7%
github.com/0xrawsec/whids/api/endpoint.go:45: Copy 100.0%
github.com/0xrawsec/whids/api/endpoint.go:51: UpdateLastConnection 100.0%
github.com/0xrawsec/whids/api/forwarder.go:63: NewForwarder 72.7%
github.com/0xrawsec/whids/api/forwarder.go:103: LogfilePath 100.0%
github.com/0xrawsec/whids/api/forwarder.go:111: ArchiveLogs 0.0%
github.com/0xrawsec/whids/api/forwarder.go:128: PipeEvent 100.0%
github.com/0xrawsec/whids/api/forwarder.go:137: Save 84.6%
github.com/0xrawsec/whids/api/forwarder.go:164: HasQueuedEvents 100.0%
github.com/0xrawsec/whids/api/forwarder.go:174: CleanOlderQueued 94.4%
github.com/0xrawsec/whids/api/forwarder.go:204: DiskSpaceQueue 100.0%
github.com/0xrawsec/whids/api/forwarder.go:217: listLogfiles 100.0%
github.com/0xrawsec/whids/api/forwarder.go:230: ProcessQueue 61.8%
github.com/0xrawsec/whids/api/forwarder.go:300: Reset 100.0%
github.com/0xrawsec/whids/api/forwarder.go:306: Collect 90.0%
github.com/0xrawsec/whids/api/forwarder.go:332: Run 100.0%
github.com/0xrawsec/whids/api/forwarder.go:366: Close 100.0%
github.com/0xrawsec/whids/api/forwarder.go:64: NewForwarder 75.0%
github.com/0xrawsec/whids/api/forwarder.go:106: LogfilePath 100.0%
github.com/0xrawsec/whids/api/forwarder.go:114: ArchiveLogs 0.0%
github.com/0xrawsec/whids/api/forwarder.go:131: PipeEvent 100.0%
github.com/0xrawsec/whids/api/forwarder.go:140: Save 84.6%
github.com/0xrawsec/whids/api/forwarder.go:167: HasQueuedEvents 100.0%
github.com/0xrawsec/whids/api/forwarder.go:177: CleanOlderQueued 94.4%
github.com/0xrawsec/whids/api/forwarder.go:207: DiskSpaceQueue 100.0%
github.com/0xrawsec/whids/api/forwarder.go:220: listLogfiles 100.0%
github.com/0xrawsec/whids/api/forwarder.go:233: ProcessQueue 61.8%
github.com/0xrawsec/whids/api/forwarder.go:303: Reset 100.0%
github.com/0xrawsec/whids/api/forwarder.go:309: Collect 90.0%
github.com/0xrawsec/whids/api/forwarder.go:335: Run 100.0%
github.com/0xrawsec/whids/api/forwarder.go:365: Close 100.0%
github.com/0xrawsec/whids/api/log_streamer.go:18: Queue 75.0%
github.com/0xrawsec/whids/api/log_streamer.go:26: Stream 100.0%
github.com/0xrawsec/whids/api/log_streamer.go:40: Close 0.0%
Expand Down Expand Up @@ -208,12 +208,12 @@ github.com/0xrawsec/whids/hids/actions.go:94: shouldDump 100.0%
github.com/0xrawsec/whids/hids/actions.go:99: writeReader 100.0%
github.com/0xrawsec/whids/hids/actions.go:104: dumpAsJson 66.7%
github.com/0xrawsec/whids/hids/actions.go:117: dumpBinFile 100.0%
github.com/0xrawsec/whids/hids/actions.go:121: dumpFile 77.8%
github.com/0xrawsec/whids/hids/actions.go:156: listFilesFromCommandLine 100.0%
github.com/0xrawsec/whids/hids/actions.go:179: filedumpSet 55.6%
github.com/0xrawsec/whids/hids/actions.go:121: dumpFile 72.2%
github.com/0xrawsec/whids/hids/actions.go:156: listFilesFromCommandLine 81.8%
github.com/0xrawsec/whids/hids/actions.go:179: filedumpSet 44.4%
github.com/0xrawsec/whids/hids/actions.go:232: filedump 80.0%
github.com/0xrawsec/whids/hids/actions.go:242: memdump 0.0%
github.com/0xrawsec/whids/hids/actions.go:271: regdump 93.3%
github.com/0xrawsec/whids/hids/actions.go:271: regdump 26.7%
github.com/0xrawsec/whids/hids/actions.go:302: suspend_process 0.0%
github.com/0xrawsec/whids/hids/actions.go:312: kill_process 0.0%
github.com/0xrawsec/whids/hids/actions.go:325: Queue 100.0%
Expand Down Expand Up @@ -280,26 +280,26 @@ github.com/0xrawsec/whids/hids/hids.go:516: updateSystemInfo 0.0%
github.com/0xrawsec/whids/hids/hids.go:544: updateSysmon 0.0%
github.com/0xrawsec/whids/hids/hids.go:590: updateSysmonConfig 0.0%
github.com/0xrawsec/whids/hids/hids.go:650: cleanup 33.3%
github.com/0xrawsec/whids/hids/hids.go:666: IsHIDSEvent 93.8%
github.com/0xrawsec/whids/hids/hids.go:666: IsHIDSEvent 87.5%
github.com/0xrawsec/whids/hids/hids.go:700: Report 0.0%
github.com/0xrawsec/whids/hids/hids.go:727: Run 58.8%
github.com/0xrawsec/whids/hids/hids.go:850: LogStats 0.0%
github.com/0xrawsec/whids/hids/hids.go:859: Stop 68.8%
github.com/0xrawsec/whids/hids/hids.go:895: Wait 0.0%
github.com/0xrawsec/whids/hids/hids.go:900: WaitWithTimeout 0.0%
github.com/0xrawsec/whids/hids/hookdefs.go:39: hookSetImageSize 94.1%
github.com/0xrawsec/whids/hids/hids.go:727: Run 64.7%
github.com/0xrawsec/whids/hids/hids.go:847: LogStats 0.0%
github.com/0xrawsec/whids/hids/hids.go:856: Stop 68.8%
github.com/0xrawsec/whids/hids/hids.go:892: Wait 0.0%
github.com/0xrawsec/whids/hids/hids.go:897: WaitWithTimeout 0.0%
github.com/0xrawsec/whids/hids/hookdefs.go:39: hookSetImageSize 82.4%
github.com/0xrawsec/whids/hids/hookdefs.go:71: hookImageLoad 95.0%
github.com/0xrawsec/whids/hids/hookdefs.go:108: trackSysmonProcessCreate 76.1%
github.com/0xrawsec/whids/hids/hookdefs.go:108: trackSysmonProcessCreate 62.7%
github.com/0xrawsec/whids/hids/hookdefs.go:229: hookTrack 50.0%
github.com/0xrawsec/whids/hids/hookdefs.go:242: hookStats 76.4%
github.com/0xrawsec/whids/hids/hookdefs.go:242: hookStats 21.8%
github.com/0xrawsec/whids/hids/hookdefs.go:353: hookUpdateGeneScore 0.0%
github.com/0xrawsec/whids/hids/hookdefs.go:370: hookTerminator 76.9%
github.com/0xrawsec/whids/hids/hookdefs.go:370: hookTerminator 53.8%
github.com/0xrawsec/whids/hids/hookdefs.go:398: hookProcTerm 87.5%
github.com/0xrawsec/whids/hids/hookdefs.go:414: hookSelfGUID 75.0%
github.com/0xrawsec/whids/hids/hookdefs.go:448: hookFileSystemAudit 0.0%
github.com/0xrawsec/whids/hids/hookdefs.go:478: hookProcessIntegrityProcTamp 0.0%
github.com/0xrawsec/whids/hids/hookdefs.go:554: hookEnrichServices 80.6%
github.com/0xrawsec/whids/hids/hookdefs.go:632: hookEnrichAnySysmon 100.0%
github.com/0xrawsec/whids/hids/hookdefs.go:554: hookEnrichServices 77.8%
github.com/0xrawsec/whids/hids/hookdefs.go:632: hookEnrichAnySysmon 86.7%
github.com/0xrawsec/whids/hids/hookdefs.go:754: hookClipboardEvents 0.0%
github.com/0xrawsec/whids/hids/hookdefs.go:781: hookKernelFiles 0.0%
github.com/0xrawsec/whids/hids/hooks.go:23: newHookCache 100.0%
Expand All @@ -311,7 +311,7 @@ github.com/0xrawsec/whids/hids/hooks.go:84: RunHooksOn 93.8%
github.com/0xrawsec/whids/hids/hooks.go:123: getFunctionName 0.0%
github.com/0xrawsec/whids/hids/hookutils.go:13: toString 100.0%
github.com/0xrawsec/whids/hids/hookutils.go:17: toHex 66.7%
github.com/0xrawsec/whids/hids/hookutils.go:25: terminate 100.0%
github.com/0xrawsec/whids/hids/hookutils.go:25: terminate 0.0%
github.com/0xrawsec/whids/hids/hookutils.go:41: isSysmonProcessTerminate 100.0%
github.com/0xrawsec/whids/hids/hookutils.go:45: srcPIDFromEvent 0.0%
github.com/0xrawsec/whids/hids/hookutils.go:58: hasAction 0.0%
Expand All @@ -320,9 +320,9 @@ github.com/0xrawsec/whids/hids/iocs.go:17: ruleHashIoC 100.0%
github.com/0xrawsec/whids/hids/iocs.go:32: ruleDomainIoC 100.0%
github.com/0xrawsec/whids/hids/paths.go:11: EventDataPath 100.0%
github.com/0xrawsec/whids/hids/ptrack.go:41: NewProcStats 100.0%
github.com/0xrawsec/whids/hids/ptrack.go:52: UpdateNetResolve 0.0%
github.com/0xrawsec/whids/hids/ptrack.go:52: UpdateNetResolve 100.0%
github.com/0xrawsec/whids/hids/ptrack.go:61: UpdateCon 0.0%
github.com/0xrawsec/whids/hids/ptrack.go:71: ConStat 0.0%
github.com/0xrawsec/whids/hids/ptrack.go:71: ConStat 100.0%
github.com/0xrawsec/whids/hids/ptrack.go:83: NewGeneScore 100.0%
github.com/0xrawsec/whids/hids/ptrack.go:87: Update 0.0%
github.com/0xrawsec/whids/hids/ptrack.go:96: sysmonHashesToMap 100.0%
Expand All @@ -339,10 +339,10 @@ github.com/0xrawsec/whids/hids/ptrack.go:301: KernelFileFromEvent 0.0%
github.com/0xrawsec/whids/hids/ptrack.go:313: sourceGUIDFromEvent 88.9%
github.com/0xrawsec/whids/hids/ptrack.go:334: targetGUIDFromEvent 70.0%
github.com/0xrawsec/whids/hids/ptrack.go:376: NewActivityTracker 100.0%
github.com/0xrawsec/whids/hids/ptrack.go:393: delete 100.0%
github.com/0xrawsec/whids/hids/ptrack.go:393: delete 83.3%
github.com/0xrawsec/whids/hids/ptrack.go:406: freeRtn 80.0%
github.com/0xrawsec/whids/hids/ptrack.go:444: CheckDumpCountOrInc 100.0%
github.com/0xrawsec/whids/hids/ptrack.go:458: Add 100.0%
github.com/0xrawsec/whids/hids/ptrack.go:458: Add 83.3%
github.com/0xrawsec/whids/hids/ptrack.go:469: PS 0.0%
github.com/0xrawsec/whids/hids/ptrack.go:480: Blacklist 100.0%
github.com/0xrawsec/whids/hids/ptrack.go:484: IsBlacklisted 100.0%
Expand All @@ -367,16 +367,16 @@ github.com/0xrawsec/whids/hids/reports.go:104: PrepareCommands 0.0%
github.com/0xrawsec/whids/hids/stats.go:29: NewEventStats 100.0%
github.com/0xrawsec/whids/hids/stats.go:39: SinceStart 0.0%
github.com/0xrawsec/whids/hids/stats.go:43: Start 100.0%
github.com/0xrawsec/whids/hids/stats.go:48: Threshold 0.0%
github.com/0xrawsec/whids/hids/stats.go:52: Duration 0.0%
github.com/0xrawsec/whids/hids/stats.go:48: Threshold 100.0%
github.com/0xrawsec/whids/hids/stats.go:52: Duration 100.0%
github.com/0xrawsec/whids/hids/stats.go:56: Update 75.0%
github.com/0xrawsec/whids/hids/stats.go:65: Events 100.0%
github.com/0xrawsec/whids/hids/stats.go:69: Detections 0.0%
github.com/0xrawsec/whids/hids/stats.go:73: EPS 0.0%
github.com/0xrawsec/whids/hids/stats.go:81: CriticalEPS 0.0%
github.com/0xrawsec/whids/hids/stats.go:81: CriticalEPS 100.0%
github.com/0xrawsec/whids/hids/stats.go:85: DynEPS 75.0%
github.com/0xrawsec/whids/hids/stats.go:93: HasPerfIssue 38.5%
github.com/0xrawsec/whids/hids/stats.go:113: HasCriticalPerfIssue 0.0%
github.com/0xrawsec/whids/hids/stats.go:93: HasPerfIssue 69.2%
github.com/0xrawsec/whids/hids/stats.go:113: HasCriticalPerfIssue 100.0%
github.com/0xrawsec/whids/hids/sysinfo/sysinfo.go:15: RegisterEdrInfo 0.0%
github.com/0xrawsec/whids/hids/sysinfo/windows_sysinfo.go:31: NewSystemInfo 100.0%
github.com/0xrawsec/whids/ioc/ioc.go:24: FromObjects 0.0%
Expand Down Expand Up @@ -526,4 +526,4 @@ github.com/0xrawsec/whids/utils/windows.go:53: ResolveCDrive 0.0%
github.com/0xrawsec/whids/utils/windows.go:76: RegValue 0.0%
github.com/0xrawsec/whids/utils/windows.go:91: RegJoin 0.0%
github.com/0xrawsec/whids/utils/windows.go:98: RegValueToString 0.0%
total: (statements) 60.2%
total: (statements) 58.8%
51 changes: 28 additions & 23 deletions api/forwarder.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package api
import (
"bytes"
"compress/gzip"
"context"
"fmt"
"os"
"path/filepath"
Expand Down Expand Up @@ -43,12 +44,12 @@ type ForwarderConfig struct {
// Forwarder structure definition
type Forwarder struct {
sync.Mutex
sync.WaitGroup
ctx context.Context
cancel context.CancelFunc
fwdConfig *ForwarderConfig
stop chan bool
done chan bool
logfile logfile.LogFile
sleep time.Duration
closed bool

Client *ManagerClient
TimeTresh time.Duration
Expand All @@ -60,15 +61,17 @@ type Forwarder struct {

// NewForwarder creates a new Forwarder structure
// Todo: needs update with client
func NewForwarder(c *ForwarderConfig) (*Forwarder, error) {
func NewForwarder(ctx context.Context, c *ForwarderConfig) (*Forwarder, error) {
var err error

cctx, cancel := context.WithCancel(ctx)

// Initialize the Forwarder
// TODO: better organize forwarder configuration
co := Forwarder{
ctx: cctx,
cancel: cancel,
fwdConfig: c,
stop: make(chan bool),
done: make(chan bool),
sleep: time.Second,
TimeTresh: time.Second * 10,
// Writing events too quickly has a perf impact
Expand Down Expand Up @@ -330,17 +333,13 @@ func (f *Forwarder) Collect() {

// Run starts the Forwarder worker function
func (f *Forwarder) Run() {
f.Add(1)
// Process Piped Events
go func() {
// defer signal that we are done
defer func() { f.done <- true }()
defer f.Done()

timer := time.Now()
for {
select {
case <-f.stop:
return
default:
}
for f.ctx.Err() == nil {
// We have queued events so we try to send them before sending pending events
// We check if server is up not to close the current logfile if not needed
if f.HasQueuedEvents() {
Expand All @@ -364,23 +363,29 @@ func (f *Forwarder) Run() {

// Close closes the forwarder properly
func (f *Forwarder) Close() {
if f.closed {

// forwarder is already closed -> nothing to do
if f.ctx.Err() != nil {
return
}

// Close idle connections if not local
if !f.Local {
defer f.Client.Close()
}
f.stop <- true
// Waiting forwarder stopped routine is done
<-f.done
// we cancel forwarder's context
f.cancel()
// we wait for forwarding routine to terminate
f.Wait()

// we collect last events if needed
if f.EventsPiped > 0 {
f.Collect()
}

// we close logfile
if f.logfile != nil {
f.logfile.Close()
}

f.closed = true
// Close idle connections if not local
if !f.Local {
defer f.Client.Close()
}
}
37 changes: 30 additions & 7 deletions api/forwarder_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package api

import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
Expand All @@ -25,7 +26,6 @@ import (
)

var (

eventFile = "./data/events.json"
events = make([]event.EdrEvent, 0)
)
Expand Down Expand Up @@ -162,7 +162,10 @@ func TestForwarderBasic(t *testing.T) {
r.Run()

fconf.Client.Key = key
f, err := NewForwarder(&fconf)
ctx, cancel := context.WithCancel(context.Background())
defer cancel()

f, err := NewForwarder(ctx, &fconf)
if err != nil {
t.Errorf("Failed to create collector: %s", err)
t.FailNow()
Expand Down Expand Up @@ -210,7 +213,11 @@ func TestCollectorAuthFailure(t *testing.T) {

fconf.Client.Key = key
fconf.Client.ServerKey = utils.UnsafeKeyGen(DefaultKeySize)
f, err := NewForwarder(&fconf)

ctx, cancel := context.WithCancel(context.Background())
defer cancel()

f, err := NewForwarder(ctx, &fconf)
if err != nil {
t.Errorf("Failed to create collector: %s", err)
t.FailNow()
Expand Down Expand Up @@ -255,7 +262,11 @@ func TestCollectorAuthSuccess(t *testing.T) {

fconf.Client.Key = key
fconf.Client.ServerKey = serverKey
f, err := NewForwarder(&fconf)

ctx, cancel := context.WithCancel(context.Background())
defer cancel()

f, err := NewForwarder(ctx, &fconf)
if err != nil {
t.Errorf("Failed to create collector: %s", err)
t.FailNow()
Expand Down Expand Up @@ -310,7 +321,11 @@ func TestForwarderParallel(t *testing.T) {
defer jobs.Release()
defer wg.Done()
fconf.Client.Key = key
c, err := NewForwarder(&fconf)

ctx, cancel := context.WithCancel(context.Background())
defer cancel()

c, err := NewForwarder(ctx, &fconf)
if err != nil {
t.Errorf("Failed to create collector: %s", err)
t.FailNow()
Expand Down Expand Up @@ -355,7 +370,11 @@ func TestForwarderQueueBasic(t *testing.T) {

// Inititialize the forwarder
fconf.Client.Key = key
f, err := NewForwarder(&fconf)

ctx, cancel := context.WithCancel(context.Background())
defer cancel()

f, err := NewForwarder(ctx, &fconf)
if err != nil {
t.Errorf("Failed to create collector: %s", err)
t.FailNow()
Expand Down Expand Up @@ -412,8 +431,12 @@ func TestForwarderCleanup(t *testing.T) {

// Change rotation interval not to create unexpected number of files
fconf.Logging.RotationInterval = time.Hour

ctx, cancel := context.WithCancel(context.Background())
defer cancel()

// Inititialize the forwarder
f, err := NewForwarder(&fconf)
f, err := NewForwarder(ctx, &fconf)
tt.CheckErr(err)
// decreases sleep time to speed up test
f.sleep = time.Millisecond * 500
Expand Down
Loading

0 comments on commit 28e1e43

Please sign in to comment.