From c32647ec928b2f87dacb76b6754283f613dcd24e Mon Sep 17 00:00:00 2001
From: Flavio Crisciani <flavio.crisciani@docker.com>
Date: Tue, 3 Oct 2017 17:12:57 -0700
Subject: [PATCH] Fix IPMask marshalling

Fix marshalling and add test

Signed-off-by: Flavio Crisciani <flavio.crisciani@docker.com>
---
 drivers/overlay/peerdb.go      | 21 ++++++++++++---------
 drivers/overlay/peerdb_test.go | 30 ++++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+), 9 deletions(-)
 create mode 100644 drivers/overlay/peerdb_test.go

diff --git a/drivers/overlay/peerdb.go b/drivers/overlay/peerdb.go
index a27687bcbe..bdd3cb12af 100644
--- a/drivers/overlay/peerdb.go
+++ b/drivers/overlay/peerdb.go
@@ -27,11 +27,13 @@ type peerEntry struct {
 }
 
 func (p *peerEntry) MarshalDB() peerEntryDB {
+	ones, bits := p.peerIPMask.Size()
 	return peerEntryDB{
-		eid:        p.eid,
-		vtep:       p.vtep.String(),
-		peerIPMask: p.peerIPMask.String(),
-		isLocal:    p.isLocal,
+		eid:            p.eid,
+		vtep:           p.vtep.String(),
+		peerIPMaskOnes: ones,
+		peerIPMaskBits: bits,
+		isLocal:        p.isLocal,
 	}
 }
 
@@ -39,17 +41,18 @@ func (p *peerEntry) MarshalDB() peerEntryDB {
 // the value inserted in the set has to be Hashable so the []byte had to be converted into
 // strings
 type peerEntryDB struct {
-	eid        string
-	vtep       string
-	peerIPMask string
-	isLocal    bool
+	eid            string
+	vtep           string
+	peerIPMaskOnes int
+	peerIPMaskBits int
+	isLocal        bool
 }
 
 func (p *peerEntryDB) UnMarshalDB() peerEntry {
 	return peerEntry{
 		eid:        p.eid,
 		vtep:       net.ParseIP(p.vtep),
-		peerIPMask: net.IPMask(net.ParseIP(p.peerIPMask)),
+		peerIPMask: net.CIDRMask(p.peerIPMaskOnes, p.peerIPMaskBits),
 		isLocal:    p.isLocal,
 	}
 }
diff --git a/drivers/overlay/peerdb_test.go b/drivers/overlay/peerdb_test.go
new file mode 100644
index 0000000000..13842b0462
--- /dev/null
+++ b/drivers/overlay/peerdb_test.go
@@ -0,0 +1,30 @@
+package overlay
+
+import (
+	"net"
+	"testing"
+
+	_ "github.com/docker/libnetwork/testutils"
+)
+
+func TestPeerMarshal(t *testing.T) {
+	_, ipNet, _ := net.ParseCIDR("192.168.0.1/24")
+	p := &peerEntry{eid: "eid",
+		isLocal:    true,
+		peerIPMask: ipNet.Mask,
+		vtep:       ipNet.IP}
+	entryDB := p.MarshalDB()
+	x := entryDB.UnMarshalDB()
+	if x.eid != p.eid {
+		t.Fatalf("Incorrect Unmarshalling for eid: %v != %v", x.eid, p.eid)
+	}
+	if x.isLocal != p.isLocal {
+		t.Fatalf("Incorrect Unmarshalling for isLocal: %v != %v", x.isLocal, p.isLocal)
+	}
+	if x.peerIPMask.String() != p.peerIPMask.String() {
+		t.Fatalf("Incorrect Unmarshalling for eid: %v != %v", x.peerIPMask, p.peerIPMask)
+	}
+	if x.vtep.String() != p.vtep.String() {
+		t.Fatalf("Incorrect Unmarshalling for eid: %v != %v", x.vtep, p.vtep)
+	}
+}