Jet Bridge example for JWT auth

jet_bridge_jwt_auth.java

import com.google.gson.Gson;
import com.google.gson.reflect.TypeToken;
import io.jsonwebtoken.Claims;
import io.jsonwebtoken.ExpiredJwtException;
import io.jsonwebtoken.Jws;
import io.jsonwebtoken.Jwts;

import java.io.BufferedReader;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.lang.reflect.Type;
import java.nio.charset.StandardCharsets;
import java.security.KeyFactory;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.security.PublicKey;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.X509EncodedKeySpec;
import java.util.Arrays;
import java.util.Base64;
import java.util.Map;
import java.util.zip.GZIPInputStream;

public class JetAuth {
    String JWT_VERIFY_KEY = "-----BEGIN PUBLIC KEY-----\nMIICIjANBgkqhkiG9w0BAQEFAAOCAg8AMIICCgKCAgEAyfJablscZmsh7eHisWHi\n/x2gJUjc9juL4VLkEpp6PH9Ah+wBtGmPF9nfoVlQsgH9hra63POiaLUzIwW/ewOl\njyPF0FJngsxreCxNs8lmd/WNXrcnazknFItkFFeXJuMTfoBPQGiZOmVYb14jmvkc\n9vMmeXYjrbT+95SVayN3E6DzLoHDhny4Mka1OsxvIP5s77s0dOo68TzoEfBVeuto\nI/dopG86DVu4wYVtYPITzJ4z47OFVPKCyYVyy5aR3+DUnmdK7xTRVr+iWmHpcr7e\nhoeVcL4CqAILZ0gd54kQmnHbg7Bu6x8JtQkiLU5TQvWzjiN00io4eydvIAkQTAaR\nmdd32O1vJbSHmLyCR2tEW/uV7P25naPUlkApxuLzh5C21S0XJxNJ/P07KSMymt5U\n1lWqt4CInpjAwMI8qs9MkEwJev5+yumxqIrDKcQLMR3TBLJZIb+rL1teCLOW28qB\nL6VSKhfKRIaXUdLpRwAcSuXraTzwa9oCCZa19tw3uizMeMFrCrv43YbyOsS9h7JQ\n8ixj/a1R/ud0fCrhXWUl7nKlz0b15koILLG1Ts+MUTmIaEnHTVEY74CfJVq7waw9\nx2kyzSzbsmMXvFkrVzTmyImTN631+gatU+npJ3vtcD9SooEZLOCLa4pb+DIsv9P1\nEeIEAh1VZC7s2qsQZsiYTG0CAwEAAQ==\n-----END PUBLIC KEY-----\n";

    public boolean check_token(String project_unique_name, String project_token, String user_token, String permission_type, String permission_object, String[] permission_actions) {
        String[] parts = user_token.split(" ", 2);
        String tokenType = parts[0];
        String tokenValue = parts[1];

        try {
            PublicKey key = this.get_verify_key();
            Jws<Claims> claims = Jwts.parserBuilder().setSigningKey(key).build().parseClaimsJws(tokenValue);

            var userId = (String)claims.getBody().get("user");
            var projects = (Map<String, Object>)claims.getBody().get("projects");
            var user_permissions = (Map<String, Object>)projects.get(project_unique_name);

            if (user_permissions == null) {
                return false;
            }

            return this.has_permissions(user_permissions, project_token, permission_type, permission_object, permission_actions);
        } catch (ExpiredJwtException e) {
            return false;
        } catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
            e.printStackTrace();
            return false;
        }
    }

    public String get_sha256_hash(String base) {
        try{
            MessageDigest digest = MessageDigest.getInstance("SHA-256");
            byte[] hash = digest.digest(base.getBytes("UTF-8"));
            StringBuffer hexString = new StringBuffer();

            for (byte b : hash) {
                String hex = Integer.toHexString(0xff & b);
                if (hex.length() == 1) hexString.append('0');
                hexString.append(hex);
            }

            return hexString.toString();
        } catch (Exception ex) {
            throw new RuntimeException(ex);
        }
    }

    private Map<String, Object>[] decompress_data(String value) {
        byte[] name = Base64.getEncoder().encode(value.getBytes());
        byte[] decodedString = Base64.getDecoder().decode(new String(name).getBytes(StandardCharsets.UTF_8));

        if ((decodedString == null) || (decodedString.length == 0)) {
            return null;
        }

        final StringBuilder outStr = new StringBuilder();

        try (GZIPInputStream gis = new GZIPInputStream(new ByteArrayInputStream(decodedString))) {
            final BufferedReader bufferedReader = new BufferedReader(new InputStreamReader(gis, StandardCharsets.UTF_8));
            String line;

            while ((line = bufferedReader.readLine()) != null) {
                outStr.append(line);
            }

            String str = outStr.toString();

            Type mapType = new TypeToken<Map<String, Object>[]>(){}.getType();
            return new Gson().fromJson(str, mapType);
        } catch (IOException e) {
            e.printStackTrace();
            return null;
        }
    }

    private boolean has_permissions(Map<String, Object> user_permissions, String project_token, String permission_type, String permission_object, String[] permission_actions) {
        if ((boolean)user_permissions.get("owner")) {
            return true;
        } else if ((boolean)user_permissions.get("super_group")) {
            return true;
        }

        Map<String, Object>[] permissions;

        if (user_permissions.containsKey("permissions")) {
            permissions = this.decompress_data((String)user_permissions.get("permissions"));
        } else {
            permissions = new Map[]{};
        }

        if (permissions == null) {
            return false;
        }

        if ((boolean)user_permissions.get("read_only")) {
            if (permission_type.equals("model") && Arrays.stream(permission_actions).allMatch(item -> item.equals("r"))) {
                return true;
            }
        } else {
            return false;
        }

        String token_hash = get_sha256_hash(project_token.replace("-", "").toLowerCase());

        for (var item : permissions) {
            var item_type = (String)item.getOrDefault("permission_type", "");
            var item_object = (String)item.getOrDefault("permission_object", "");
            var item_actions = (String)item.getOrDefault("permission_actions", "");

            if (permission_type.equals("model")) {
                var resource_token_hash = item.getOrDefault("resource_token_hash", "");
                var item_object_model = item_object.split("\\.", 2)[1];

                if (resource_token_hash != "" && resource_token_hash != token_hash) {
                    continue;
                }

                if (!item_type.equals(permission_type) || !item_object_model.equals(permission_object)) {
                    continue;
                }
            } else {
                if (!item_type.equals(permission_type) || !item_object.equals(permission_object)) {
                    continue;
                }
            }

            return Arrays.stream(permission_actions).allMatch(action -> item_actions.contains(action));
        }

        return false;
    }

    private PublicKey get_verify_key() throws NoSuchAlgorithmException, InvalidKeySpecException {
        KeyFactory keyFactory = KeyFactory.getInstance("RSA");
        String keyString = JWT_VERIFY_KEY
                .replaceAll("\n", "")
                .replace("-----BEGIN PUBLIC KEY-----", "")
                .replace("-----END PUBLIC KEY-----", "");

        byte[] byteKey = Base64.getDecoder().decode(keyString);
        X509EncodedKeySpec keySpec = new X509EncodedKeySpec(byteKey);
        return keyFactory.generatePublic(keySpec);
    }
}