Disco client now follows HTTP redirects

disco/client.go

@@ -7,7 +7,9 @@ import (
 	"errors"
 	"fmt"
 	"io/ioutil"
+	"log"
 	"net/http"
+	"os"
 )
 
 // Response represents the response returned by a Discovery Service.
@@ -19,13 +21,15 @@ type Response struct {
 
 // Client provides a Discovery Service client.
 type Client struct {
-	url string
+	url    string
+	logger *log.Logger
 }
 
 // New returns an initialized Discovery Service client.
 func New(url string) *Client {
 	return &Client{
-		url: url,
+		url:    url,
+		logger: log.New(os.Stderr, "[discovery] ", log.LstdFlags),
 	}
 }
 
@@ -40,35 +44,44 @@ func (c *Client) Register(id, addr string) (*Response, error) {
 	m := map[string]string{
 		"addr": addr,
 	}
-	b, err := json.Marshal(m)
-	if err != nil {
-		return nil, err
-	}
 
-	url := c.registrationURL(id)
-	resp, err := http.Post(url, "application-type/json", bytes.NewReader(b))
-	if err != nil {
-		return nil, err
-	}
-	defer resp.Body.Close()
+	url := c.registrationURL(c.url, id)
 
-	b, err = ioutil.ReadAll(resp.Body)
-	if err != nil {
-		return nil, err
-	}
+	for {
+		b, err := json.Marshal(m)
+		if err != nil {
+			return nil, err
+		}
 
-	if resp.StatusCode != http.StatusOK {
-		return nil, errors.New(resp.Status)
-	}
+		c.logger.Println("discovery client attempting registration at", url)
+		resp, err := http.Post(url, "application-type/json", bytes.NewReader(b))
+		if err != nil {
+			return nil, err
+		}
+		defer resp.Body.Close()
 
-	r := &Response{}
-	if err := json.Unmarshal(b, r); err != nil {
-		return nil, err
-	}
+		b, err = ioutil.ReadAll(resp.Body)
+		if err != nil {
+			return nil, err
+		}
 
-	return r, nil
+		switch resp.StatusCode {
+		case http.StatusOK:
+			r := &Response{}
+			if err := json.Unmarshal(b, r); err != nil {
+				return nil, err
+			}
+			return r, nil
+		case http.StatusMovedPermanently:
+			url = c.registrationURL(resp.Header.Get("location"), id)
+			c.logger.Println("discovery client redirecting to", url)
+			continue
+		default:
+			return nil, errors.New(resp.Status)
+		}
+	}
 }
 
-func (c *Client) registrationURL(id string) string {
-	return fmt.Sprintf("%s/%s", c.url, id)
+func (c *Client) registrationURL(url, id string) string {
+	return fmt.Sprintf("%s/%s", url, id)
 }

disco/client_test.go

@@ -106,6 +106,47 @@ func Test_ClientRegisterRequestOK(t *testing.T) {
 	}
 }
 
+func Test_ClientRegisterRequestRedirectOK(t *testing.T) {
+	ts1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.Method != "POST" {
+			t.Fatalf("Client did not use POST")
+		}
+
+		if r.URL.String() != "/1234" {
+			t.Fatalf("Request URL is wrong, got: %s", r.URL.String())
+		}
+
+		b, err := ioutil.ReadAll(r.Body)
+		if err != nil {
+			t.Fatalf("failed to read request from client: %s", err.Error())
+		}
+
+		m := map[string]string{}
+		if err := json.Unmarshal(b, &m); err != nil {
+			t.Fatalf("failed to unmarshal request from client: %s", err.Error())
+		}
+
+		if m["addr"] != "http://127.0.0.1" {
+			t.Fatalf("incorrect join address supplied by client: %s", m["addr"])
+		}
+
+		fmt.Fprintln(w, `{"created_at": "2017-02-17 04:49:05.079125", "disco_id": "68d6c7cc-f4cc-11e6-a170-2e79ea0be7b1", "nodes": ["http://127.0.0.1"]}`)
+	}))
+	defer ts1.Close()
+	ts2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		http.Redirect(w, r, ts1.URL, http.StatusMovedPermanently)
+	}))
+
+	c := New(ts2.URL)
+	disco, err := c.Register("1234", "http://127.0.0.1")
+	if err != nil {
+		t.Fatalf("failed to register: %s", err.Error())
+	}
+	if len(disco.Nodes) != 1 {
+		t.Fatalf("failed to receive correct list of nodes, got %v", disco.Nodes)
+	}
+}
+
 func Test_ClientRegisterFollowerOK(t *testing.T) {
 	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if r.Method != "POST" {