OAuth2 vs JWT.

package com.example.demo;

import com.fasterxml.jackson.databind.ObjectMapper;
import jakarta.servlet.FilterChain;
import jakarta.servlet.ServletException;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import org.springframework.boot.SpringApplication;
import org.springframework.boot.autoconfigure.SpringBootApplication;
import org.springframework.context.annotation.Bean;
import org.springframework.http.HttpStatus;
import org.springframework.http.ResponseEntity;
import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
import org.springframework.security.config.annotation.web.builders.HttpSecurity;
import org.springframework.security.config.http.SessionCreationPolicy;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.security.web.SecurityFilterChain;
import org.springframework.security.web.authentication.UsernamePasswordAuthenticationFilter;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.filter.OncePerRequestFilter;

import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.time.Instant;
import java.util.Collections;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.Optional;

@SpringBootApplication
public class OAuth2VsJwtDemoApplication {

    public static void main(String[] args) {
        SpringApplication.run(OAuth2VsJwtDemoApplication.class, args);
    }

    @Bean
    SecurityFilterChain securityFilterChain(HttpSecurity http, JwtAuthenticationFilter jwtAuthenticationFilter) throws Exception {
        return http
                // Stateless APIs usually disable CSRF because bearer tokens are sent in the Authorization header,
                // not as browser cookies.
                .csrf(csrf -> csrf.disable())
                .sessionManagement(session -> session.sessionCreationPolicy(SessionCreationPolicy.STATELESS))
                .authorizeHttpRequests(auth -> auth
                        .requestMatchers("/auth/login", "/api/public").permitAll()
                        .anyRequest().authenticated())
                .addFilterBefore(jwtAuthenticationFilter, UsernamePasswordAuthenticationFilter.class)
                .build();
    }

    @Bean
    JwtService jwtService(ObjectMapper objectMapper) {
        return new JwtService(objectMapper, "change-me-in-prod-change-me-in-prod");
    }

    @Bean
    JwtAuthenticationFilter jwtAuthenticationFilter(JwtService jwtService) {
        return new JwtAuthenticationFilter(jwtService);
    }
}

record LoginRequest(String username, String password) {}

@RestController
class AuthController {
    private final JwtService jwtService;
    private final Map<String, String> users = Map.of(
            "alice", "password123",
            "bob", "s3cret"
    );

    AuthController(JwtService jwtService) {
        this.jwtService = jwtService;
    }

    @PostMapping("/auth/login")
    ResponseEntity<Map<String, Object>> login(@RequestBody LoginRequest request) {
        String expected = users.get(request.username());
        if (expected == null || !expected.equals(request.password())) {
            return ResponseEntity.status(HttpStatus.UNAUTHORIZED)
                    .body(Map.of("error", "bad_credentials"));
        }

        String token = jwtService.issueToken(request.username());
        return ResponseEntity.ok(Map.of(
                "access_token", token,
                "token_type", "Bearer",
                "expires_in_seconds", 900
        ));
    }
}

@RestController
class ApiController {
    @GetMapping("/api/public")
    Map<String, String> publicEndpoint() {
        return Map.of("message", "public endpoint");
    }

    @GetMapping("/api/me")
    Map<String, Object> me(Authentication authentication) {
        return Map.of(
                "user", authentication.getName(),
                "message", "JWT accepted"
        );
    }
}

class JwtAuthenticationFilter extends OncePerRequestFilter {
    private final JwtService jwtService;

    JwtAuthenticationFilter(JwtService jwtService) {
        this.jwtService = jwtService;
    }

    @Override
    protected void doFilterInternal(HttpServletRequest request,
                                    HttpServletResponse response,
                                    FilterChain filterChain) throws ServletException, IOException {
        String header = request.getHeader("Authorization");

        // If no token is present, let Spring Security handle the anonymous/public path.
        if (header != null && header.startsWith("Bearer ")) {
            String token = header.substring(7).trim();
            Optional<String> subject = jwtService.validateAndGetSubject(token);

            // This is the failure path interviewers care about: bad signature, bad format, or expired token.
            if (subject.isEmpty()) {
                response.setStatus(HttpStatus.UNAUTHORIZED.value());
                response.setContentType("application/json");
                response.getWriter().write("{\"error\":\"invalid_or_expired_token\"}");
                return;
            }

            UsernamePasswordAuthenticationToken authentication =
                    new UsernamePasswordAuthenticationToken(subject.get(), null, Collections.emptyList());
            SecurityContextHolder.getContext().setAuthentication(authentication);
        }

        filterChain.doFilter(request, response);
    }
}

class JwtService {
    private static final long TTL_SECONDS = 900;
    private final ObjectMapper objectMapper;
    private final SecretKeySpec keySpec;

    JwtService(ObjectMapper objectMapper, String secret) {
        this.objectMapper = objectMapper;
        // HMAC keeps this demo self-contained. Real systems often use RSA or EC keys so resource servers
        // can verify with a public key and the auth server keeps the private key.
        this.keySpec = new SecretKeySpec(secret.getBytes(StandardCharsets.UTF_8), "HmacSHA256");
    }

    String issueToken(String subject) {
        long now = Instant.now().getEpochSecond();

        Map<String, Object> header = Map.of(
                "alg", "HS256",
                "typ", "JWT"
        );

        Map<String, Object> payload = new LinkedHashMap<>();
        payload.put("sub", subject);
        payload.put("iat", now);
        payload.put("exp", now + TTL_SECONDS);
        payload.put("iss", "spring-boot-demo");

        String headerPart = base64Url(json(header));
        String payloadPart = base64Url(json(payload));
        String signaturePart = sign(headerPart + "." + payloadPart);

        return headerPart + "." + payloadPart + "." + signaturePart;
    }

    Optional<String> validateAndGetSubject(String token) {
        try {
            int firstDot = token.indexOf('.');
            int secondDot = token.indexOf('.', firstDot + 1);
            if (firstDot <= 0 || secondDot <= firstDot + 1 || token.indexOf('.', secondDot + 1) != -1) {
                return Optional.empty();
            }

            String headerPart = token.substring(0, firstDot);
            String payloadPart = token.substring(firstDot + 1, secondDot);
            String signaturePart = token.substring(secondDot + 1);

            String expectedSignature = sign(headerPart + "." + payloadPart);
            if (!MessageDigest.isEqual(base64UrlDecode(expectedSignature), base64UrlDecode(signaturePart))) {
                return Optional.empty();
            }

            Map<String, Object> header = readMap(headerPart);
            if (!"HS256".equals(header.get("alg")) || !"JWT".equals(header.get("typ"))) {
                return Optional.empty();
            }

            Map<String, Object> payload = readMap(payloadPart);
            Object expValue = payload.get("exp");
            Object subValue = payload.get("sub");
            if (!(expValue instanceof Number) || !(subValue instanceof String)) {
                return Optional.empty();
            }

            long exp = ((Number) expValue).longValue();
            if (Instant.now().getEpochSecond() >= exp) {
                return Optional.empty();
            }

            return Optional.of((String) subValue);
        } catch (Exception ex) {
            return Optional.empty();
        }
    }

    private byte[] json(Map<String, Object> data) {
        try {
            return objectMapper.writeValueAsBytes(data);
        } catch (IOException ex) {
            throw new IllegalStateException("Could not serialize JWT data", ex);
        }
    }

    @SuppressWarnings("unchecked")
    private Map<String, Object> readMap(String base64UrlPart) throws IOException {
        return objectMapper.readValue(base64UrlDecode(base64UrlPart), Map.class);
    }

    private String sign(String signingInput) {
        try {
            Mac mac = Mac.getInstance("HmacSHA256");
            mac.init(keySpec);
            byte[] signed = mac.doFinal(signingInput.getBytes(StandardCharsets.UTF_8));
            return base64Url(signed);
        } catch (Exception ex) {
            throw new IllegalStateException("Could not sign JWT", ex);
        }
    }

    private String base64Url(byte[] bytes) {
        return java.util.Base64.getUrlEncoder().withoutPadding().encodeToString(bytes);
    }

    private byte[] base64UrlDecode(String value) {
        return java.util.Base64.getUrlDecoder().decode(value);
    }
}