Fix ca* checks

firewall.go

@@ -83,19 +83,23 @@ func newFirewallTable() *FirewallTable {
 	}
 }
 
+type FirewallCA struct {
+	Any     *FirewallRule
+	CANames map[string]*FirewallRule
+	CAShas  map[string]*FirewallRule
+}
+
 type FirewallRule struct {
-	// Any makes Hosts, Groups, and CIDR irrelevant. CAName and CASha still need to be checked
-	Any     bool
-	Hosts   map[string]struct{}
-	Groups  [][]string
-	CIDR    *CIDRTree
-	CANames map[string]struct{}
-	CAShas  map[string]struct{}
+	// Any makes Hosts, Groups, and CIDR irrelevant
+	Any    bool
+	Hosts  map[string]struct{}
+	Groups [][]string
+	CIDR   *CIDRTree
 }
 
 // Even though ports are uint16, int32 maps are faster for lookup
 // Plus we can use `-1` for fragment rules
-type firewallPort map[int32]*FirewallRule
+type firewallPort map[int32]*FirewallCA
 
 type FirewallPacket struct {
 	LocalIP    uint32
@@ -182,9 +186,9 @@ func NewFirewall(tcpTimeout, UDPTimeout, defaultTimeout time.Duration, c *cert.N
 
 func NewFirewallFromConfig(nc *cert.NebulaCertificate, c *Config) (*Firewall, error) {
 	fw := NewFirewall(
-		c.GetDuration("firewall.conntrack.tcp_timeout", time.Duration(time.Minute*12)),
-		c.GetDuration("firewall.conntrack.udp_timeout", time.Duration(time.Minute*3)),
-		c.GetDuration("firewall.conntrack.default_timeout", time.Duration(time.Minute*10)),
+		c.GetDuration("firewall.conntrack.tcp_timeout", time.Minute*12),
+		c.GetDuration("firewall.conntrack.udp_timeout", time.Minute*3),
+		c.GetDuration("firewall.conntrack.default_timeout", time.Minute*10),
 		nc,
 		//TODO: max_connections
 	)
@@ -499,12 +503,9 @@ func (fp firewallPort) addRule(startPort int32, endPort int32, groups []string,
 
 	for i := startPort; i <= endPort; i++ {
 		if _, ok := fp[i]; !ok {
-			fp[i] = &FirewallRule{
-				Groups:  make([][]string, 0),
-				Hosts:   make(map[string]struct{}),
-				CIDR:    NewCIDRTree(),
-				CANames: make(map[string]struct{}),
-				CAShas:  make(map[string]struct{}),
+			fp[i] = &FirewallCA{
+				CANames: make(map[string]*FirewallRule),
+				CAShas:  make(map[string]*FirewallRule),
 			}
 		}
 
@@ -539,15 +540,83 @@ func (fp firewallPort) match(p FirewallPacket, incoming bool, c *cert.NebulaCert
 	return fp[fwPortAny].match(p, c, caPool)
 }
 
-func (fr *FirewallRule) addRule(groups []string, host string, ip *net.IPNet, caName string, caSha string) error {
-	if caName != "" {
-		fr.CANames[caName] = struct{}{}
+func (fc *FirewallCA) addRule(groups []string, host string, ip *net.IPNet, caName, caSha string) error {
+	// If there is an any rule then there is no need to establish specific ca rules
+	if fc.Any != nil {
+		return fc.Any.addRule(groups, host, ip)
+	}
+
+	fr := func() *FirewallRule {
+		return &FirewallRule{
+			Hosts:  make(map[string]struct{}),
+			Groups: make([][]string, 0),
+			CIDR:   NewCIDRTree(),
+		}
+	}
+
+	any := false
+	if caSha == "" && caName == "" {
+		any = true
+	}
+
+	if any {
+		if fc.Any == nil {
+			fc.Any = fr()
+		}
+
+		// If it's any we need to wipe out any pre-existing rules to save on memory
+		fc.CAShas = make(map[string]*FirewallRule)
+		fc.CANames = make(map[string]*FirewallRule)
+		return fc.Any.addRule(groups, host, ip)
 	}
 
 	if caSha != "" {
-		fr.CAShas[caSha] = struct{}{}
+		if _, ok := fc.CAShas[caSha]; !ok {
+			fc.CAShas[caSha] = fr()
+		}
+		err := fc.CAShas[caSha].addRule(groups, host, ip)
+		if err != nil {
+			return err
+		}
 	}
 
+	if caName != "" {
+		if _, ok := fc.CANames[caName]; !ok {
+			fc.CANames[caName] = fr()
+		}
+		err := fc.CANames[caName].addRule(groups, host, ip)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func (fc *FirewallCA) match(p FirewallPacket, c *cert.NebulaCertificate, caPool *cert.NebulaCAPool) bool {
+	if fc == nil {
+		return false
+	}
+
+	if fc.Any != nil {
+		return fc.Any.match(p, c)
+	}
+
+	if t, ok := fc.CAShas[c.Details.Issuer]; ok {
+		if t.match(p, c) {
+			return true
+		}
+	}
+
+	s, err := caPool.GetCAForCert(c)
+	if err != nil {
+		return false
+	}
+
+	return fc.CANames[s.Details.Name].match(p, c)
+}
+
+func (fr *FirewallRule) addRule(groups []string, host string, ip *net.IPNet) error {
 	if fr.Any {
 		return nil
 	}
@@ -593,28 +662,11 @@ func (fr *FirewallRule) isAny(groups []string, host string, ip *net.IPNet) bool
 	return false
 }
 
-func (fr *FirewallRule) match(p FirewallPacket, c *cert.NebulaCertificate, caPool *cert.NebulaCAPool) bool {
+func (fr *FirewallRule) match(p FirewallPacket, c *cert.NebulaCertificate) bool {
 	if fr == nil {
 		return false
 	}
 
-	// CASha and CAName always need to be checked
-	if len(fr.CAShas) > 0 {
-		if _, ok := fr.CAShas[c.Details.Issuer]; !ok {
-			return false
-		}
-	}
-
-	if len(fr.CANames) > 0 {
-		s, err := caPool.GetCAForCert(c)
-		if err != nil {
-			return false
-		}
-		if _, ok := fr.CANames[s.Details.Name]; !ok {
-			return false
-		}
-	}
-
 	// Shortcut path for if groups, hosts, or cidr contained an `any`
 	if fr.Any {
 		return true
@@ -773,7 +825,7 @@ func setTCPRTTTracking(c *conn, p []byte) {
 	ihl := int(p[0]&0x0f) << 2
 
 	// Don't track FIN packets
-	if uint8(p[ihl+13])&tcpFIN != 0 {
+	if p[ihl+13]&tcpFIN != 0 {
 		return
 	}
 
@@ -787,7 +839,7 @@ func (f *Firewall) checkTCPRTT(c *conn, p []byte) bool {
 	}
 
 	ihl := int(p[0]&0x0f) << 2
-	if uint8(p[ihl+13])&tcpACK == 0 {
+	if p[ihl+13]&tcpACK == 0 {
 		return false
 	}