Merge pull request #391 from step-security/fix-wildcard-resolution

Fix bug with wildcard resolution

Author: varunsh-coder

agent.go

@@ -270,7 +270,7 @@ func refreshDNSEntries(ctx context.Context, iptables *Firewall, allowedEndpoints
 							}
 
 							// add to cache with new TTL
-							dnsProxy.Cache.Set(domainName, answer)
+							dnsProxy.Cache.Set(domainName, answer, false)
 
 							WriteLog(fmt.Sprintf("domain resolved: %s, ip address: %s, TTL: %d", domainName, answer.Data, answer.TTL))
 						}

cache.go

@@ -6,8 +6,9 @@ import (
 )
 
 type Element struct {
-	Value     *Answer
-	TimeAdded int64
+	Value            *Answer
+	TimeAdded        int64
+	IsWildcardDomain bool
 }
 
 type Cache struct {
@@ -32,7 +33,7 @@ func (cache *Cache) Get(k string) (*Element, bool) {
 		return nil, false
 	}
 
-	if cache.egressPolicy == EgressPolicyAudit {
+	if cache.egressPolicy == EgressPolicyAudit || element.IsWildcardDomain {
 		// TTL is in seconds
 		// if now minus time added is greater than TTL, return nil, so new DNS request is made
 		if time.Now().Unix()-element.TimeAdded > int64(element.Value.TTL) {
@@ -51,12 +52,13 @@ func (cache *Cache) Get(k string) (*Element, bool) {
 	}
 }
 
-func (cache *Cache) Set(k string, v *Answer) {
+func (cache *Cache) Set(k string, v *Answer, isWildcardDomain bool) {
 	cache.mutex.Lock()
 
 	cache.elements[k] = Element{
-		Value:     v,
-		TimeAdded: time.Now().Unix(),
+		Value:            v,
+		TimeAdded:        time.Now().Unix(),
+		IsWildcardDomain: isWildcardDomain,
 	}
 
 	cache.mutex.Unlock()

dnsproxy.go

@@ -206,7 +206,7 @@ func (proxy *DNSProxy) getIPByDomain(domain string) (string, error) {
 
 				// return an ip address, so calling process calls the ip address
 				// the call will be blocked by the firewall
-				proxy.Cache.Set(domain, &Answer{Name: domain, TTL: math.MaxInt32, Data: StepSecuritySinkHoleIPAddress})
+				proxy.Cache.Set(domain, &Answer{Name: domain, TTL: math.MaxInt32, Data: StepSecuritySinkHoleIPAddress}, false)
 
 				go proxy.ApiClient.sendDNSRecord(proxy.CorrelationId, proxy.Repo, domain, StepSecuritySinkHoleIPAddress)
 
@@ -227,7 +227,7 @@ func (proxy *DNSProxy) getIPByDomain(domain string) (string, error) {
 		}
 	}
 
-	proxy.Cache.Set(domain, answer)
+	proxy.Cache.Set(domain, answer, matchesAnyWildcard)
 
 	go WriteLog(fmt.Sprintf("domain resolved: %s, ip address: %s, TTL: %d", domain, answer.Data, answer.TTL))
 
@@ -259,7 +259,7 @@ func (proxy *DNSProxy) processTypeA(q *dns.Question, requestMsg *dns.Msg) (*dns.
 			return nil, err
 		}
 
-		proxy.Cache.Set(q.Name, &Answer{Name: q.Name, TTL: math.MaxInt32, Data: "8.8.8.8"})
+		proxy.Cache.Set(q.Name, &Answer{Name: q.Name, TTL: math.MaxInt32, Data: "8.8.8.8"}, false)
 
 		return &rr, nil
 	}