From c6e1651d9e313ce0dc603fea3a5b174745f74070 Mon Sep 17 00:00:00 2001
From: Ce Gao <ce.gao@outlook.com>
Date: Mon, 11 Oct 2021 15:58:29 +0800
Subject: [PATCH] fix: Fix the empty key error

---
 middlewares/authenticate.go      | 19 +++++++----
 middlewares/authenticate_test.go | 56 ++++++++++++++++++++++++++++----
 2 files changed, 63 insertions(+), 12 deletions(-)

diff --git a/middlewares/authenticate.go b/middlewares/authenticate.go
index d99a614..fc25c7a 100644
--- a/middlewares/authenticate.go
+++ b/middlewares/authenticate.go
@@ -1,6 +1,7 @@
 package middlewares
 
 import (
+	"fmt"
 	"net/http"
 	"strings"
 
@@ -10,6 +11,10 @@ import (
 	"github.com/muety/wakapi/utils"
 )
 
+var (
+	errEmptyKey = fmt.Errorf("the api_key is empty")
+)
+
 type AuthenticateMiddleware struct {
 	config           *conf.Config
 	userSrvc         services.IUserService
@@ -46,10 +51,10 @@ func (m *AuthenticateMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Reques
 	user, err := m.tryGetUserByCookie(r)
 
 	if err != nil {
-		user, err = m.tryGetUserByApiKey(r)
+		user, err = m.tryGetUserByApiKeyHeader(r)
 	}
 	if err != nil {
-		user, err = m.tryGetUserByQueryParameter(r)
+		user, err = m.tryGetUserByApiKeyQuery(r)
 	}
 
 	if err != nil || user == nil {
@@ -81,7 +86,7 @@ func (m *AuthenticateMiddleware) isOptional(requestPath string) bool {
 	return false
 }
 
-func (m *AuthenticateMiddleware) tryGetUserByApiKey(r *http.Request) (*models.User, error) {
+func (m *AuthenticateMiddleware) tryGetUserByApiKeyHeader(r *http.Request) (*models.User, error) {
 	key, err := utils.ExtractBearerAuth(r)
 	if err != nil {
 		return nil, err
@@ -96,11 +101,13 @@ func (m *AuthenticateMiddleware) tryGetUserByApiKey(r *http.Request) (*models.Us
 	return user, nil
 }
 
-func (m *AuthenticateMiddleware) tryGetUserByQueryParameter(r *http.Request) (*models.User, error) {
-	key := r.URL.Query().Get("token")
-
+func (m *AuthenticateMiddleware) tryGetUserByApiKeyQuery(r *http.Request) (*models.User, error) {
+	key := r.URL.Query().Get("api_token")
 	var user *models.User
 	userKey := strings.TrimSpace(key)
+	if userKey == "" {
+		return nil, errEmptyKey
+	}
 	user, err := m.userSrvc.GetUserByKey(userKey)
 	if err != nil {
 		return nil, err
diff --git a/middlewares/authenticate_test.go b/middlewares/authenticate_test.go
index 2d4201f..019a5b8 100644
--- a/middlewares/authenticate_test.go
+++ b/middlewares/authenticate_test.go
@@ -3,14 +3,16 @@ package middlewares
 import (
 	"encoding/base64"
 	"fmt"
+	"net/http"
+	"net/url"
+	"testing"
+
 	"github.com/muety/wakapi/mocks"
 	"github.com/muety/wakapi/models"
 	"github.com/stretchr/testify/assert"
-	"net/http"
-	"testing"
 )
 
-func TestAuthenticateMiddleware_tryGetUserByApiKey_Success(t *testing.T) {
+func TestAuthenticateMiddleware_tryGetUserByApiKeyHeader_Success(t *testing.T) {
 	testApiKey := "z5uig69cn9ut93n"
 	testToken := base64.StdEncoding.EncodeToString([]byte(testApiKey))
 	testUser := &models.User{ApiKey: testApiKey}
@@ -26,13 +28,13 @@ func TestAuthenticateMiddleware_tryGetUserByApiKey_Success(t *testing.T) {
 
 	sut := NewAuthenticateMiddleware(userServiceMock)
 
-	result, err := sut.tryGetUserByApiKey(mockRequest)
+	result, err := sut.tryGetUserByApiKeyHeader(mockRequest)
 
 	assert.Nil(t, err)
 	assert.Equal(t, testUser, result)
 }
 
-func TestAuthenticateMiddleware_tryGetUserByApiKey_InvalidHeader(t *testing.T) {
+func TestAuthenticateMiddleware_tryGetUserByApiKeyHeader_Invalid(t *testing.T) {
 	testApiKey := "z5uig69cn9ut93n"
 	testToken := base64.StdEncoding.EncodeToString([]byte(testApiKey))
 
@@ -47,10 +49,52 @@ func TestAuthenticateMiddleware_tryGetUserByApiKey_InvalidHeader(t *testing.T) {
 
 	sut := NewAuthenticateMiddleware(userServiceMock)
 
-	result, err := sut.tryGetUserByApiKey(mockRequest)
+	result, err := sut.tryGetUserByApiKeyHeader(mockRequest)
 
 	assert.Error(t, err)
 	assert.Nil(t, result)
 }
 
+func TestAuthenticateMiddleware_tryGetUserByApiKeyQuery_Success(t *testing.T) {
+	testApiKey := "z5uig69cn9ut93n"
+	testUser := &models.User{ApiKey: testApiKey}
+
+	mockRequest := &http.Request{
+		URL: &url.URL{
+			RawQuery: fmt.Sprintf("api_token=%s", testApiKey),
+		},
+	}
+
+	userServiceMock := new(mocks.UserServiceMock)
+	userServiceMock.On("GetUserByKey", testApiKey).Return(testUser, nil)
+
+	sut := NewAuthenticateMiddleware(userServiceMock)
+
+	result, err := sut.tryGetUserByApiKeyQuery(mockRequest)
+
+	assert.Nil(t, err)
+	assert.Equal(t, testUser, result)
+}
+
+func TestAuthenticateMiddleware_tryGetUserByApiKeyQuery_Invalid(t *testing.T) {
+	testApiKey := "z5uig69cn9ut93n"
+
+	mockRequest := &http.Request{
+		URL: &url.URL{
+			// Use the wrong parameter name.
+			RawQuery: fmt.Sprintf("token=%s", testApiKey),
+		},
+	}
+
+	userServiceMock := new(mocks.UserServiceMock)
+
+	sut := NewAuthenticateMiddleware(userServiceMock)
+
+	result, actualErr := sut.tryGetUserByApiKeyQuery(mockRequest)
+
+	assert.Error(t, actualErr)
+	assert.Equal(t, errEmptyKey, actualErr)
+	assert.Nil(t, result)
+}
+
 // TODO: somehow test cookie auth function