From 15c4666ace65c187d57fb1f64171fdcbf4437211 Mon Sep 17 00:00:00 2001
From: ae <git@golfed.xyz>
Date: Wed, 2 Apr 2025 13:18:30 +0300
Subject: [PATCH] feat: includeUser URL parameter for login handler

---
 server/pkg/service/users.go      | 30 +++++++++++++++++++---
 server/pkg/service/users_test.go | 44 +++++++++++++++++++++++++++++---
 2 files changed, 68 insertions(+), 6 deletions(-)

diff --git a/server/pkg/service/users.go b/server/pkg/service/users.go
index 1a178b7..fcb8cff 100644
--- a/server/pkg/service/users.go
+++ b/server/pkg/service/users.go
@@ -5,6 +5,8 @@ import (
 	"encoding/json"
 	"errors"
 	"net/http"
+	"strconv"
+	"time"
 
 	"git.umbrella.haus/ae/notatest/pkg/data"
 	"github.com/go-chi/chi/v5"
@@ -16,6 +18,15 @@ import (
 
 type userCtxKey struct{}
 
+// Stripped object that only contains non-critical data
+type userResponse struct {
+	ID        uuid.UUID  `json:"id"`
+	Username  string     `json:"username"`
+	CreatedAt *time.Time `json:"created_at"`
+	UpdatedAt *time.Time `json:"updated_at"`
+	IsAdmin   bool       `json:"is_admin"`
+}
+
 // Mockable database operations interface
 type UserStore interface {
 	CreateUser(ctx context.Context, arg data.CreateUserParams) (data.User, error)
@@ -160,10 +171,23 @@ func (rs usersResource) Login(w http.ResponseWriter, r *http.Request) {
 		SameSite: http.SameSiteStrictMode,
 	})
 
-	// Return the access token in the response body (it should be stored in browser's memory client-side)
-	respondJSON(w, http.StatusOK, map[string]string{
+	// Build response
+	response := map[string]any{
 		"access_token": tokenPair.AccessToken,
-	})
+	}
+
+	// Include user data if the client has requested it (`?includeUser=true`)
+	if includeUser, _ := strconv.ParseBool(r.URL.Query().Get("includeUser")); includeUser {
+		response["user"] = userResponse{
+			ID:        user.ID,
+			Username:  user.Username,
+			CreatedAt: user.CreatedAt,
+			UpdatedAt: user.UpdatedAt,
+			IsAdmin:   user.IsAdmin,
+		}
+	}
+
+	respondJSON(w, http.StatusOK, response)
 }
 
 func (rs usersResource) List(w http.ResponseWriter, r *http.Request) {
diff --git a/server/pkg/service/users_test.go b/server/pkg/service/users_test.go
index e12e9dc..23bf8ab 100644
--- a/server/pkg/service/users_test.go
+++ b/server/pkg/service/users_test.go
@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"net/http"
 	"net/http/httptest"
+	"net/url"
 	"strings"
 	"testing"
 
@@ -295,7 +296,9 @@ func TestUsersLogin(t *testing.T) {
 
 	tests := []struct {
 		name         string
-		requestBody  interface{}
+		includeUser  string
+		wantUserData bool
+		requestBody  any
 		mockSetup    func(*mockUserStore)
 		wantStatus   int
 		wantResponse string
@@ -337,7 +340,25 @@ func TestUsersLogin(t *testing.T) {
 			wantResponse: `{"error":"Invalid credentials"}`,
 		},
 		{
-			name: "successful login",
+			name:         "successful login with user data",
+			includeUser:  "true",
+			wantUserData: true,
+			requestBody: map[string]string{
+				"username": testUser.Username,
+				"password": validPassword,
+			},
+			mockSetup: func(m *mockUserStore) {
+				m.GetUserByUsernameFunc = func(_ context.Context, username string) (data.User, error) {
+					return testUser, nil
+				}
+			},
+			wantStatus:  http.StatusOK,
+			checkCookie: true,
+		},
+		{
+			name:         "successful login without user data",
+			includeUser:  "false",
+			wantUserData: false,
 			requestBody: map[string]string{
 				"username": testUser.Username,
 				"password": validPassword,
@@ -366,6 +387,11 @@ func TestUsersLogin(t *testing.T) {
 			req := httptest.NewRequest("POST", "/login", bytes.NewReader(body))
 			req.Header.Set("Content-Type", "application/json")
 
+			// Add the necessary query parameters
+			q := url.Values{}
+			q.Add("includeUser", tt.includeUser)
+			req.URL.RawQuery = q.Encode()
+
 			w := httptest.NewRecorder()
 			rs.Login(w, req)
 
@@ -377,6 +403,18 @@ func TestUsersLogin(t *testing.T) {
 				t.Errorf("expected response %q, got %q", tt.wantResponse, w.Body.String())
 			}
 
+			if tt.wantUserData {
+				var response struct {
+					AccessToken string    `json:"access_token"`
+					User        data.User `json:"user"` // Cast to the "raw" type to allow checking for sensitive data fields
+				}
+				json.Unmarshal(w.Body.Bytes(), &response)
+
+				assert.Equal(t, testUser.ID, response.User.ID)
+				assert.Equal(t, testUser.Username, response.User.Username)
+				assert.Empty(t, response.User.PasswordHash) // Ensure sensitive data excluded
+			}
+
 			if tt.checkCookie {
 				cookies := w.Result().Cookies()
 				var refreshCookie *http.Cookie
@@ -407,7 +445,7 @@ func TestUsersLogin(t *testing.T) {
 				token, err := jwt.ParseWithClaims(
 					response["access_token"],
 					&userClaims{},
-					func(token *jwt.Token) (interface{}, error) {
+					func(token *jwt.Token) (any, error) {
 						return []byte(jwtSecret), nil
 					},
 				)