JWT Authentication and Token Flow

@RestController
@RequestMapping("/api/auth")
public class ExternalAuthController {

    private final AdminServiceClient adminServiceClient;
    private final HashValidationService hashService;

    @PostMapping("/external-login")
    public ResponseEntity<Void> externalLogin(
            @RequestBody ExternalLoginRequest request,
            HttpSession session) {

        // Validate hash server-side
        if (!hashService.validate(request.getUserId(), request.getUserHash())) {
            throw new UnauthorizedException("Invalid credentials");
        }

        // Exchange credentials with Admin Service
        TokenResponse tokenResponse = adminServiceClient.exchangeCredentials(
            ExchangeRequest.builder()
                .userId(request.getUserId())
                .build()
        );

        // Store JWT in HTTP session (NOT sent to client)
        session.setAttribute("JWT_TOKEN", tokenResponse.getToken());
        session.setAttribute("JWT_EXPIRES_AT", tokenResponse.getExpiresAt());

        return ResponseEntity.ok().build();
    }

    @PostMapping("/refresh")
    public ResponseEntity<Void> refreshToken(HttpSession session) {
        String currentJwt = (String) session.getAttribute("JWT_TOKEN");

        if (currentJwt == null) {
            throw new UnauthorizedException("No active session");
        }

        // Request new token from Admin Service
        TokenResponse newToken = adminServiceClient.refreshToken(currentJwt);

        // Update session with new token
        session.setAttribute("JWT_TOKEN", newToken.getToken());
        session.setAttribute("JWT_EXPIRES_AT", newToken.getExpiresAt());

        return ResponseEntity.ok().build();
    }

    @PostMapping("/logout")
    public ResponseEntity<Void> logout(HttpSession session) {
        session.invalidate();
        return ResponseEntity.ok().build();
    }
}

@Service
public class AdminServiceClient {

    private final WebClient webClient;
    private final String apiKey;

    public AdminServiceClient(WebClient.Builder builder,
                             @Value("${security.admin-service.api-key}") String apiKey,
                             @Value("${security.admin-service.base-url}") String baseUrl) {
        this.webClient = builder.baseUrl(baseUrl).build();
        this.apiKey = apiKey;
    }

    public TokenResponse exchangeCredentials(ExchangeRequest request) {
        return webClient.post()
            .uri("/api/auth/exchange")
            .header("Authorization", "ApiKey " + apiKey)
            .bodyValue(request)
            .retrieve()
            .onStatus(HttpStatusCode::is4xxClientError,
                response -> Mono.error(new UnauthorizedException("Invalid credentials")))
            .bodyToMono(TokenResponse.class)
            .block();
    }

    public TokenResponse refreshToken(String currentToken) {
        return webClient.post()
            .uri("/api/auth/refresh")
            .header("Authorization", "Bearer " + currentToken)
            .retrieve()
            .bodyToMono(TokenResponse.class)
            .block();
    }
}

@Configuration
public class GatewayFilterConfig {

    @Bean
    public GlobalFilter jwtInjectionFilter() {
        return new JwtGatewayFilter();
    }
}

@Component
@Order(-100) // High priority
public class JwtGatewayFilter implements GlobalFilter {

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        String path = exchange.getRequest().getPath().value();

        // Only inject JWT for admin-service routes
        if (!path.startsWith("/api/") && !path.startsWith("/services/")) {
            return chain.filter(exchange);
        }

        return exchange.getSession()
            .flatMap(session -> {
                String jwt = session.getAttribute("JWT_TOKEN");

                if (jwt != null) {
                    ServerHttpRequest modifiedRequest = exchange.getRequest()
                        .mutate()
                        .header("Authorization", "Bearer " + jwt)
                        .build();

                    return chain.filter(
                        exchange.mutate().request(modifiedRequest).build()
                    );
                }

                return chain.filter(exchange);
            });
    }
}

@RestController
@RequestMapping("/api/auth")
public class AuthController {

    private final JwtTokenProvider tokenProvider;
    private final ApiKeyValidator apiKeyValidator;
    private final UserService userService;

    @PostMapping("/exchange")
    public ResponseEntity<TokenResponse> exchangeCredentials(
            @RequestHeader("Authorization") String apiKeyHeader,
            @RequestBody ExchangeRequest request) {

        // Validate API key from gateway
        if (!apiKeyValidator.validate(apiKeyHeader)) {
            throw new UnauthorizedException("Invalid API key");
        }

        // Verify user exists
        User user = userService.findByUserId(request.getUserId())
            .orElseThrow(() -> new NotFoundException("User not found"));

        // Generate JWT with user claims
        String jwt = tokenProvider.createToken(
            JwtClaims.builder()
                .subject(user.getId().toString())
                .userId(user.getUserId())
                .scope("public")
                .authorities(user.getAuthorities())
                .build()
        );

        return ResponseEntity.ok(
            TokenResponse.builder()
                .token(jwt)
                .expiresIn(tokenProvider.getTokenValidity())
                .expiresAt(Instant.now().plusSeconds(tokenProvider.getTokenValidity()))
                .build()
        );
    }

    @PostMapping("/refresh")
    public ResponseEntity<TokenResponse> refreshToken(
            @AuthenticationPrincipal Jwt jwt) {

        // Extract user from current JWT
        String userId = jwt.getClaimAsString("userId");

        User user = userService.findByUserId(userId)
            .orElseThrow(() -> new NotFoundException("User not found"));

        // Issue new JWT
        String newJwt = tokenProvider.createToken(
            JwtClaims.builder()
                .subject(user.getId().toString())
                .userId(user.getUserId())
                .scope("public")
                .authorities(user.getAuthorities())
                .build()
        );

        return ResponseEntity.ok(
            TokenResponse.builder()
                .token(newJwt)
                .expiresIn(tokenProvider.getTokenValidity())
                .expiresAt(Instant.now().plusSeconds(tokenProvider.getTokenValidity()))
                .build()
        );
    }
}

@Service
public class JwtTokenProvider {

    private final Key signingKey;
    private final long tokenValidityInSeconds;

    public JwtTokenProvider(
            @Value("${security.jwt.secret}") String secret,
            @Value("${security.jwt.token-validity-seconds:1800}") long validity) {
        this.signingKey = Keys.hmacShaKeyFor(secret.getBytes());
        this.tokenValidityInSeconds = validity;
    }

    public String createToken(JwtClaims claims) {
        Instant now = Instant.now();
        Instant expiration = now.plusSeconds(tokenValidityInSeconds);

        return Jwts.builder()
            .setSubject(claims.getSubject())
            .claim("userId", claims.getUserId())
            .claim("scope", claims.getScope())
            .claim("authorities", claims.getAuthorities())
            .setIssuer("event-admin-service")
            .setAudience("registration-portal")
            .setIssuedAt(Date.from(now))
            .setExpiration(Date.from(expiration))
            .signWith(signingKey, SignatureAlgorithm.HS512)
            .compact();
    }

    public boolean validateToken(String token) {
        try {
            Jwts.parserBuilder()
                .setSigningKey(signingKey)
                .build()
                .parseClaimsJws(token);
            return true;
        } catch (JwtException | IllegalArgumentException e) {
            return false;
        }
    }

    public long getTokenValidity() {
        return tokenValidityInSeconds;
    }
}

@Configuration
@EnableWebSecurity
@EnableMethodSecurity
public class SecurityConfiguration {

    @Bean
    public SecurityFilterChain filterChain(HttpSecurity http) throws Exception {
        http
            .csrf(csrf -> csrf.disable()) // Stateless API
            .sessionManagement(session -> session
                .sessionCreationPolicy(SessionCreationPolicy.STATELESS)
            )
            .oauth2ResourceServer(oauth2 -> oauth2
                .jwt(jwt -> jwt
                    .jwtAuthenticationConverter(jwtAuthenticationConverter())
                )
            )
            .authorizeHttpRequests(authz -> authz
                // Token exchange requires API key
                .requestMatchers("/api/auth/exchange").hasAuthority("SCOPE_api-key")
                // All other API requests require JWT
                .requestMatchers("/api/**").authenticated()
                // Public endpoints
                .requestMatchers("/actuator/health").permitAll()
                .anyRequest().permitAll()
            );

        return http.build();
    }

    @Bean
    public JwtAuthenticationConverter jwtAuthenticationConverter() {
        JwtGrantedAuthoritiesConverter grantedAuthoritiesConverter =
            new JwtGrantedAuthoritiesConverter();
        grantedAuthoritiesConverter.setAuthorityPrefix("");

        JwtAuthenticationConverter jwtAuthenticationConverter =
            new JwtAuthenticationConverter();
        jwtAuthenticationConverter.setJwtGrantedAuthoritiesConverter(
            grantedAuthoritiesConverter
        );

        return jwtAuthenticationConverter;
    }

    @Bean
    public JwtDecoder jwtDecoder(@Value("${security.jwt.secret}") String secret) {
        SecretKey key = Keys.hmacShaKeyFor(secret.getBytes());
        return NimbusJwtDecoder.withSecretKey(key).build();
    }
}