From 44b8157572fbc049b04e15dc5be1c0acb644f6f1 Mon Sep 17 00:00:00 2001
From: spacexnice <yaoyao.xyy@alibaba-inc.com>
Date: Tue, 23 Aug 2016 22:09:18 +0800
Subject: [PATCH] fix simpleAuthChallge concurrent problem

Signed-off-by: spacexnice <yaoyao.xyy@alibaba-inc.com>
---
 registry/client/auth/authchallenge.go      | 22 ++++++---
 registry/client/auth/authchallenge_test.go | 57 ++++++++++++++++++++++
 2 files changed, 73 insertions(+), 6 deletions(-)

diff --git a/registry/client/auth/authchallenge.go b/registry/client/auth/authchallenge.go
index c8cd83bb..4dbda9e0 100644
--- a/registry/client/auth/authchallenge.go
+++ b/registry/client/auth/authchallenge.go
@@ -5,6 +5,7 @@ import (
 	"net/http"
 	"net/url"
 	"strings"
+	"sync"
 )
 
 // Challenge carries information from a WWW-Authenticate response header.
@@ -43,19 +44,26 @@ type ChallengeManager interface {
 // perform requests on the endpoints or cache the responses
 // to a backend.
 func NewSimpleChallengeManager() ChallengeManager {
-	return simpleChallengeManager{}
+	return &simpleChallengeManager{
+		Challanges: make(map[string][]Challenge),
+	}
 }
 
-type simpleChallengeManager map[string][]Challenge
+type simpleChallengeManager struct{
+	sync.RWMutex
+	Challanges  map[string][]Challenge
+}
 
-func (m simpleChallengeManager) GetChallenges(endpoint url.URL) ([]Challenge, error) {
+func (m *simpleChallengeManager) GetChallenges(endpoint url.URL) ([]Challenge, error) {
 	endpoint.Host = strings.ToLower(endpoint.Host)
 
-	challenges := m[endpoint.String()]
+	m.RLock()
+	defer m.RUnlock()
+	challenges := m.Challanges[endpoint.String()]
 	return challenges, nil
 }
 
-func (m simpleChallengeManager) AddResponse(resp *http.Response) error {
+func (m *simpleChallengeManager) AddResponse(resp *http.Response) error {
 	challenges := ResponseChallenges(resp)
 	if resp.Request == nil {
 		return fmt.Errorf("missing request reference")
@@ -65,7 +73,9 @@ func (m simpleChallengeManager) AddResponse(resp *http.Response) error {
 		Host:   strings.ToLower(resp.Request.URL.Host),
 		Scheme: resp.Request.URL.Scheme,
 	}
-	m[urlCopy.String()] = challenges
+	m.Lock()
+	defer m.Unlock()
+	m.Challanges[urlCopy.String()] = challenges
 	return nil
 }
 
diff --git a/registry/client/auth/authchallenge_test.go b/registry/client/auth/authchallenge_test.go
index 953ed5b4..e4f81d28 100644
--- a/registry/client/auth/authchallenge_test.go
+++ b/registry/client/auth/authchallenge_test.go
@@ -6,6 +6,7 @@ import (
 	"net/url"
 	"strings"
 	"testing"
+	"sync"
 )
 
 func TestAuthChallengeParse(t *testing.T) {
@@ -43,6 +44,7 @@ func TestAuthChallengeParse(t *testing.T) {
 func TestAuthChallengeNormalization(t *testing.T) {
 	testAuthChallengeNormalization(t, "reg.EXAMPLE.com")
 	testAuthChallengeNormalization(t, "bɿɒʜɔiɿ-ɿɘƚƨim-ƚol-ɒ-ƨʞnɒʜƚ.com")
+	testAuthChallengeConcurrent(t, "reg.EXAMPLE.com")
 }
 
 func testAuthChallengeNormalization(t *testing.T, host string) {
@@ -79,3 +81,58 @@ func testAuthChallengeNormalization(t *testing.T, host string) {
 		t.Fatal("Expected challenge for lower-cased-host URL")
 	}
 }
+
+
+func testAuthChallengeConcurrent(t *testing.T, host string) {
+
+	scm := NewSimpleChallengeManager()
+
+	url, err := url.Parse(fmt.Sprintf("http://%s/v2/", host))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	resp := &http.Response{
+		Request: &http.Request{
+			URL: url,
+		},
+		Header:     make(http.Header),
+		StatusCode: http.StatusUnauthorized,
+	}
+	resp.Header.Add("WWW-Authenticate", fmt.Sprintf("Bearer realm=\"https://%s/token\",service=\"registry.example.com\"", host))
+	var s sync.WaitGroup
+	s.Add(2)
+	go func() {
+		i := 200
+		for {
+			//time.Sleep(500 * time.Millisecond)
+			err = scm.AddResponse(resp)
+			if err != nil {
+				t.Fatal(err)
+			}
+			i = i -1
+			if i < 0{
+				break
+			}
+		}
+		s.Done()
+	}()
+	go func() {
+		lowered := *url
+		lowered.Host = strings.ToLower(lowered.Host)
+		k:= 200
+		for {
+			_, err := scm.GetChallenges(lowered)
+			if err != nil {
+				t.Fatal(err)
+			}
+			k = k -1
+			if k < 0{
+				break
+			}
+		}
+
+		s.Done()
+	}()
+	s.Wait()
+}