KeyCloak Security

KeyCloak/OAuth

The nature of OAuth in web usage includes a frontend client like a web browser and a backend client which is the application (for example, Apache Polaris). 

Realms are isolated namespaces of security. They can be federated but are not by default.

The JWT (JSON Web Token) is associated with a realm that is passed to the backend client. It contains an ISSued field that indicates who issued it. The KeyCloak issuing it signs it with its private key.

The ISS is a URL that the embeds both the realm and the domain of the ISS. The backend client can use this URL to find a public key of KeyCloak with which it can verify the provenance of the JWT. 

This endpoint is called the Realm JWKS (JSON Web Key Set).

If a claim in the JWT is azp (authorized party), it identifies to whom the token was issued. If it's audience, it identifies for whom the token was issued.

So, an example would be that the user logs in and their web browser is the azp. They ask KeyCloak for a token on behalf of the back end, the aud. KeyCloak signs both so the backend knows the token is genuinely for it and who requested it. If the azp is not in its whitelist, the backend can reject it. 

KeyCloak Configuration

To relax this constraint (say, for integration tests), Polaris can set the environment variable quarkus.oidc.token.issuer=any.

Interestingly, OAuth2 servers have a handy endpoint at .well-known/openid-configuration listing salient details. So, for KeyCloak, run:

curl -s https://YOUR_DOMAIN/realms/YOUR_REALM/.well-known/openid-configuration | jq

If your backend is in the same private network as KeyCloak, you can set KC_HOSTNAME_STRICT_BACKCHANNEL=false to make KeyCloak dynamically generate backchannel URLs based on the request's Host header - the mandatory HTTP header that defines which virtual host is the endpoint.

Note that an attacker could use Host Header Injection to trick KeyCloak into certifying a malicious domain. To stop this, the load balancer must be correctly configured to ignore inappropriate Host values plus KeyCloak should have KC_HOSTNAME_STRICT=true and KC_HOSTNAME set so even if the hacker gets past the load balancer, KeyCloak says "nice try but I know what public facing name I have". Again, you might like to relax this for testing.

Talking of the load balancer, traffic between it and KeyCloak is unencrypted so we need to tell KeyCloak communication is plain HTTP after TLS termination with KC_HTTP_ENABLED=true and set a KC_PROXY_HEADERS strategy to tell it what the load balancer is doing with the traffic. We use xforwarded so it knows our AWS ALB injects headers in a standard way.

LLMs in Java - part 1

Interpretability of LLMs is important in many real world situations. For instance, there is an EU law that says if a bank's model refuses a loan then the client has the right to know why. But neural nets are notoriously hard to interpret.

There are a few techniques:

Activation Patching: this is basically A/B testing a neural net where the activations from a clean prompt are transplanted to that of a corrupted prompt and we check if that fixes the corruption.

Logit lens: take the residual stream vector (that is, the vector that represents the working memory) and apply it to the vocabulary logits at each layer. This way, we can see where an output word (probably) starts to form.

Sparse Auto Encoders: the auto encoder is trained on the model's input and output, and deliberately made sparse for low values so prominent features become clear.

The neural net itself

Most of the Java code that iterates over the tokens lives in InferenceEngine.generateTokensGPULlama. In turn, this delegates to the Tornado graphs that are executed in TornadoVMMasterPlan.tornadoVMForwardExecuteLayered. Note that the positionHolder is an IntArray that has a single element.

However, the real heavy lifting is done in LlamaFP16FFNLayers.setupSingleFFNLayer and all the functions in TransformerComputeKernelsLayered it uses. Note, that it doesn't call them. Instead, it builds them up into a graph of lazily evaluated functions that are executed in the code mentioned above.

Note that values whose name ends with Cache are a distillation of all that has come before.

The functions that are called are (in this order):

reductionOneBlockWithLayer

reductionFinalNormalization

fusedQKVMatmulX

ropeRotationWithCacheCopy

processHeadsFlashAttention

processHeadsParallel

matrixVectorGenericWithResidual

reductionOneBlockWithLayer

reductionFinalNormalization

fusedRmsNormFFNGateUp

matrixVectorGenericWithResidual

The italicised functions configure the attention.

Let's go through them:

reductionOneBlockWithLayer calculates the root mean square of the input and takes its inverse. This is calculated per work group.

reductionFinalNormalization reduces these calculations of all the workgroups to a single value.

 

fusedQKVMatmulX multiplies the vector x by the vectors wq, wk, and wv, putting the results in q, k, and v respectively. 

ropeRotationWithCacheCopy here we see a standard, two-dimensional rotation on pairs of elements in q and k. The angle we rotate by is a function of both the word's position and which attention head this is for (which is just the global thread ID mod the head size). Note we're imposing a geometry on vectors q and k that doesn't exist before this point. 

processHeadsFlashAttention this calculates softmax(Q K^T) V. It uses tiling [previous post] as this is a massive matrix operation.

This is done for all the FFN layers plus similar calls in LogitsFP16Layer.setupLogitsTaskGraph for the logit layer. Here, the functions are (in this order):

reductionOneBlockWithLayer

reductionFinalNormalization

mapContextWithQuantizeLogits

matrixVectorGeneric

Note that these are all called for each of the layers. There appear to be 16 for Llama.

There are more layers but this is enough for my first post on the subject. More in Part 2.

Rope

Upon loading the model, RoPE.precomputeFreqsCis is called when just before the weights are loaded.

We do a 2-d rotation on pair elements of the vector, imposing a geometry that was not there before. The choice of a 2-d rotation is arbitrary but it's computationally the cheapest dimension to rotate as the number of elements in a rotation matrix is d² (where d is the number of dimensions we're working in).

Production ready Polaris

Persisting the metadata

First, create a DB. I created a Postgres RDS database in AWS and allowed it to create the VPC, subnets etc. It took me a while to work out why I could connect from my laptop but not a Polaris running in Azure: the source in the security group AWS automatically generated allowed my IP address but not Microsoft's.

You can check what IP address the greater internet sees you as with:

curl -s https://checkip.amazonaws.com

Bootstrap Polaris with:

docker run --rm -it   --env="polaris.persistence.type=relational-jdbc"   --env="quarkus.datasource.username=$DB_USERNAME"   --env="quarkus.datasource.password=$DB_PASSWORD"   --env="quarkus.datasource.jdbc.url=jdbc:postgresql://$DB_HOST:5432/polaris_db"   apache/polaris-admin-tool:latest bootstrap -r POLARIS -c POLARIS,root,s3cr3t

(You can purge the database by using the above but with arguments purge -r POLARIS)

Then you can see in Postgres:

postgres=> \c polaris_db

polaris_db=> SELECT * FROM pg_catalog.pg_tables;

     schemaname     |           tablename           | tableowner | tablespace | hasindexes | hasrules | hastriggers | rowsecurity 

--------------------+-------------------------------+------------+------------+------------+----------+-------------+-------------

 polaris_schema     | version                       | postgres   |            | t          | f        | f           | f

 polaris_schema     | entities                      | postgres   |            | t          | f        | f           | f

 polaris_schema     | grant_records                 | postgres   |            | t          | f        | f           | f

 polaris_schema     | principal_authentication_data | postgres   |            | t          | f        | f           | f

Create the database with:

CREATE DATABASE polaris_db;

CREATE USER polaris_user WITH PASSWORD 'your_secure_password';

GRANT ALL PRIVILEGES ON DATABASE polaris_db TO polaris_user;

\c polaris_db

GRANT ALL ON SCHEMA public TO polaris_user;

If you mess up your Polaris, just run:

kubectl rollout restart deployment polaris-deployment

as now the data is all in the database.

Access Control

For integration tests, you can just use the client_id and client_secret with which you set up Polaris. But if you put it in production, you'll want to create users (Principals).

"At the most basic level, Polaris' persistence layer stores Entities and Grants, where Grants define the access-control-related relationship between entities." [Apache Polaris Catalog Federation Proposal]

To access Polaris, you need a Principal. This will have Principal Roles. They need to be associated with the Catalog Roles that in turn belong to a Catalog.

REST via curl

Effective debugging of Polaris can be done by poking its REST API. There is a command line tool to do this more easily but it's a bit buggy. The latest code for me had trouble deserializing SigV4 objects.

First, you need a token

POLARIS_TOKEN=$(curl -X POST "https://$HOST/api/catalog/v1/oauth/tokens"   -H "Content-Type: application/x-www-form-urlencoded"   -d "grant_type=client_credentials&client_id=$CLIENT_ID&client_secret=$CLIENT_SECRET&scope=PRINCIPAL_ROLE:ALL" | jq -r '.access_token')

View namespaces

curl -X GET "https://$HOST/api/catalog/v1/$CATALOG_NAME/namespaces/$NAMESPACE" -H "Authorization: Bearer ${POLARIS_TOKEN}" | jq

{

  "namespace": [

    "samples"

  ],

  "properties": {

    "owner": "henryp",

    "location": "s3a://emrys-afon-bucket/samples/"

  }

}

View tables

curl -X GET "https://$HOST/api/catalog/v1/$CATALOG_NAME/namespaces/$NAMESPACE/tables"   -H "Authorization: Bearer ${POLARIS_TOKEN}"   -H "Accept: application/json"   -s | jq .

or given a table:

curl -X GET "https://$HOST/api/catalog/v1/$CATALOG_NAME/namespaces/$NAMESPACE/tables/$TABLE" -H "Authorization: Bearer ${POLARIS_TOKEN}" | jq

Clean up

Remove lingering details in the namespace with something like:

curl -X POST -H "Content-Type: application/json" -H "Authorization: Bearer ${POLARIS_TOKEN}"  "https://$HOST/api/catalog/v1/aws/namespaces/samples/properties" --data '{"removals": ["owner"] }' | jq

Delete Catalog

curl -X DELETE "https://$HOST/api/management/v1/catalogs/$CATALOG_NAME"   -H "Authorization: Bearer ${POLARIS_TOKEN}"   -H "Content-Type: application/json" -o /dev/null -s -w "%{http_code}\n"

R-squared

R² is a good way to test how well your model describes the data. It is literally

1 - (variance you predict / total variance)

It is therefore model agnostic. (For more about the maths of variance, see a previous post).

A value of 1.0 means the model describes the data perfectly, 0.0 means it is the same as guessing the mean and less than zero means it's worse than useless.

So, here's an interesting question: if my prediction for the next value in a series is the last known value, what is R²? In theory, it should be 1 and there is a mathematical argument for this. 

Some general points.

1. The formula for covariance is Σ(X_t-X̅)(Y_t-Ȳ)/N . 

Expand it and you'll see it's the same as E[XY] - E[X]E[Y].

2. We assume that a value is the sum of previous shocks mutliplied by a factor that geometrically reduces with each step. We'll call this φ, the autoregressive coefficient.

So:

y_t = ε_t + φε_t-1 + φ²ε_t-2 + ... where |φ|<1

And with a self-substitution, this becomes:

y_t = ε_t + φ y_t-1

We generally add a constant to this formula but for mathematical simplicity the rest of the argument assumes we've centred the variable on its mean - ie E[Y]=0.

3. The covariance of two unrelated distributions is 0

1. note that the joint probability distribution of x and y if they're independent is f(x,y) = f(x)f(y)

2. integrate the expected values, that is integrate xy.f(x,y) = x f(x) y f(y)

3. When you do that, you'll find E[XY] = E[X]E[Y].

4. Substituing this into the formula for covariance above, Cov(X,Y) = 0 if X and Y are independent.

4. The covariance of two distributions that are the same is the variance.

Substitute X=Y into the formula for covariance and you get Cov(X,Y)=Cov(X,X)=E[X²]-E[X]²=Var(X) as you can see in a previous blog post.

5. Take that formula in step 2, multiply both sides by y_t-1 and take the covariance of all the terms. 

Cov(ε_ty_t-1) is zero because of step 3.

Cov((y_t-1)2)=Var(Y)= E[Y²]-E[Y]² = E[Y²] because E[Y]=0 since we centred Y on its mean.

So, rearranging φ=Cov(y_t,y_t-1)/Var(y_t-1)

For a random walk:

y_t = ε_t + y_t-1

Necessarily φ=1 because the most recent step is the last plus some randomness.

So, the variance in our model is Var(ε). This is because the above formula can equivalently be written as:

y_t = y₀ + ε₁ + ε₂ + ε₃ +...

and choosing y₀=0 then:

Var(y_t) = Var(ε₁) + Var(ε₂) + Var(ε₃)+...

Var(y_t) = N Var(ε)

The total variance changes over time. Centering our walk at zero, y_t = N ε_t after N steps. So, here we find the variance is N Var(ε).

Putting these two values into the equation at the top of this post,

R² = 1 - Var(ε) / N Var(ε) = 1 - (1/N)

Naturally, as N tends to infinity, R² tends to 1.

So, a model that predicts the next value based on the last is a good one because R² is 1, right? No! This is a spurious regression. They may be useful for diagnostics (maybe) but not for models.

Linear Regression tip 

I built a decent sales model with:

MAE      :         1439240.3746               6.3954 % of mean

MAPE     :              17.0928 %

R²       :               0.9646

If you plot the residuals against the target values, you should see a random scatter plot around y=0.

This was not the case for me on my first try. There was a pattern. The residuals tended to be below zero for small target values and above zero for large. Basically, the plots were saying I was overestimating low values and underestimating large values. 

Since this was a linear regression, adding some non-linearity really helped. I both squared and loged all my regressors:

MAE      :         1293045.8398               5.1369 % of mean

MAPE     :              10.7136 %

R²       :               0.9717

This was great but the improvement didn't appear to actually come from the non-linearity. Instead, it was this line removing lots of rows:

        df = df.replace([np.inf, -np.inf], np.nan).dropna()

We need this because linear regression is intolerant to anything but a real number. Why removing the lines that had negative values (and consequently infinite loged regressors) remains a mystery. 

But it also caused another problem. As more columns were added, the probability of a row being excluded because it had a nan increased. Therefore, this little line of Python in our data pipeline had the unintended consequence of reducing the number of rows as the number of columns increased!

Polaris and Cloud tokens

Polaris rather pleasingly mints cloud tokens that are subscoped to a directory in a bucket or blob container for AWS and GCP. That is, even if the token has been hijacked, the blast radius is limited by:

• the token only allowing access to a single directory and its subfolders, not the whole bucket
• the token is no good after X minutes (where the default valu of X is 60)

There's currently an outstanding ticket to give subscoping to Azure.

The code for vending for the different clouds belongs in the implementations of PolarisStorageIntegration.getSubscopedCreds and this is where the tokens are created. You could put breakpoints in the breakpoints of:

com.google.auth.oauth2.AccessToken

com.azure.core.credential.AccessToken

software.amazon.awssdk.auth.credentials.AwsSessionCredentials 

and grab the credentials and use them on the command line (that is, entirely outside of Polaris) thus:

# AWS

AWS_ACCESS_KEY_ID=... AWS_SECRET_ACCESS_KEY=... AWS_SESSION_TOKEN=...  aws s3 ls s3://YOUR_BUCKET/DIRECTORY_FOR_TOKEN

# Azure

az storage blob list  --account-name $STORAGE_ACCOUNT   --container-name $CONTAINER --sas-token $SAS_TOKEN --prefix YOUR_DIRECTORY

# GCP

CLOUDSDK_AUTH_ACCESS_TOKEN=ya29... gcloud storage ls gs://YOUR_BUCKET/DIRECTORY_FOR_TOKEN

But even if you did, the tokens no longer work after 60 minutes and in the case of AWS and GCP, you cannot even view directories for which the token was not defined.

The state of Polaris

Huzzah! My PR for Apache Polaris has been accepted and merged with its main branch! Here are some miscellaneous notes I made as I looked at what needed to be done and how to test my code.

Federated catalogs

How Polaris handles vended credentials in federated catalogs is still an ongoing concern [Polaris mailing list]. The issue concerns who has say over what is vended. If the external catalog does not allows user X but the Polaris instance that defers to it does, is user X allowed to use that data or not?

In the ticket "Does Polaris support credential vending for external REST Catalogs?", Polaris maintainer, Alex Dutra, says:

"When the client requests credential vending, Polaris forwards the request to the remote catalog, but mints temporary credentials itself and vends them to the client. IOW, a PolarisStorageConfigurationInfo must have been configured when declaring the external catalog in Polaris, and it's this storage config that will be used for vending credentials."

Integration tests 

My GcpCatalogFederationIntegrationIT lives in runtime/service/src/cloudTest/ and unlike their counterparts in runtime/service/src/, they need to run against an already started Polaris instance (the latter start their own and run out of the box).

Note, you'll have to set:

polaris.features."ENABLE_CATALOG_FEDERATION"=true

polaris.features."ALLOW_OVERLAPPING_CATALOG_URLS"=true

in runtime/defaults/src/main/resources/application.properties and

"-Dpolaris.bootstrap.credentials=POLARIS,test-admin,test-secret",

in runtime/server/build.gradle.kts if you're running your Polaris from the source code because ServerManager hard codes the ClientCredentials.

Note that your Polaris will need to Google credentials for my code - for example:

export GOOGLE_APPLICATION_CREDENTIALS=/home/henryp/gcp.json 

Now you can run Polaris with:

./gradlew --stop && ./gradlew run

Cloud Topologies and Azure

All clouds follow a similar pattern when it comes to networking. All of them have:

• a VPN (virtual private network) that is an entirely isolated environment
• subnets that are divided into public and private IP spaces by using non-overlapping CIDR blocks
• an internet gateway that allows access to the outside world and performs the NAT
• Route tables that permit traffic flows by typically bringing together subnets and security groups. If subnets are rooms in a building, route tables are the corridors between them and the Network ACLs (NACLs) are the bouncers on the door. Security groups are more identity based and they do a similar job to NACLs albeit at the vNIC level.
• a load balancer that allows incoming traffic from the internet. This differs from the load balancer as it's more the receptionist directing people rather than the security guard on the front door.
• the computers/VMs on which Kubernetes runs, called nodes.

The terminology may slightly differ but this is true for all the main cloud providers

Azure

Now, bearing in mind that this is the general concept for all cloud/K8s clusters, the recipe in Azure is:

1. Create the Resource Group
2. Create an identity for this group.
3. Create a network for this group and identity.

Now, when it comes to the network, the steps are:

1. Create the Virtual Network compatible with the resource group
2. Create the subnet for the Virtual Network
3. Create a Network Security Group for the resource group. This defines the inbound and outbound rules. Funnily enough, the rules in the NSG are of higher priority the lower the number.
4. We associate the subnet with the Network Security Group
5. We assign a role to the subnet.

Now we can set up the Kubernetes cluster if we:

1. Assign the subnet to a K8s Node Pool
2. Create a network profile

This last one is interesting. It's the "the reference to the network interface card" [Azure for Architects, Packt]. The network car is the actual physical hardware the packets must use but generally they're only referenced in Azure (and GCP) as virtual network interface cards (vNIC). In AWS, they're called ENIs.