server: fix radix key bug in rpki test

also clean up functions to create a radix key.

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

2 files changed, 24 insertions, 6 deletions

diff --git a/table/destination.go b/table/destination.go
index 4176a1ae..7d3fe487 100644
--- a/table/destination.go
+++ b/table/destination.go
@@ -41,14 +41,26 @@ const (
 	BPR_ROUTER_ID          = "Router ID"
 )
 
+func toRadixkey(b []byte, max uint8) string {
+	var buffer bytes.Buffer
+	for i := 0; i < len(b) && i < int(max); i++ {
+		buffer.WriteString(fmt.Sprintf("%08b", b[i]))
+	}
+	return buffer.String()[:max]
+}
+
+func IpToRadixkey(prefix net.IP, prefixLen uint8) string {
+	b := prefix.To4()
+	if b == nil {
+		b = prefix.To16()
+	}
+	return toRadixkey(b, prefixLen)
+}
+
 func CidrToRadixkey(cidr string) string {
 	_, n, _ := net.ParseCIDR(cidr)
 	ones, _ := n.Mask.Size()
-	var buffer bytes.Buffer
-	for i := 0; i < len(n.IP) && i < ones; i++ {
-		buffer.WriteString(fmt.Sprintf("%08b", n.IP[i]))
-	}
-	return buffer.String()[:ones]
+	return toRadixkey(n.IP, uint8(ones))
 }
 
 type PeerInfo struct {
diff --git a/table/destination_test.go b/table/destination_test.go
index abf1c652..e8edf562 100644
--- a/table/destination_test.go
+++ b/table/destination_test.go
@@ -19,7 +19,7 @@ import (
 	//"fmt"
 	"github.com/osrg/gobgp/packet"
 	"github.com/stretchr/testify/assert"
-	//"net"
+	"net"
 	"testing"
 	"time"
 )
@@ -191,3 +191,9 @@ func updateMsgD3() *bgp.BGPMessage {
 	UpdatePathAttrs4ByteAs(updateMsg.Body.(*bgp.BGPUpdate))
 	return updateMsg
 }
+
+func TestRadixkey(t *testing.T) {
+	assert.Equal(t, "000010100000001100100000", CidrToRadixkey("10.3.32.0/24"))
+	assert.Equal(t, "000010100000001100100000", IpToRadixkey(net.ParseIP("10.3.32.0"), 24))
+	assert.Equal(t, "000010100000001100100000", IpToRadixkey(net.ParseIP("10.3.32.0").To4(), 24))
+}