From 2a00b0092d826c1a96dc6afdc138b1589a24a267 Mon Sep 17 00:00:00 2001
From: Bryan Boreham <bjboreham@gmail.com>
Date: Wed, 3 Jun 2015 12:10:28 +0100
Subject: [PATCH] Extend 'Lookup' to multi-subnet

---
 ipam/allocator.go | 11 ++++++-----
 ipam/http.go      | 18 +++++++++++++++++-
 ipam/http_test.go |  5 +++++
 3 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/ipam/allocator.go b/ipam/allocator.go
index aeef0f73..0444ff69 100644
--- a/ipam/allocator.go
+++ b/ipam/allocator.go
@@ -200,13 +200,14 @@ func (alloc *Allocator) Allocate(ident string, subnet address.CIDR, cancelChan <
 	return result.addr, result.err
 }
 
-func (alloc *Allocator) Lookup(ident string) (address.Address, error) {
+func (alloc *Allocator) Lookup(ident string, subnet address.CIDR) (address.Address, error) {
 	resultChan := make(chan allocateResult)
 	alloc.actionChan <- func() {
-		addr, found := alloc.owned[ident]
-		if found {
-			resultChan <- allocateResult{addr: addr[0]}
-			return
+		for _, addr := range alloc.owned[ident] {
+			if subnet.Contains(addr) {
+				resultChan <- allocateResult{addr: addr}
+				return
+			}
 		}
 		resultChan <- allocateResult{err: fmt.Errorf("lookup: no address found")}
 	}
diff --git a/ipam/http.go b/ipam/http.go
index a851f35d..35cdb437 100644
--- a/ipam/http.go
+++ b/ipam/http.go
@@ -33,12 +33,28 @@ func (alloc *Allocator) HandleHTTP(router *mux.Router) {
 		w.WriteHeader(204)
 	})
 
+	router.Methods("GET").Path("/ip/{id}/{ip}/{prefixlen}").HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		vars := mux.Vars(r)
+		cidr := vars["ip"] + "/" + vars["prefixlen"]
+		_, subnet, err := address.ParseCIDR(cidr)
+		if err != nil {
+			badRequest(w, err)
+			return
+		}
+		addr, err := alloc.Lookup(vars["id"], subnet)
+		if err != nil {
+			http.NotFound(w, r)
+			return
+		}
+		fmt.Fprintf(w, "%s/%d", addr.String(), subnet.PrefixLen)
+	})
+
 	router.Methods("GET").Path("/ip/{id}").HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if alloc.defaultSubnet.Blank() {
 			badRequest(w, fmt.Errorf("No subnet specified"))
 			return
 		}
-		addr, err := alloc.Lookup(mux.Vars(r)["id"])
+		addr, err := alloc.Lookup(mux.Vars(r)["id"], alloc.defaultSubnet)
 		if err != nil {
 			http.NotFound(w, r)
 			return
diff --git a/ipam/http_test.go b/ipam/http_test.go
index c66b2d85..19efaa95 100644
--- a/ipam/http_test.go
+++ b/ipam/http_test.go
@@ -89,6 +89,11 @@ func TestHttp(t *testing.T) {
 	cidr2a := HTTPPost(t, allocURL(port, testCIDR2, containerID))
 	wt.AssertEqualString(t, cidr2a, testAddr2, "address")
 
+	check := HTTPGet(t, allocURL(port, testCIDR1, containerID))
+	wt.AssertEqualString(t, check, cidr1a, "address")
+	check = HTTPGet(t, allocURL(port, testCIDR2, containerID))
+	wt.AssertEqualString(t, check, cidr2a, "address")
+
 	// Ask the http server for a pair of addresses for another container and check they're different
 	cidr1b := HTTPPost(t, allocURL(port, testCIDR1, container2))
 	wt.AssertFalse(t, cidr1b == testAddr1, "address")