From ca7a865f58944a34f39e7da89ffcd7bbdbca4d93 Mon Sep 17 00:00:00 2001
From: Mark Sagi-Kazar <mark.sagikazar@gmail.com>
Date: Tue, 26 Sep 2023 00:13:43 +0200
Subject: [PATCH] feat: AuthorizationServer implements http.Handler

Signed-off-by: Mark Sagi-Kazar <mark.sagikazar@gmail.com>
---
 auth/server.go      | 19 +++++++++++++++++++
 auth/server_test.go | 20 +++-----------------
 2 files changed, 22 insertions(+), 17 deletions(-)

diff --git a/auth/server.go b/auth/server.go
index 78d45d5..02d709a 100644
--- a/auth/server.go
+++ b/auth/server.go
@@ -175,3 +175,22 @@ type rawOAuth2Request struct {
 	Password     string `schema:"password"`
 	RefreshToken string `schema:"refresh_token"`
 }
+
+// ServeHTTP implements the [http.Handler] interface.
+//
+// Use it to register the AuthorizationServer directly as an HTTP handler.
+// Otherwise, register the handler in an HTTP router directly:
+//   - GET / -> [AuthorizationServer.TokenHandler]
+//   - POST / -> [AuthorizationServer.OAuth2Handler]
+func (s AuthorizationServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	switch r.Method {
+	case http.MethodGet:
+		s.TokenHandler(w, r)
+
+	case http.MethodPost:
+		s.OAuth2Handler(w, r)
+
+	default:
+		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
+	}
+}
diff --git a/auth/server_test.go b/auth/server_test.go
index 3c36f17..85397df 100644
--- a/auth/server_test.go
+++ b/auth/server_test.go
@@ -109,26 +109,12 @@ func TestAuthorizationServer(t *testing.T) {
 		Service: service,
 	}
 
-	router := http.NewServeMux()
-	router.HandleFunc("/token", func(w http.ResponseWriter, r *http.Request) {
-		switch r.Method {
-		case http.MethodGet:
-			server.TokenHandler(w, r)
-
-		case http.MethodPost:
-			server.OAuth2Handler(w, r)
-
-		default:
-			http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
-		}
-	})
-
-	httpServer := httptest.NewServer(router)
+	httpServer := httptest.NewServer(server)
 
 	t.Run("TokenHandler", func(t *testing.T) {
 		t.Parallel()
 
-		request, err := http.NewRequest(http.MethodGet, httpServer.URL+"/token", nil)
+		request, err := http.NewRequest(http.MethodGet, httpServer.URL, nil)
 		require.NoError(t, err)
 
 		request.SetBasicAuth("user", "password")
@@ -275,7 +261,7 @@ func TestAuthorizationServer(t *testing.T) {
 	t.Run("Oauth2Handler", func(t *testing.T) {
 		t.Parallel()
 
-		request, err := http.NewRequest(http.MethodPost, httpServer.URL+"/token", nil)
+		request, err := http.NewRequest(http.MethodPost, httpServer.URL, nil)
 		require.NoError(t, err)
 
 		request.Header.Add("Content-Type", "application/x-www-form-urlencoded")