From 50c018af627daff9416fe0afdcb7400a9ccd4595 Mon Sep 17 00:00:00 2001 From: Mazyar Date: Tue, 30 Jun 2026 22:30:28 +0330 Subject: [PATCH] Enhance company context handling in customer middleware and routing - Introduced CompanyContextMiddleware to resolve the active company context for customer requests, ensuring that tender recommendations and company-scoped APIs remain in sync with the database. - Updated public routes to utilize the new CompanyContextMiddleware alongside the existing AuthMiddleware, improving the handling of company-specific requests. - Added unit tests for the pickActiveCompanyID function to validate the logic for selecting the appropriate company context based on customer assignments and requested company IDs. This update enhances the accuracy and reliability of company context management in the application, improving user experience and data consistency. --- cmd/web/router/routes.go | 8 +-- internal/customer/company_context.go | 49 +++++++++++++++ internal/customer/company_context_test.go | 77 +++++++++++++++++++++++ internal/customer/middleware.go | 34 ++++++++++ internal/customer/service.go | 57 ++++++++++------- 5 files changed, 197 insertions(+), 28 deletions(-) create mode 100644 internal/customer/company_context.go create mode 100644 internal/customer/company_context_test.go diff --git a/cmd/web/router/routes.go b/cmd/web/router/routes.go index 996e312..9e144b1 100644 --- a/cmd/web/router/routes.go +++ b/cmd/web/router/routes.go @@ -281,7 +281,7 @@ func RegisterPublicRoutes(e *echo.Echo, customerHandler *customer.Handler, tende // Public Tenders Routes tendersGP := v1.Group("/tenders") - tendersGP.Use(customerHandler.AuthMiddleware()) + tendersGP.Use(customerHandler.AuthMiddleware(), customerHandler.CompanyContextMiddleware()) { tendersGP.GET("", tenderHandler.GetPublicTenders) tendersGP.GET("/recommend", tenderHandler.RecommendTenders) @@ -308,7 +308,7 @@ func RegisterPublicRoutes(e *echo.Echo, customerHandler *customer.Handler, tende // Public Tender Approvals Routes tenderApprovalGP := v1.Group("/tender-approvals") { - tenderApprovalGP.Use(customerHandler.AuthMiddleware()) + tenderApprovalGP.Use(customerHandler.AuthMiddleware(), customerHandler.CompanyContextMiddleware()) tenderApprovalGP.POST("", tenderApprovalHandler.ToggleTenderApproval) tenderApprovalGP.GET("", tenderApprovalHandler.PublicGetTenderApprovals) tenderApprovalGP.GET("/:id", tenderApprovalHandler.GetPublicTenderApproval) @@ -319,13 +319,13 @@ func RegisterPublicRoutes(e *echo.Echo, customerHandler *customer.Handler, tende // Public Company Routes companiesGP := v1.Group("/companies") { - companiesGP.Use(customerHandler.AuthMiddleware()) + companiesGP.Use(customerHandler.AuthMiddleware(), customerHandler.CompanyContextMiddleware()) companiesGP.GET("", companyHandler.MyCompany) } // AI tender recommendation routes recommendationGP := v1.Group("") - recommendationGP.Use(customerHandler.AuthMiddleware()) + recommendationGP.Use(customerHandler.AuthMiddleware(), customerHandler.CompanyContextMiddleware()) { recommendationGP.POST("/onboarding", companyHandler.StartOnboarding) recommendationGP.POST("/recommend", companyHandler.RecommendTenders) diff --git a/internal/customer/company_context.go b/internal/customer/company_context.go new file mode 100644 index 0000000..58a6e63 --- /dev/null +++ b/internal/customer/company_context.go @@ -0,0 +1,49 @@ +package customer + +import ( + "context" + "errors" + "strings" +) + +var errCompanyNotAssigned = errors.New("company is not assigned to customer") + +// pickActiveCompanyID chooses the company context for a customer request. +// When requestedCompanyID is set it must be in assigned; otherwise a still-valid +// token company is kept; if the token company was removed, the first assignment is used. +func pickActiveCompanyID(assigned []string, tokenCompanyID, requestedCompanyID string) (string, error) { + if len(assigned) == 0 { + return "", nil + } + + requested := strings.TrimSpace(requestedCompanyID) + if requested != "" { + for _, id := range assigned { + if id == requested { + return requested, nil + } + } + return "", errCompanyNotAssigned + } + + token := strings.TrimSpace(tokenCompanyID) + if token != "" { + for _, id := range assigned { + if id == token { + return token, nil + } + } + } + + return assigned[0], nil +} + +// ResolveActiveCompanyID resolves the active company for an authenticated customer. +func (s *customerService) ResolveActiveCompanyID(ctx context.Context, customerID, tokenCompanyID, requestedCompanyID string) (string, error) { + customer, err := s.repository.GetByID(ctx, customerID) + if err != nil { + return "", errors.New("customer not found") + } + + return pickActiveCompanyID(customer.Companies, tokenCompanyID, requestedCompanyID) +} diff --git a/internal/customer/company_context_test.go b/internal/customer/company_context_test.go new file mode 100644 index 0000000..f0321ed --- /dev/null +++ b/internal/customer/company_context_test.go @@ -0,0 +1,77 @@ +package customer + +import "testing" + +func TestPickActiveCompanyID(t *testing.T) { + assigned := []string{"company-a", "company-b"} + + tests := []struct { + name string + assigned []string + token string + requested string + want string + wantErr bool + wantErrType error + }{ + { + name: "uses requested company when assigned", + assigned: assigned, + token: "company-a", + requested: "company-b", + want: "company-b", + }, + { + name: "rejects requested company that is not assigned", + assigned: assigned, + token: "company-a", + requested: "company-c", + wantErr: true, + wantErrType: errCompanyNotAssigned, + }, + { + name: "keeps valid token company", + assigned: assigned, + token: "company-b", + want: "company-b", + }, + { + name: "falls back to first assignment when token company was removed", + assigned: []string{"company-b"}, + token: "company-a", + want: "company-b", + }, + { + name: "returns empty when customer has no companies", + assigned: nil, + token: "company-a", + want: "", + }, + { + name: "uses first assignment when token is empty", + assigned: assigned, + want: "company-a", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := pickActiveCompanyID(tt.assigned, tt.token, tt.requested) + if tt.wantErr { + if err == nil { + t.Fatalf("expected error, got nil") + } + if tt.wantErrType != nil && err != tt.wantErrType { + t.Fatalf("expected error %v, got %v", tt.wantErrType, err) + } + return + } + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if got != tt.want { + t.Fatalf("pickActiveCompanyID() = %q, want %q", got, tt.want) + } + }) + } +} diff --git a/internal/customer/middleware.go b/internal/customer/middleware.go index bd6b78e..76647b4 100644 --- a/internal/customer/middleware.go +++ b/internal/customer/middleware.go @@ -1,6 +1,7 @@ package customer import ( + "errors" "net/http" "strings" "tm/pkg/audit" @@ -61,6 +62,39 @@ func (h *Handler) AuthMiddleware() echo.MiddlewareFunc { } } +// CompanyContextMiddleware resolves company_id from current customer assignments. +// JWT company_id can be stale after admin reassigns companies; this keeps tender +// recommendations and other company-scoped APIs in sync with the database. +func (h *Handler) CompanyContextMiddleware() echo.MiddlewareFunc { + return func(next echo.HandlerFunc) echo.HandlerFunc { + return func(c echo.Context) error { + customerID, err := GetCustomerIDFromContext(c) + if err != nil { + return response.Unauthorized(c, "User not authenticated") + } + + tokenCompanyID, _ := c.Get("company_id").(string) + requestedCompanyID := strings.TrimSpace(c.QueryParam("company_id")) + + companyID, err := h.service.ResolveActiveCompanyID( + c.Request().Context(), + customerID, + tokenCompanyID, + requestedCompanyID, + ) + if err != nil { + if errors.Is(err, errCompanyNotAssigned) { + return response.Forbidden(c, "Company is not assigned to customer") + } + return response.Unauthorized(c, "User not authenticated") + } + + c.Set("company_id", companyID) + return next(c) + } + } +} + // GetCustomerIDFromContext extracts customer ID from Echo context func GetCustomerIDFromContext(c echo.Context) (string, error) { customerID, ok := c.Get("customer_id").(string) diff --git a/internal/customer/service.go b/internal/customer/service.go index 5d1521d..d917974 100644 --- a/internal/customer/service.go +++ b/internal/customer/service.go @@ -53,6 +53,9 @@ type Service interface { // Role assignment operations AssignRole(ctx context.Context, customerID string, form *AssignRoleForm) (*AssignRoleResponse, error) AssignCompanies(ctx context.Context, customerID string, Companies []string) error + + // ResolveActiveCompanyID returns the company context for an authenticated customer. + ResolveActiveCompanyID(ctx context.Context, customerID, tokenCompanyID, requestedCompanyID string) (string, error) } // customerService implements the Service interface @@ -797,10 +800,9 @@ func (s *customerService) Login(ctx context.Context, form *LoginForm) (*AuthResp return nil, errors.New("invalid credentials") } - // Generate JWT tokens using authorization service - companyID := "" - if len(customer.Companies) > 0 { - companyID = customer.Companies[0] // Use first company ID for JWT token + companyID, err := pickActiveCompanyID(customer.Companies, "", "") + if err != nil { + return nil, errors.New("failed to resolve customer company") } tokenPair, err := s.authService.GenerateTokenPair( @@ -863,32 +865,20 @@ func (s *customerService) RefreshToken(ctx context.Context, refreshToken string) "refresh_token": refreshToken[:10] + "...", // Log only first 10 chars for security }) - // Validate refresh token and generate new access token - newAccessToken, err := s.authService.RefreshAccessToken(refreshToken) + validationResult, err := s.authService.ValidateRefreshToken(refreshToken) if err != nil { - s.logger.Warn("Failed to refresh access token", map[string]interface{}{ + s.logger.Warn("Failed to validate refresh token", map[string]interface{}{ "error": err.Error(), }) return nil, errors.New("invalid or expired refresh token") } - - // Parse the new access token to get customer information - validationResult, err := s.authService.ValidateAccessToken(newAccessToken) - if err != nil { - s.logger.Error("Failed to validate new access token", map[string]interface{}{ - "error": err.Error(), - }) - return nil, errors.New("failed to generate valid access token") - } - if !validationResult.Valid { - s.logger.Error("New access token validation failed", map[string]interface{}{ + s.logger.Warn("Invalid refresh token provided", map[string]interface{}{ "error": validationResult.Error, }) - return nil, errors.New("failed to generate valid access token") + return nil, errors.New("invalid or expired refresh token") } - // Get customer information customerID, err := bson.ObjectIDFromHex(validationResult.Claims.UserID) if err != nil { s.logger.Error("Failed to parse customer ID from token", map[string]interface{}{ @@ -906,11 +896,29 @@ func (s *customerService) RefreshToken(ctx context.Context, refreshToken string) return nil, errors.New("customer not found") } - // Check if customer is still active if customer.Status != CustomerStatusActive { return nil, errors.New("customer account is inactive") } + companyID, err := pickActiveCompanyID(customer.Companies, validationResult.Claims.CompanyID, "") + if err != nil { + return nil, errors.New("failed to resolve customer company") + } + + tokenPair, err := s.authService.GenerateTokenPair( + customer.ID.Hex(), + customer.Email, + string(customer.Role), + companyID, + ) + if err != nil { + s.logger.Error("Failed to generate JWT tokens during refresh", map[string]interface{}{ + "error": err.Error(), + "customer_id": customer.ID, + }) + return nil, errors.New("failed to generate authentication tokens") + } + companies, err := s.loadCompaniesForCustomer(ctx, customer) if err != nil { s.logger.Error("Failed to load companies for customer", map[string]interface{}{ @@ -922,13 +930,14 @@ func (s *customerService) RefreshToken(ctx context.Context, refreshToken string) s.logger.Info("Access token refreshed successfully", map[string]interface{}{ "customer_id": customer.ID, "customer_type": string(customer.Type), + "company_id": companyID, }) return &AuthResponse{ Customer: customer.ToResponse(companies), - AccessToken: newAccessToken, - RefreshToken: refreshToken, // Return the same refresh token - ExpiresAt: time.Now().Add(15 * time.Minute).Unix(), // 15 minutes from now + AccessToken: tokenPair.AccessToken, + RefreshToken: tokenPair.RefreshToken, + ExpiresAt: time.Now().Add(15 * time.Minute).Unix(), }, nil }