From 6389c8cf75c3794fa75b41dcb7974afb00f6e24c Mon Sep 17 00:00:00 2001
From: David Lawrence <david.lawrence@docker.com>
Date: Thu, 21 Jan 2016 11:44:00 -0800
Subject: [PATCH] tokenAuth should also 'succeed' if we get a 401, which will
 result in attempting futher authentication later Signed-off-by: David
 Lawrence <david.lawrence@docker.com> (github: endophage)

---
 cmd/notary/tuf.go      |  5 +++--
 cmd/notary/tuf_test.go | 35 ++++++++++++++++++++++++++++++++++-
 2 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/cmd/notary/tuf.go b/cmd/notary/tuf.go
index 2ffbfa5acc..2e794f1191 100644
--- a/cmd/notary/tuf.go
+++ b/cmd/notary/tuf.go
@@ -432,8 +432,9 @@ func tokenAuth(trustServerURL string, baseTransport *http.Transport, gun string,
 	}
 	// non-nil err means we must close body
 	defer resp.Body.Close()
-	if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
-		// If we didn't get a 2XX range status code, we're not talking to a notary server.
+	if (resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices) &&
+		resp.StatusCode != http.StatusUnauthorized {
+		// If we didn't get a 2XX range or 401 status code, we're not talking to a notary server.
 		// The http client should be configured to handle redirects so at this point, 3XX is
 		// not a valid status code.
 		logrus.Errorf("could not reach %s: %d", trustServerURL, resp.StatusCode)
diff --git a/cmd/notary/tuf_test.go b/cmd/notary/tuf_test.go
index 9b42d720c1..a51874604b 100644
--- a/cmd/notary/tuf_test.go
+++ b/cmd/notary/tuf_test.go
@@ -17,11 +17,44 @@ func TestTokenAuth(t *testing.T) {
 	require.Nil(t, tokenAuth("https://localhost:9999", baseTransport, gun, readOnly))
 }
 
+func StatusOKTestHandler(w http.ResponseWriter, r *http.Request) {
+	w.WriteHeader(200)
+	w.Write([]byte("{}"))
+}
+
+func TestTokenAuth200Status(t *testing.T) {
+	var (
+		readOnly      bool
+		baseTransport = &http.Transport{}
+		gun           = "test"
+	)
+	s := httptest.NewServer(http.HandlerFunc(NotAuthorizedTestHandler))
+	defer s.Close()
+
+	require.NotNil(t, tokenAuth(s.URL, baseTransport, gun, readOnly))
+}
+
+func NotAuthorizedTestHandler(w http.ResponseWriter, r *http.Request) {
+	w.WriteHeader(401)
+}
+
+func TestTokenAuth401Status(t *testing.T) {
+	var (
+		readOnly      bool
+		baseTransport = &http.Transport{}
+		gun           = "test"
+	)
+	s := httptest.NewServer(http.HandlerFunc(NotAuthorizedTestHandler))
+	defer s.Close()
+
+	require.NotNil(t, tokenAuth(s.URL, baseTransport, gun, readOnly))
+}
+
 func NotFoundTestHandler(w http.ResponseWriter, r *http.Request) {
 	w.WriteHeader(404)
 }
 
-func TestTokenAuthNon200Status(t *testing.T) {
+func TestTokenAuthNon200Non401Status(t *testing.T) {
 	var (
 		readOnly      bool
 		baseTransport = &http.Transport{}