feat: authenticate agent call (#296)

## Description

Calling the agent API now requires a token.

Closes #292

.env.example

@@ -32,3 +32,4 @@ LEADERBOARD_PROGRAM_NAME="leaderboard.aleo"
 AWS_ACCESS_KEY="<your_access_key>"
 AWS_SECRET="<your_secret>"
 AWS_INFERENCE_URL="https://nez5g29b86.execute-api.eu-central-1.amazonaws.com/staging/generate"
+LAMBDA_TOKEN="<your_token>"

.github/workflows/deploy-staging-gke.yaml

@@ -48,6 +48,7 @@ jobs:
           OBSCURA_API_KEY: ${{ secrets.OBSCURA_API_KEY }}
           AWS_ACCESS_KEY: ${{ secrets.AWS_ACCESS_KEY }}
           AWS_SECRET: ${{ secrets.AWS_SECRET }}
+          LAMBDA_TOKEN: ${{ secrets.LAMBDA_TOKEN }}
         run: |-
           envsubst < deployment/staging/secrets.template.yaml > deployment/staging/secrets.yaml
           skaffold run --filename skaffold.staging.yml --cache-artifacts=false

deployment/staging/deployment.yaml

@@ -61,5 +61,10 @@ spec:
                 secretKeyRef:
                   name: backend-secrets
                   key: aws_secret
+            - name: LAMBDA_TOKEN
+              valueFrom:
+                secretKeyRef:
+                  name: backend-secrets
+                  key: lambda_token
             - name: AWS_INFERENCE_URL
               value: https://nez5g29b86.execute-api.eu-central-1.amazonaws.com/staging/generate

deployment/staging/secrets.template.yaml

@@ -10,3 +10,4 @@ data:
   obscura_api_key: $OBSCURA_API_KEY
   aws_access_key: $AWS_ACCESS_KEY
   aws_secret: $AWS_SECRET
+  lambda_token: $LAMBDA_TOKEN

infra/auth-lambda.py

@@ -0,0 +1,45 @@
+# A simple token-based authorizer example to demonstrate how to use an authorization token
+# to allow or deny a request. In this example, the caller named 'user' is allowed to invoke
+# a request if the client-supplied token value is 'allow'. The caller is not allowed to invoke
+# the request if the token value is 'deny'. If the token value is 'unauthorized' or an empty
+# string, the authorizer function returns an HTTP 401 status code. For any other token value,
+# the authorizer returns an HTTP 500 status code.
+# Note that token values are case-sensitive.
+
+import json
+import os
+
+
+TOKEN = os.environ["TOKEN"]
+
+
+def lambda_handler(event, context):
+    token = event["authorizationToken"]
+    if token == TOKEN:
+        print("authorized")
+        response = generatePolicy("user", "Allow", event["methodArn"])
+    else:
+        print("unauthorized")
+        raise Exception("Unauthorized")  # Return a 401 Unauthorized response
+    try:
+        return json.loads(response)
+    except BaseException:
+        print("unauthorized")
+        return "unauthorized"  # Return a 500 error
+
+
+def generatePolicy(principalId, effect, resource):
+    authResponse = {}
+    authResponse["principalId"] = principalId
+    if effect and resource:
+        policyDocument = {}
+        policyDocument["Version"] = "2012-10-17"
+        policyDocument["Statement"] = []
+        statementOne = {}
+        statementOne["Action"] = "execute-api:Invoke"
+        statementOne["Effect"] = effect
+        statementOne["Resource"] = resource
+        policyDocument["Statement"] = [statementOne]
+        authResponse["policyDocument"] = policyDocument
+    authResponse_JSON = json.dumps(authResponse)
+    return authResponse_JSON

src/env.mjs

@@ -30,6 +30,7 @@ export const env = createEnv({
     AWS_ACCESS_KEY: z.string(),
     AWS_SECRET: z.string(),
     AWS_INFERENCE_URL: z.string().url(),
+    LAMBDA_TOKEN: z.string(),
   },
 
   /**
@@ -71,6 +72,7 @@ export const env = createEnv({
     AWS_ACCESS_KEY: process.env.AWS_ACCESS_KEY,
     AWS_SECRET: process.env.AWS_SECRET,
     AWS_INFERENCE_URL: process.env.AWS_INFERENCE_URL,
+    LAMBDA_TOKEN: process.env.LAMBDA_TOKEN,
   },
   /**
    * Run `build` or `dev` with `SKIP_ENV_VALIDATION` to skip env validation.

src/server/service/agent.ts

@@ -18,7 +18,7 @@ export class Agent {
   private _systemPrompt: string;
 
   private _triggeredAt = Date.now();
-  private _silenceToken = "001001";
+  private readonly _silenceToken = "001001";
 
   get id() {
     return this._id;
@@ -103,24 +103,29 @@ export class Agent {
       parameters: { max_new_tokens: 58, top_p: 1, temperature: 0.8 }, // TODO define final parameters as constants
     });
 
-    const response = await fetch(env.AWS_INFERENCE_URL, {
-      headers: {
-        "Content-Type": "application/json",
-        // Authorization: `Bearer ${env.AWS_TOKEN}`,
-      },
-      method: "POST",
-      body,
-      signal: AbortSignal.timeout(10000),
-    });
-
-    const textRes = await response.text();
-
-    if (!textRes) return this._silenceToken;
-
-    const result = JSON.parse(textRes) as { body: string };
-    const responseBody = JSON.parse(result.body) as string;
-
-    return responseBody;
+    try {
+      const response = await fetch(env.AWS_INFERENCE_URL, {
+        headers: {
+          "Content-Type": "application/json",
+          authorizationToken: env.LAMBDA_TOKEN,
+        },
+        method: "POST",
+        body,
+        signal: AbortSignal.timeout(10000),
+      });
+
+      const textRes = await response.text();
+
+      if (!textRes) return this._silenceToken;
+
+      const result = JSON.parse(textRes) as { body: string };
+      const responseBody = JSON.parse(result.body) as string;
+
+      return responseBody;
+    } catch (error) {
+      console.error(error);
+      return this._silenceToken;
+    }
   }
 
   cleanup() {