Merge pull request #117 from step-security/feature-28

varunsh-coder · web-flow · commit c117f150332c · 2021-12-23T14:09:18.000-08:00
Pass ctx to netmon
diff --git a/agent.go b/agent.go
@@ -143,7 +143,7 @@ func Run(ctx context.Context, configFilePath string, hostDNSServer DNSServer,
 		}
 
 		// Start network monitor
-		go netMonitor.MonitorNetwork(nflog, errc) // listens for NFLOG messages
+		go netMonitor.MonitorNetwork(ctx, nflog, errc) // listens for NFLOG messages
 
 		WriteLog("before audit rules")
 
@@ -167,7 +167,7 @@ func Run(ctx context.Context, configFilePath string, hostDNSServer DNSServer,
 		}
 
 		// Start network monitor
-		go netMonitor.MonitorNetwork(nflog, errc) // listens for NFLOG messages
+		go netMonitor.MonitorNetwork(ctx, nflog, errc) // listens for NFLOG messages
 
 		if err := addBlockRulesForGitHubHostedRunner(iptables, ipAddressEndpoints); err != nil {
 			WriteLog(fmt.Sprintf("Error setting firewall for allowed domains %v", err))
diff --git a/agent_test.go b/agent_test.go
@@ -3,6 +3,7 @@ package main
 import (
 	"context"
 	"fmt"
+	"net/http"
 	"os"
 	"path"
 	"testing"
@@ -124,7 +125,15 @@ func TestRun(t *testing.T) {
 		httpmock.NewStringResponder(200, `{"Status":0,"TC":false,"RD":true,"RA":true,"AD":false,"CD":false,"Question":[{"name":"domain1.com.","type":1}],"Answer":[{"name":"domain1.com.","type":1,"TTL":30,"data":"67.67.67.67"}]}`))
 
 	httpmock.RegisterResponder("GET", "https://dns.google/resolve?name=domain2.com.&type=a",
-		httpmock.NewStringResponder(200, `{"Status":0,"TC":false,"RD":true,"RA":true,"AD":false,"CD":false,"Question":[{"name":"domain2.com.","type":1}],"Answer":[{"name":"domain2.com.","type":1,"TTL":30,"data":"68.68.68.68"}]}`))
+		httpmock.ResponderFromMultipleResponses(
+			[]*http.Response{
+				httpmock.NewStringResponse(200, `{"Status":0,"TC":false,"RD":true,"RA":true,"AD":false,"CD":false,"Question":[{"name":"domain2.com.","type":1}],"Answer":[{"name":"domain2.com.","type":1,"TTL":30,"data":"68.68.68.68"}]}`),
+				httpmock.NewStringResponse(200, `{"Status":0,"TC":false,"RD":true,"RA":true,"AD":false,"CD":false,"Question":[{"name":"domain2.com.","type":1}],"Answer":[{"name":"domain2.com.","type":1,"TTL":30,"data":"68.68.68.68"}]}`),
+				httpmock.NewStringResponse(200, `{"Status":0,"TC":false,"RD":true,"RA":true,"AD":false,"CD":false,"Question":[{"name":"domain2.com.","type":1}],"Answer":[{"name":"domain2.com.","type":1,"TTL":30,"data":"70.70.70.70"}]}`),
+				httpmock.NewStringResponse(200, `{"Status":0,"TC":false,"RD":true,"RA":true,"AD":false,"CD":false,"Question":[{"name":"domain2.com.","type":1}],"Answer":[{"name":"domain2.com.","type":1,"TTL":30,"data":"68.68.68.68"}]}`),
+				httpmock.NewStringResponse(200, `{"Status":0,"TC":false,"RD":true,"RA":true,"AD":false,"CD":false,"Question":[{"name":"domain2.com.","type":1}],"Answer":[{"name":"domain2.com.","type":1,"TTL":30,"data":"70.70.70.70"}]}`),
+			},
+			t.Log))
 
 	httpmock.RegisterResponder("GET", "https://dns.google/resolve", // no query params to match all other requests
 		httpmock.NewStringResponder(200, `{"Status":0,"TC":false,"RD":true,"RA":true,"AD":false,"CD":false,"Question":[{"name":"requesteddomain.com.","type":1}],"Answer":[{"name":"requesteddomain.com.","type":1,"TTL":300,"data":"69.69.69.69"}]}`))
@@ -134,28 +143,21 @@ func TestRun(t *testing.T) {
 		args    args
 		wantErr bool
 	}{
-		{name: "success", args: args{ctxCancelDuration: 2, configFilePath: "./testfiles/agent.json", hostDNSServer: &mockDNSServer{}, dockerDNSServer: &mockDNSServer{},
+		{name: "success egress audit", args: args{ctxCancelDuration: 2, configFilePath: "./testfiles/agent.json", hostDNSServer: &mockDNSServer{}, dockerDNSServer: &mockDNSServer{},
 			iptables: &Firewall{&MockIPTables{}}, nflog: &MockAgentNflogger{}, cmd: &MockCommand{}, resolvdConfigPath: createTempFileWithContents(""),
 			dockerDaemonConfigPath: createTempFileWithContents("{}")}, wantErr: false},
 
-		{name: "success monitor process", args: args{ctxCancelDuration: 2, configFilePath: "./testfiles/agent.json", hostDNSServer: &mockDNSServer{}, dockerDNSServer: &mockDNSServer{},
-			iptables: &Firewall{&MockIPTables{}}, nflog: &MockAgentNflogger{}, cmd: nil, resolvdConfigPath: createTempFileWithContents(""),
-			dockerDaemonConfigPath: createTempFileWithContents("{}"), ciTestOnly: true}, wantErr: false},
-
-		{name: "success allowed endpoints", args: args{ctxCancelDuration: 2, configFilePath: "./testfiles/agent-allowed-endpoints.json",
+		{name: "success egress blocked", args: args{ctxCancelDuration: 2, configFilePath: "./testfiles/agent-allowed-endpoints.json",
 			hostDNSServer: &mockDNSServer{}, dockerDNSServer: &mockDNSServer{},
 			iptables: &Firewall{&MockIPTables{}}, nflog: &MockAgentNflogger{}, cmd: &MockCommand{}, resolvdConfigPath: createTempFileWithContents(""),
 			dockerDaemonConfigPath: createTempFileWithContents("{}")}, wantErr: false},
 
-		{name: "success allowed endpoints CI Test", args: args{ctxCancelDuration: 2, configFilePath: "./testfiles/agent-allowed-endpoints.json",
+		// ctx will cancel after 35 seconds
+		// DNS refresh will be done after 30 seconds
+		{name: "success egress blocked DNS refresh", args: args{ctxCancelDuration: 35, configFilePath: "./testfiles/agent-allowed-endpoints.json",
 			hostDNSServer: &mockDNSServer{}, dockerDNSServer: &mockDNSServer{},
-			iptables: nil, nflog: &MockAgentNflogger{}, cmd: &MockCommand{}, resolvdConfigPath: createTempFileWithContents(""),
-			dockerDaemonConfigPath: createTempFileWithContents("{}"), ciTestOnly: true}, wantErr: false},
-
-		{name: "success allowed endpoints DNS refresh CI Test", args: args{ctxCancelDuration: 60, configFilePath: "./testfiles/agent-allowed-endpoints.json",
-			hostDNSServer: &mockDNSServer{}, dockerDNSServer: &mockDNSServer{},
-			iptables: nil, nflog: &MockAgentNflogger{}, cmd: &MockCommand{}, resolvdConfigPath: createTempFileWithContents(""),
-			dockerDaemonConfigPath: createTempFileWithContents("{}"), ciTestOnly: true}, wantErr: false},
+			iptables: &Firewall{&MockIPTables{}}, nflog: &MockAgentNflogger{}, cmd: &MockCommand{}, resolvdConfigPath: createTempFileWithContents(""),
+			dockerDaemonConfigPath: createTempFileWithContents("{}")}, wantErr: false},
 
 		{name: "dns failure", args: args{ctxCancelDuration: 5, configFilePath: "./testfiles/agent.json", hostDNSServer: &mockDNSServer{}, dockerDNSServer: &mockDNSServerWithError{},
 			iptables: &Firewall{&MockIPTables{}}, nflog: &MockAgentNflogger{}, cmd: &MockCommand{}, resolvdConfigPath: createTempFileWithContents(""),
@@ -168,6 +170,21 @@ func TestRun(t *testing.T) {
 		{name: "nflog failure", args: args{ctxCancelDuration: 5, configFilePath: "./testfiles/agent.json", hostDNSServer: &mockDNSServer{}, dockerDNSServer: &mockDNSServer{},
 			iptables: &Firewall{&MockIPTables{}}, nflog: &MockAgentNfloggerWithErr{}, cmd: &MockCommand{}, resolvdConfigPath: createTempFileWithContents(""),
 			dockerDaemonConfigPath: createTempFileWithContents("{}")}, wantErr: true},
+
+		// CI only tests
+		{name: "success monitor process CI Test", args: args{ctxCancelDuration: 2, configFilePath: "./testfiles/agent.json", hostDNSServer: &mockDNSServer{}, dockerDNSServer: &mockDNSServer{},
+			iptables: &Firewall{&MockIPTables{}}, nflog: &MockAgentNflogger{}, cmd: nil, resolvdConfigPath: createTempFileWithContents(""),
+			dockerDaemonConfigPath: createTempFileWithContents("{}"), ciTestOnly: true}, wantErr: false},
+
+		{name: "success allowed endpoints CI Test", args: args{ctxCancelDuration: 2, configFilePath: "./testfiles/agent-allowed-endpoints.json",
+			hostDNSServer: &mockDNSServer{}, dockerDNSServer: &mockDNSServer{},
+			iptables: nil, nflog: &MockAgentNflogger{}, cmd: &MockCommand{}, resolvdConfigPath: createTempFileWithContents(""),
+			dockerDaemonConfigPath: createTempFileWithContents("{}"), ciTestOnly: true}, wantErr: false},
+
+		{name: "success allowed endpoints DNS refresh CI Test", args: args{ctxCancelDuration: 60, configFilePath: "./testfiles/agent-allowed-endpoints.json",
+			hostDNSServer: &mockDNSServer{}, dockerDNSServer: &mockDNSServer{},
+			iptables: nil, nflog: &MockAgentNflogger{}, cmd: &MockCommand{}, resolvdConfigPath: createTempFileWithContents(""),
+			dockerDaemonConfigPath: createTempFileWithContents("{}"), ciTestOnly: true}, wantErr: false},
 	}
 	_, ciTest := os.LookupEnv("CI")
 	fmt.Printf("ci-test: %t\n", ciTest)
diff --git a/dnsproxy.go b/dnsproxy.go
@@ -113,6 +113,9 @@ func (proxy *DNSProxy) ResolveDomain(domain string) (*Answer, error) {
 
 	for _, answer := range dnsReponse.Answer {
 		if answer.Type == 1 {
+			if answer.TTL < 30 {
+				answer.TTL = 30 // less than 30 will cause too frequent DNS requests in audit scenario
+			}
 			return &answer, nil
 		}
 	}
diff --git a/dnsproxy_test.go b/dnsproxy_test.go
@@ -5,6 +5,7 @@ import (
 	"net/http"
 	"reflect"
 	"testing"
+	"time"
 
 	"github.com/jarcoal/httpmock"
 	"github.com/miekg/dns"
@@ -113,3 +114,44 @@ func TestDNSProxy_getResponse(t *testing.T) {
 		})
 	}
 }
+
+func TestDNSProxy_auditCacheTTL(t *testing.T) {
+	apiclient := &ApiClient{Client: &http.Client{}, APIURL: agentApiBaseUrl}
+
+	httpmock.ActivateNonDefault(apiclient.Client)
+
+	httpmock.RegisterResponder("GET", "https://dns.google/resolve?name=domain12345.com.&type=a",
+		httpmock.ResponderFromMultipleResponses(
+			[]*http.Response{
+				httpmock.NewStringResponse(200, `{"Status":0,"TC":false,"RD":true,"RA":true,"AD":false,"CD":false,"Question":[{"name":"domain12345.com.","type":1}],"Answer":[{"name":"domain12345.com.","type":1,"TTL":30,"data":"68.68.68.68"}]}`),
+				httpmock.NewStringResponse(200, `{"Status":0,"TC":false,"RD":true,"RA":true,"AD":false,"CD":false,"Question":[{"name":"domain12345.com.","type":1}],"Answer":[{"name":"domain12345.com.","type":1,"TTL":30,"data":"68.68.68.68"}]}`),
+			},
+			t.Log))
+
+	cache := InitCache(EgressPolicyAudit)
+
+	proxy := &DNSProxy{
+		Cache:        &cache,
+		ApiClient:    apiclient,
+		EgressPolicy: EgressPolicyAudit,
+	}
+
+	// should call httpmock
+	proxy.getResponse(&dns.Msg{Question: []dns.Question{{Name: "domain12345.com.", Qtype: dns.TypeA}}})
+
+	time.Sleep(2 * time.Second)
+
+	// should get from cache
+	proxy.getResponse(&dns.Msg{Question: []dns.Question{{Name: "domain12345.com.", Qtype: dns.TypeA}}})
+
+	time.Sleep(30 * time.Second)
+
+	// should call httpmock
+	proxy.getResponse(&dns.Msg{Question: []dns.Question{{Name: "domain12345.com.", Qtype: dns.TypeA}}})
+
+	info := httpmock.GetCallCountInfo()
+	count := info["GET https://dns.google/resolve?name=domain12345.com.&type=a"]
+	if count != 2 {
+		t.Errorf("incorrect call count %d, expected 2, %v", count, info)
+	}
+}
diff --git a/netmon.go b/netmon.go
@@ -24,7 +24,7 @@ type NetworkMonitor struct {
 
 var ipAddresses = make(map[string]int)
 
-func (netMonitor *NetworkMonitor) MonitorNetwork(nflogger AgentNflogger, errc chan error) []string {
+func (netMonitor *NetworkMonitor) MonitorNetwork(ctx context.Context, nflogger AgentNflogger, errc chan error) []string {
 
 	//sysLogger, err := syslog.NewLogger(syslog.LOG_INFO|syslog.LOG_USER, 1)
 	var err error
@@ -48,9 +48,6 @@ func (netMonitor *NetworkMonitor) MonitorNetwork(nflogger AgentNflogger, errc ch
 	}
 	defer nf.Close()
 
-	ctx, cancel := context.WithCancel(context.Background()) // TODO: Pass context from the top
-	defer cancel()
-
 	fn := func(attrs nflog.Attribute) int {
 		go netMonitor.handlePacket(attrs)
 		return 0

Original file line number	Diff line number	Diff line change
`@@ -143,7 +143,7 @@ func Run(ctx context.Context, configFilePath string, hostDNSServer DNSServer,`
`143`	`143`	`}`
`144`	`144`
`145`	`145`	`// Start network monitor`
`146`		`- go netMonitor.MonitorNetwork(nflog, errc) // listens for NFLOG messages`
	`146`	`+ go netMonitor.MonitorNetwork(ctx, nflog, errc) // listens for NFLOG messages`
`147`	`147`
`148`	`148`	`WriteLog("before audit rules")`
`149`	`149`
`@@ -167,7 +167,7 @@ func Run(ctx context.Context, configFilePath string, hostDNSServer DNSServer,`
`167`	`167`	`}`
`168`	`168`
`169`	`169`	`// Start network monitor`
`170`		`- go netMonitor.MonitorNetwork(nflog, errc) // listens for NFLOG messages`
	`170`	`+ go netMonitor.MonitorNetwork(ctx, nflog, errc) // listens for NFLOG messages`
`171`	`171`
`172`	`172`	`if err := addBlockRulesForGitHubHostedRunner(iptables, ipAddressEndpoints); err != nil {`
`173`	`173`	`WriteLog(fmt.Sprintf("Error setting firewall for allowed domains %v", err))`
Original file line number	Diff line number	Diff line change
`@@ -113,6 +113,9 @@ func (proxy DNSProxy) ResolveDomain(domain string) (Answer, error) {`
`113`	`113`
`114`	`114`	`for _, answer := range dnsReponse.Answer {`
`115`	`115`	`if answer.Type == 1 {`
	`116`	`+ if answer.TTL < 30 {`
	`117`	`+ answer.TTL = 30 // less than 30 will cause too frequent DNS requests in audit scenario`
	`118`	`+ }`
`116`	`119`	`return &answer, nil`
`117`	`120`	`}`
`118`	`121`	`}`