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 }