rpki: improve statistics API

Signed-off-by: FUJITA Tomonori <fujita.tomonori@lab.ntt.co.jp>

2 files changed, 132 insertions, 35 deletions

diff --git a/server/rpki.go b/server/rpki.go
index 492587ae..cf05b011 100644
--- a/server/rpki.go
+++ b/server/rpki.go
@@ -126,8 +126,10 @@ func newROAManager(as uint32, servers []config.RpkiServer) (*roaManager, error)
 	for _, entry := range servers {
 		c := entry.Config
 		client := &roaClient{
-			host:    net.JoinHostPort(c.Address, strconv.Itoa(int(c.Port))),
-			eventCh: m.eventCh,
+			host:     net.JoinHostPort(c.Address, strconv.Itoa(int(c.Port))),
+			eventCh:  m.eventCh,
+			records:  make(map[int]uint32),
+			prefixes: make(map[int]uint32),
 		}
 		m.clientMap[client.host] = client
 		client.t.Go(client.tryConnect)
@@ -188,30 +190,54 @@ func (m *roaManager) handleROAEvent(ev *roaClientEvent) {
 	}
 }
 
-func deleteROA(host string, tree *radix.Tree, as uint32, prefix []byte, prefixLen, maxLen uint8) {
+func deleteROA(client *roaClient, family int, tree *radix.Tree, as uint32, prefix []byte, prefixLen, maxLen uint8) {
+	host := client.host
 	key := table.IpToRadixkey(prefix, prefixLen)
 	b, _ := tree.Get(key)
-	if b != nil {
-		bucket := b.(*roaBucket)
-		for _, r := range bucket.entries {
-			if r.MaxLen == maxLen && r.Src == host {
-				for idx, a := range r.AS {
-					if a == as {
-						r.AS = append(r.AS[:idx], r.AS[idx+1:]...)
-						if len(bucket.entries) == 0 {
-							tree.Delete(key)
+	isDeleted := func() bool {
+		if b != nil {
+			bucket := b.(*roaBucket)
+			for _, r := range bucket.entries {
+				if r.MaxLen == maxLen && r.Src == host {
+					for idx, a := range r.AS {
+						if a == as {
+							r.AS = append(r.AS[:idx], r.AS[idx+1:]...)
+							if len(bucket.entries) == 0 {
+								tree.Delete(key)
+							}
+							return true
 						}
-						return
 					}
 				}
 			}
 		}
+		return false
+	}()
+	if isDeleted {
+		client.records[family]--
+		isNoPrefix := func() bool {
+			if b, _ := tree.Get(key); b != nil {
+				bucket := b.(*roaBucket)
+				for _, r := range bucket.entries {
+					if r.Src == host {
+						return false
+					}
+				}
+				return true
+			} else {
+				return true
+			}
+		}()
+		if isNoPrefix {
+			client.prefixes[family]--
+		}
+	} else {
+		log.Info("can't withdraw a roa", net.IP(prefix).String(), as, prefixLen, maxLen)
 	}
-	p := net.IP(prefix)
-	log.Info("can't withdraw a roa", p.String(), as, prefixLen, maxLen)
 }
 
-func addROA(host string, tree *radix.Tree, as uint32, prefix []byte, prefixLen, maxLen uint8) {
+func addROA(client *roaClient, family int, tree *radix.Tree, as uint32, prefix []byte, prefixLen, maxLen uint8) {
+	host := client.host
 	key := table.IpToRadixkey(prefix, prefixLen)
 	b, _ := tree.Get(key)
 	if b == nil {
@@ -232,8 +258,22 @@ func addROA(host string, tree *radix.Tree, as uint32, prefix []byte, prefixLen,
 		r.bucket = b
 
 		tree.Insert(key, b)
+		client.prefixes[family]++
+		client.records[family]++
 	} else {
 		bucket := b.(*roaBucket)
+		isNewPrefix := func() bool {
+			for _, r := range bucket.entries {
+				if r.Src == host {
+					return false
+				}
+			}
+			return true
+		}()
+		if isNewPrefix {
+			client.prefixes[family]++
+		}
+
 		for _, r := range bucket.entries {
 			if r.MaxLen == maxLen && r.Src == host {
 				// we already have?
@@ -243,6 +283,7 @@ func addROA(host string, tree *radix.Tree, as uint32, prefix []byte, prefixLen,
 					}
 				}
 				r.AS = append(r.AS, as)
+				client.records[family]++
 				return
 			}
 		}
@@ -253,6 +294,7 @@ func addROA(host string, tree *radix.Tree, as uint32, prefix []byte, prefixLen,
 			Src:    host,
 		}
 		bucket.entries = append(bucket.entries, r)
+		client.records[family]++
 	}
 }
 
@@ -276,17 +318,19 @@ func (c *roaManager) handleRTRMsg(client *roaClient, state *config.RpkiServerSta
 			received.CacheResponse++
 		case *bgp.RTRIPPrefix:
 			var tree *radix.Tree
+			family := bgp.AFI_IP
 			if msg.Type == bgp.RTR_IPV4_PREFIX {
 				received.Ipv4Prefix++
 				tree = c.roas[bgp.RF_IPv4_UC]
 			} else {
+				family = bgp.AFI_IP6
 				received.Ipv6Prefix++
 				tree = c.roas[bgp.RF_IPv6_UC]
 			}
 			if (msg.Flags & 1) == 1 {
-				addROA(client.host, tree, msg.AS, msg.Prefix, msg.PrefixLen, msg.MaxLen)
+				addROA(client, family, tree, msg.AS, msg.Prefix, msg.PrefixLen, msg.MaxLen)
 			} else {
-				deleteROA(client.host, tree, msg.AS, msg.Prefix, msg.PrefixLen, msg.MaxLen)
+				deleteROA(client, family, tree, msg.AS, msg.Prefix, msg.PrefixLen, msg.MaxLen)
 			}
 		case *bgp.RTREndOfData:
 			received.EndOfData++
@@ -321,17 +365,36 @@ func (c *roaManager) handleGRPC(grpcReq *GrpcRequest) {
 		results := make([]*GrpcResponse, 0)
 		for _, client := range c.clientMap {
 			state := client.state
-			received := &state.RpkiMessages.RpkiReceived
 			addr, port := splitHostPort(client.host)
+			received := &state.RpkiMessages.RpkiReceived
+			sent := client.state.RpkiMessages.RpkiSent
+			up := true
+			if client.conn == nil {
+				up = false
+			}
 			rpki := &api.RPKI{
 				Conf: &api.RPKIConf{
 					Address:    addr,
 					RemotePort: uint32(port),
 				},
 				State: &api.RPKIState{
-					Uptime:       state.Uptime,
-					ReceivedIpv4: received.Ipv4Prefix,
-					ReceivedIpv6: received.Ipv6Prefix,
+					Uptime:        state.Uptime,
+					Downtime:      state.Downtime,
+					Up:            up,
+					RecordIpv4:    client.records[bgp.AFI_IP],
+					RecordIpv6:    client.records[bgp.AFI_IP6],
+					PrefixIpv4:    client.prefixes[bgp.AFI_IP],
+					PrefixIpv6:    client.prefixes[bgp.AFI_IP6],
+					Serial:        client.serialNumber,
+					ReceivedIpv4:  received.Ipv4Prefix,
+					ReceivedIpv6:  received.Ipv6Prefix,
+					SerialNotify:  received.SerialNotify,
+					CacheReset:    received.CacheReset,
+					CacheResponse: received.CacheResponse,
+					EndOfData:     received.EndOfData,
+					Error:         received.Error,
+					SerialQuery:   sent.SerialQuery,
+					ResetQuery:    sent.ResetQuery,
 				},
 			}
 			result := &GrpcResponse{}
@@ -491,6 +554,8 @@ type roaClient struct {
 	eventCh      chan *roaClientEvent
 	sessionID    uint16
 	serialNumber uint32
+	prefixes     map[int]uint32
+	records      map[int]uint32
 }
 
 func (c *roaClient) enable(serial uint32) error {
diff --git a/server/rpki_test.go b/server/rpki_test.go
index 8be709ea..3d7be7b3 100644
--- a/server/rpki_test.go
+++ b/server/rpki_test.go
@@ -60,9 +60,13 @@ func validateOne(tree *radix.Tree, cidr, aspathStr string) config.RpkiValidation
 func TestValidate0(t *testing.T) {
 	assert := assert.New(t)
 
+	client := &roaClient{
+		records:  make(map[int]uint32),
+		prefixes: make(map[int]uint32),
+	}
 	tree := radix.New()
-	addROA("", tree, 100, net.ParseIP("192.168.0.0").To4(), 24, 32)
-	addROA("", tree, 200, net.ParseIP("192.168.0.0").To4(), 24, 24)
+	addROA(client, bgp.AFI_IP, tree, 100, net.ParseIP("192.168.0.0").To4(), 24, 32)
+	addROA(client, bgp.AFI_IP, tree, 200, net.ParseIP("192.168.0.0").To4(), 24, 24)
 
 	var r config.RpkiValidationResultType
 
@@ -88,8 +92,12 @@ func TestValidate0(t *testing.T) {
 func TestValidate1(t *testing.T) {
 	assert := assert.New(t)
 
+	client := &roaClient{
+		records:  make(map[int]uint32),
+		prefixes: make(map[int]uint32),
+	}
 	tree := radix.New()
-	addROA("", tree, 65000, net.ParseIP("10.0.0.0").To4(), 16, 16)
+	addROA(client, bgp.AFI_IP, tree, 65000, net.ParseIP("10.0.0.0").To4(), 16, 16)
 
 	var r config.RpkiValidationResultType
 
@@ -117,8 +125,12 @@ func TestValidate2(t *testing.T) {
 func TestValidate3(t *testing.T) {
 	assert := assert.New(t)
 
+	client := &roaClient{
+		records:  make(map[int]uint32),
+		prefixes: make(map[int]uint32),
+	}
 	tree1 := radix.New()
-	addROA("", tree1, 65000, net.ParseIP("10.0.0.0").To4(), 16, 16)
+	addROA(client, bgp.AFI_IP, tree1, 65000, net.ParseIP("10.0.0.0").To4(), 16, 16)
 
 	var r config.RpkiValidationResultType
 
@@ -129,7 +141,7 @@ func TestValidate3(t *testing.T) {
 	assert.Equal(r, config.RPKI_VALIDATION_RESULT_TYPE_INVALID)
 
 	tree2 := radix.New()
-	addROA("", tree2, 65000, net.ParseIP("10.0.0.0").To4(), 16, 24)
+	addROA(client, bgp.AFI_IP, tree2, 65000, net.ParseIP("10.0.0.0").To4(), 16, 24)
 
 	r = validateOne(tree2, "10.0.0.0/17", "65000")
 	assert.Equal(r, config.RPKI_VALIDATION_RESULT_TYPE_VALID)
@@ -138,9 +150,13 @@ func TestValidate3(t *testing.T) {
 func TestValidate4(t *testing.T) {
 	assert := assert.New(t)
 
+	client := &roaClient{
+		records:  make(map[int]uint32),
+		prefixes: make(map[int]uint32),
+	}
 	tree := radix.New()
-	addROA("", tree, 65000, net.ParseIP("10.0.0.0").To4(), 16, 16)
-	addROA("", tree, 65001, net.ParseIP("10.0.0.0").To4(), 16, 16)
+	addROA(client, bgp.AFI_IP, tree, 65000, net.ParseIP("10.0.0.0").To4(), 16, 16)
+	addROA(client, bgp.AFI_IP, tree, 65001, net.ParseIP("10.0.0.0").To4(), 16, 16)
 
 	var r config.RpkiValidationResultType
 
@@ -154,9 +170,13 @@ func TestValidate4(t *testing.T) {
 func TestValidate5(t *testing.T) {
 	assert := assert.New(t)
 
+	client := &roaClient{
+		records:  make(map[int]uint32),
+		prefixes: make(map[int]uint32),
+	}
 	tree := radix.New()
-	addROA("", tree, 65000, net.ParseIP("10.0.0.0").To4(), 17, 17)
-	addROA("", tree, 65000, net.ParseIP("10.0.128.0").To4(), 17, 17)
+	addROA(client, bgp.AFI_IP, tree, 65000, net.ParseIP("10.0.0.0").To4(), 17, 17)
+	addROA(client, bgp.AFI_IP, tree, 65000, net.ParseIP("10.0.128.0").To4(), 17, 17)
 
 	var r config.RpkiValidationResultType
 
@@ -167,8 +187,12 @@ func TestValidate5(t *testing.T) {
 func TestValidate6(t *testing.T) {
 	assert := assert.New(t)
 
+	client := &roaClient{
+		records:  make(map[int]uint32),
+		prefixes: make(map[int]uint32),
+	}
 	tree := radix.New()
-	addROA("", tree, 0, net.ParseIP("10.0.0.0").To4(), 8, 32)
+	addROA(client, bgp.AFI_IP, tree, 0, net.ParseIP("10.0.0.0").To4(), 8, 32)
 
 	var r config.RpkiValidationResultType
 
@@ -185,8 +209,12 @@ func TestValidate6(t *testing.T) {
 func TestValidate7(t *testing.T) {
 	assert := assert.New(t)
 
+	client := &roaClient{
+		records:  make(map[int]uint32),
+		prefixes: make(map[int]uint32),
+	}
 	tree := radix.New()
-	addROA("", tree, 65000, net.ParseIP("10.0.0.0").To4(), 16, 24)
+	addROA(client, bgp.AFI_IP, tree, 65000, net.ParseIP("10.0.0.0").To4(), 16, 24)
 
 	var r config.RpkiValidationResultType
 
@@ -203,9 +231,13 @@ func TestValidate7(t *testing.T) {
 func TestValidate8(t *testing.T) {
 	assert := assert.New(t)
 
+	client := &roaClient{
+		records:  make(map[int]uint32),
+		prefixes: make(map[int]uint32),
+	}
 	tree := radix.New()
-	addROA("", tree, 0, net.ParseIP("10.0.0.0").To4(), 16, 24)
-	addROA("", tree, 65000, net.ParseIP("10.0.0.0").To4(), 16, 24)
+	addROA(client, bgp.AFI_IP, tree, 0, net.ParseIP("10.0.0.0").To4(), 16, 24)
+	addROA(client, bgp.AFI_IP, tree, 65000, net.ParseIP("10.0.0.0").To4(), 16, 24)
 
 	var r config.RpkiValidationResultType