/*
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package io.prestosql.server;

import com.google.common.hash.Hashing;
import io.airlift.http.client.HttpRequestFilter;
import io.airlift.http.client.Request;
import io.airlift.log.Logger;
import io.airlift.node.NodeInfo;
import io.jsonwebtoken.Claims;
import io.jsonwebtoken.Jws;
import io.jsonwebtoken.JwtException;
import io.jsonwebtoken.JwtParser;
import io.jsonwebtoken.Jwts;
import io.jsonwebtoken.SignatureAlgorithm;
import io.prestosql.server.security.InternalPrincipal;

import javax.inject.Inject;
import javax.servlet.http.HttpServletRequest;

import java.security.Principal;
import java.time.ZonedDateTime;
import java.util.Date;
import java.util.Optional;
import java.util.function.Supplier;

import static io.airlift.http.client.Request.Builder.fromRequest;
import static java.nio.charset.StandardCharsets.UTF_8;
import static java.util.Objects.requireNonNull;

public class InternalAuthenticationManager
        implements HttpRequestFilter
{
    private static final Logger log = Logger.get(InternalAuthenticationManager.class);

    private static final String PRESTO_INTERNAL_BEARER = "X-Presto-Internal-Bearer";

    private final Optional<JwtParser> jwtParser;
    private final Optional<Supplier<String>> jwtGenerator;

    @Inject
    public InternalAuthenticationManager(InternalCommunicationConfig internalCommunicationConfig, NodeInfo nodeInfo)
    {
        this(requireNonNull(internalCommunicationConfig, "internalCommunicationConfig is null").getSharedSecret(), nodeInfo.getNodeId());
    }

    public InternalAuthenticationManager(Optional<String> sharedSecret, String nodeId)
    {
        if (sharedSecret.isPresent()) {
            byte[] hmac = Hashing.sha256().hashString(sharedSecret.get(), UTF_8).asBytes();
            this.jwtParser = Optional.of(Jwts.parser().setSigningKey(hmac));
            this.jwtGenerator = Optional.of(() -> generateJwt(hmac, nodeId));
        }
        else {
            this.jwtParser = Optional.empty();
            this.jwtGenerator = Optional.empty();
        }
    }

    private static String generateJwt(byte[] hmac, String nodeId)
    {
        return Jwts.builder()
                .signWith(SignatureAlgorithm.HS256, hmac)
                .setSubject(nodeId)
                .setExpiration(Date.from(ZonedDateTime.now().plusMinutes(5).toInstant()))
                .compact();
    }

    public boolean isInternalRequest(HttpServletRequest request)
    {
        return request.getHeader(PRESTO_INTERNAL_BEARER) != null;
    }

    public Principal authenticateInternalRequest(HttpServletRequest request)
    {
        if (!jwtParser.isPresent()) {
            log.error("Internal authentication in not configured");
            return null;
        }

        String internalBarer = request.getHeader(PRESTO_INTERNAL_BEARER);
        try {
            Jws<Claims> claimsJws = jwtParser.get().parseClaimsJws(internalBarer);
            String subject = claimsJws.getBody().getSubject();
            return new InternalPrincipal(subject);
        }
        catch (JwtException e) {
            log.error(e, "Internal authentication failed");
            return null;
        }
        catch (RuntimeException e) {
            throw new RuntimeException("Authentication error", e);
        }
    }

    @Override
    public Request filterRequest(Request request)
    {
        return jwtGenerator.map(Supplier::get)
                .map(jwt -> fromRequest(request)
                        .addHeader(PRESTO_INTERNAL_BEARER, jwt)
                        .build())
                .orElse(request);
    }
}

@@ -802,113 +802,20 @@
 ken.
-Claims;%0Aimport io.jsonwebtoken.Jws;%0Aimport io.jsonwebtoken.JwtException;%0Aimport io.jsonwebtoken.JwtParser
+JwtException
 ;%0Aim
@@ -1125,24 +1125,60 @@
 l.Optional;%0A
+import java.util.function.Function;%0A
 import java.
@@ -1659,17 +1659,32 @@
 nal%3C
-JwtParser
+Function%3CString, String%3E
 %3E jw
@@ -2302,40 +2302,33 @@
 .of(
-Jwts.
+jwt -%3E 
 parse
-r().setSigningKey(hmac
+Jwt(hmac, jwt
 ));%0A
@@ -2857,32 +2857,266 @@
 mpact();%0A    %7D%0A%0A
+    private static String parseJwt(byte%5B%5D hmac, String jwt)%0A    %7B%0A        return Jwts.parser()%0A                .setSigningKey(hmac)%0A                .parseClaimsJws(jwt)%0A                .getBody()%0A                .getSubject();%0A    %7D%0A%0A
     public boole
@@ -3572,139 +3572,60 @@
     
-Jws%3CClaims%3E claimsJws = jwtParser.get().parseClaimsJws(internalBarer);%0A            String subject = claimsJws.getBody().getSubject(
+String subject = jwtParser.get().apply(internalBarer
 );%0A 

Fix thread safety in internal JWT