Oauth 2.0: Locking for token requests

handler/transport/oauth2_req_auth.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"net/http"
+	"sync"
 
 	"github.com/avenga/couper/cache"
 	"github.com/avenga/couper/config"
@@ -19,6 +20,7 @@ type OAuth2ReqAuth struct {
 	oauth2Client *oauth2.Client
 	config       *config.OAuth2ReqAuth
 	memStore     *cache.MemoryStore
+	locks        sync.Map
 	next         http.RoundTripper
 }
 
@@ -29,23 +31,36 @@ func NewOAuth2ReqAuth(conf *config.OAuth2ReqAuth, memStore *cache.MemoryStore,
 		config:       conf,
 		oauth2Client: oauth2Client,
 		memStore:     memStore,
+		locks:        sync.Map{},
 		next:         next,
 	}, nil
 }
 
 // RoundTrip implements the <http.RoundTripper> interface.
 func (oa *OAuth2ReqAuth) RoundTrip(req *http.Request) (*http.Response, error) {
 	storageKey := fmt.Sprintf("%p|%s|%s", &oa.oauth2Client.Backend, oa.config.ClientID, oa.config.ClientSecret)
-	if data := oa.memStore.Get(storageKey); data != "" {
-		token, terr := oa.readAccessToken(data)
-		if terr != nil {
-			// TODO this error is not connected to the OAuth2 client's backend
-			// In fact this can only be a JSON parse error or a missing access_token,
-			// which will occur after having requested the token from the authorization
-			// server. So the erroneous response will never be stored.
-			return nil, errors.Backend.Label(oa.config.BackendName).Message("token read error").With(terr)
-		}
 
+	if token, terr := oa.readAccessToken(storageKey); terr != nil {
+		// TODO this error is not connected to the OAuth2 client's backend
+		// In fact this can only be a JSON parse error or a missing access_token,
+		// which will occur after having requested the token from the authorization
+		// server. So the erroneous response will never be stored.
+		return nil, errors.Backend.Label(oa.config.BackendName).Message("token read error").With(terr)
+	} else if token != "" {
+		req.Header.Set("Authorization", "Bearer "+token)
+
+		return oa.next.RoundTrip(req)
+	}
+
+	value, _ := oa.locks.LoadOrStore(storageKey, &sync.Mutex{})
+	mutex := value.(*sync.Mutex)
+	mutex.Lock()
+
+	if token, terr := oa.readAccessToken(storageKey); terr != nil {
+		mutex.Unlock()
+		return nil, errors.Backend.Label(oa.config.BackendName).Message("token read error").With(terr)
+	} else if token != "" {
+		mutex.Unlock()
 		req.Header.Set("Authorization", "Bearer "+token)
 
 		return oa.next.RoundTrip(req)
@@ -54,11 +69,14 @@ func (oa *OAuth2ReqAuth) RoundTrip(req *http.Request) (*http.Response, error) {
 	ctx := req.Context()
 	tokenResponse, tokenResponseData, token, err := oa.oauth2Client.GetTokenResponse(ctx)
 	if err != nil {
+		mutex.Unlock()
 		return nil, errors.Backend.Label(oa.config.BackendName).Message("token request error").With(err)
 	}
 
 	oa.updateAccessToken(tokenResponse, tokenResponseData, storageKey)
 
+	mutex.Unlock()
+
 	req.Header.Set("Authorization", "Bearer "+token)
 
 	res, err := oa.next.RoundTrip(req)
@@ -79,13 +97,17 @@ func (oa *OAuth2ReqAuth) RoundTrip(req *http.Request) (*http.Response, error) {
 	return res, err
 }
 
-func (oa *OAuth2ReqAuth) readAccessToken(data string) (string, error) {
-	_, token, err := oauth2.ParseTokenResponse([]byte(data))
-	if err != nil {
-		return "", err
+func (oa *OAuth2ReqAuth) readAccessToken(key string) (string, error) {
+	if data := oa.memStore.Get(key); data != "" {
+		_, token, err := oauth2.ParseTokenResponse([]byte(data))
+		if err != nil {
+			return "", errors.Backend.Label(oa.config.BackendName).Message("token read error").With(err)
+		}
+
+		return token, nil
 	}
 
-	return token, nil
+	return "", nil
 }
 
 func (oa *OAuth2ReqAuth) updateAccessToken(jsonBytes []byte, jData map[string]interface{}, key string) {

server/http_oauth2_test.go

@@ -7,6 +7,8 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"strings"
+	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -199,7 +201,7 @@ func TestEndpoints_OAuth2_Options(t *testing.T) {
 	}
 }
 
-func TestOAuth2AccessControl(t *testing.T) {
+func TestOAuth2_AccessControl(t *testing.T) {
 	client := newClient()
 	helper := test.New(t)
 
@@ -383,3 +385,182 @@ func TestOAuth2AccessControl(t *testing.T) {
 		})
 	}
 }
+
+func TestOAuth2_Locking(t *testing.T) {
+	helper := test.New(t)
+	token := "token-"
+	var oauthRequestCount int32
+	oauthOrigin := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		atomic.AddInt32(&oauthRequestCount, 1)
+		if req.URL.Path == "/oauth2" {
+			rw.Header().Set("Content-Type", "application/json")
+			rw.WriteHeader(http.StatusOK)
+
+			n := fmt.Sprintf("%d", atomic.LoadInt32(&oauthRequestCount))
+			body := []byte(`{
+				"access_token": "` + token + n + `",
+				"token_type": "bearer",
+				"expires_in": 1.5
+			}`)
+
+			// Slow down token request to test locking.
+			time.Sleep(1 * time.Second)
+
+			_, werr := rw.Write(body)
+			helper.Must(werr)
+
+			return
+		}
+
+		rw.WriteHeader(http.StatusBadRequest)
+	}))
+	defer oauthOrigin.Close()
+
+	ResourceOrigin := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		if req.URL.Path == "/resource" {
+			if auth := req.Header.Get("Authorization"); auth != "" {
+				rw.Header().Set("Token", auth[len("Bearer "):])
+				rw.WriteHeader(http.StatusNoContent)
+			}
+
+			return
+		}
+
+		rw.WriteHeader(http.StatusNotFound)
+	}))
+	defer ResourceOrigin.Close()
+
+	confPath := "testdata/oauth2/1_retries_couper.hcl"
+	shutdown, hook := newCouperWithTemplate(
+		confPath, test.New(t), map[string]interface{}{
+			"asOrigin": oauthOrigin.URL,
+			"rsOrigin": ResourceOrigin.URL,
+		},
+	)
+
+	defer func() {
+		if t.Failed() {
+			for _, e := range hook.Entries {
+				println(e.String())
+			}
+		}
+		shutdown()
+	}()
+
+	req, err := http.NewRequest(http.MethodGet, "http://anyserver:8080/", nil)
+	helper.Must(err)
+
+	hook.Reset()
+
+	req.URL.Path = "/"
+
+	var responses []*http.Response
+	var wg sync.WaitGroup
+
+	addLock := &sync.Mutex{}
+	// Fire 5 requests in parallel...
+	waitCh := make(chan struct{})
+	for i := 0; i < 5; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			<-waitCh
+			res, err := newClient().Do(req)
+			helper.Must(err)
+
+			addLock.Lock()
+			responses = append(responses, res)
+			addLock.Unlock()
+
+		}()
+	}
+	close(waitCh)
+	wg.Wait()
+
+	for _, res := range responses {
+		if res.StatusCode != http.StatusNoContent {
+			t.Errorf("Expected status NoContent, got: %d", res.StatusCode)
+		}
+
+		if token+"1" != res.Header.Get("Token") {
+			t.Errorf("Invalid token given: want %s1, got: %s", token, res.Header.Get("Token"))
+		}
+	}
+
+	if oauthRequestCount != 1 {
+		t.Errorf("Too many OAuth2 requests: want 1, got: %d", oauthRequestCount)
+	}
+
+	t.Run("Lock is effective", func(t *testing.T) {
+		// Wait until token has expired.
+		time.Sleep(2 * time.Second)
+
+		// Fetch new token.
+		go func() {
+			res, err := newClient().Do(req)
+			helper.Must(err)
+
+			if token+"2" != res.Header.Get("Token") {
+				t.Errorf("Received wrong token: want %s2, got: %s", token, res.Header.Get("Token"))
+			}
+		}()
+
+		// Slow response due to lock
+		go func() {
+			start := time.Now()
+			res, err := newClient().Do(req)
+			helper.Must(err)
+			timeElapsed := time.Since(start).Seconds()
+
+			if token+"2" != res.Header.Get("Token") {
+				t.Errorf("Received wrong token: want %s2, got: %s", token, res.Header.Get("Token"))
+			}
+
+			if timeElapsed < 1 {
+				t.Errorf("Response came too fast: dysfunctional lock?! (%v s)", timeElapsed)
+			}
+		}()
+	})
+
+	t.Run("Mem store expiry", func(t *testing.T) {
+		// Wait again until token has expired.
+		time.Sleep(2 * time.Second)
+
+		// Request fresh token and store in memstore
+		res, err := newClient().Do(req)
+		helper.Must(err)
+		if res.StatusCode != http.StatusNoContent {
+			t.Errorf("Unexpected response status: want %d, got: %d", http.StatusNoContent, res.StatusCode)
+		}
+
+		if token+"3" != res.Header.Get("Token") {
+			t.Errorf("Received wrong token: want %s3, got: %s", token, res.Header.Get("Token"))
+		}
+
+		if oauthRequestCount != 3 {
+			t.Errorf("Unexpected number of OAuth2 requests: want 3, got: %d", oauthRequestCount)
+		}
+
+		// Disconnect OAuth server
+		oauthOrigin.Close()
+
+		// Next request gets token from memstore
+		res, err = newClient().Do(req)
+		helper.Must(err)
+		if res.StatusCode != http.StatusNoContent {
+			t.Errorf("Unexpected response status: want %d, got: %d", http.StatusNoContent, res.StatusCode)
+		}
+
+		if token+"3" != res.Header.Get("Token") {
+			t.Errorf("Wrong token from mem store: want %s3, got: %s", token, res.Header.Get("Token"))
+		}
+
+		// Wait until token has expired. Next request accesses the OAuth server again.
+		time.Sleep(2 * time.Second)
+		res, err = newClient().Do(req)
+		helper.Must(err)
+		if res.StatusCode != http.StatusBadGateway {
+			t.Errorf("Unexpected response status: want %d, got: %d", http.StatusBadGateway, res.StatusCode)
+		}
+	})
+}