2.3 Flows Security Recipe - Securing Genkit Flows - Authentication & Authorization

In the previous recipes, we explored how to use Genkit flows to create multi-step interactions instead of normal functions. However, in real-world applications, security is a paramount concern, especially when dealing with sensitive data or operations.

It’s crucial to implement robust security measures to protect your Genkit flows from unauthorized access and potential threats. In this recipe, we will demonstrate how to secure your Genkit flows effectively.

In this recipe, we will see how to secure a Genkit flow using authentication and authorization mechanisms.

1. Setting Up Authentication

Genkit provides a built-in context provider system that allows you to manage authentication and authorization seamlessly. There are two main approaches to implement authentication in Genkit flows:

• API Key Authentication: This method involves using API keys to authenticate requests. You can then pass the API key in the request headers - Authorization.
• Custom Authentication Providers: You can create custom authentication providers by implementing the ContextProvider interface. This allows you to define your own authentication logic, such as OAuth, JWT, or any other method that suits your application’s needs. This approach provides flexibility and allows you to integrate with existing authentication systems.

API Key Authentication Example

When using API key authentication, you can either use a static API key or implement a dynamic key generation mechanism depending on your security requirements. Let’s take the following flow, which requires an API key for access:

const protectedFlowByAPIKey = ai.defineFlow(
  {
    name: 'protectedFlowByAPIKey',
    inputSchema: z.object({}),
    outputSchema: z.string(),
  },
  async (input, { context }) => {
    const apiKey = context?.auth?.apiKey;

    return apiKey;
  },
);

We want to ensure that only requests with a valid API key can access this flow. To achieve this, we can use the apiKey function to create a context provider that validates incoming requests.

The apiKey accepts either a static API key or a function that validates the API key from incoming request.

Static API Key Example

First, let’s look at how we can set up the static API key authentication for our flow:

import { startFlowServer, withFlowOptions } from '@genkit-ai/express';
import { apiKey } from 'genkit/context';

const FLOW_API_KEY = process.env.API_KEY; // The API key is stored in environment variables or secure vault

// Remember in the previous recipe, we used startFlowServer to serve our flows
startFlowServer({
  flows: [
    withFlowOptions(protectedFlowByAPIKey, {
      contextProvider: apiKey(FLOW_API_KEY),
    }),
  ],
  port: 3000,
  cors: true,
});

Now, when a request is made to the destinationExplorer flow, the server will check for the presence of the correct API key in the Authorization header. If the key is missing or incorrect, the request will be denied.

You can test this by making requests with and without the correct API key to see how the flow responds. Here’s an example of how to make a request with the API key using curl:

curl -X POST http://localhost:3000/protectedFlowByAPIKey \
  -H "Authorization: Bearer API_KEY_VALUE" \
  -H "Content-Type: application/json" \
  -d '{ "data": {}}'

Replace API_KEY_VALUE with the actual API key you set in the environment variable. If the key is valid, you should receive a successful response from the flow. Otherwise, you’ll get an authentication error.

If you want to test without an API key, simply omit the Authorization header from your request, and you should see an authentication error as expected, as shown below:

curl -X POST http://localhost:3000/protectedFlowByAPIKey \
  -H "Content-Type: application/json" \
  -d '{ "data": {}}'

Dynamic API Key Example

While static API keys are simple to implement, they may not be suitable for all applications, especially those requiring higher security levels. Think about an application where API keys need to be rotated regularly or where different users have different access levels. In such cases, you can implement a dynamic API key validation mechanism.

This is where you can provide a function to the apiKey context provider that validates the API key dynamically. Here’s an example:

import { startFlowServer, withFlowOptions } from '@genkit-ai/express';

import { apiKey, getHttpStatus } from 'genkit/context';

const validKeys = ['key1', 'key2', 'key3']; // Example list of valid keys

const protectedFlowByDynamicAPIKey = ai.defineFlow(
  {
    name: 'protectedFlowByDynamicAPIKey',
    inputSchema: z.object({}),
    outputSchema: z.string(),
  },
  async (input, { context }) => {
    const apiKey = context?.auth?.apiKey;

    return apiKey;
  },
);

startFlowServer({
  flows: [
    withFlowOptions(protectedFlowByDynamicAPIKey, {
      contextProvider: apiKey(async (context) => {
        // Implement your logic to validate the API key, maybe check against a database or an external service for more complex scenarios
        const validKeys = ['key1', 'key2', 'key3'];
        if (!context.auth.apiKey || !validKeys.includes(context.auth.apiKey)) {
          throw new UserFacingError(
            'UNAUTHENTICATED',
            'Invalid API Key provided.',
          );
        }
      }),
    }),
  ],
  port: 3000,
  cors: true,
});

In this example, the validateApiKey function checks if the provided API key is in the list of valid keys. You can extend this function to include more complex logic, such as checking against a database or an external authentication service.

You can test this dynamic API key validation in the same way as before, by making requests with different API keys to see how the flow responds based on the key’s validity. Here is a curl example with a valid key:

curl -X POST http://localhost:3000/protectedFlowByDynamicAPIKey \
  -H "Authorization: key1" \
  -H "Content-Type: application/json" \
  -d '{ "data": {}}'

You can use an invalid key to see the authentication error:

curl -X POST http://localhost:3000/protectedFlowByDynamicAPIKey \
  -H "Authorization: invalid_key" \
  -H "Content-Type: application/json" \
  -d '{ "data": {}}'

Both static and dynamic API key authentication methods provide a solid foundation for securing your Genkit flows. Depending on your application’s requirements, you can choose the method that best fits your needs.

Here is the complete example of securing a Genkit flow using static API key authentication:

import { genkit, z } from 'genkit';
import { startFlowServer, withFlowOptions } from '@genkit-ai/express';
import { apiKey } from 'genkit/context';

const FLOW_API_KEY = process.env.API_KEY; // The API key is stored in environment variables or secure vault

const protectedFlowByAPIKey = ai.defineFlow(
  {
    name: 'protectedFlowByAPIKey',
    inputSchema: z.object({}),
    outputSchema: z.string(),
  },
  async (input, { context }) => {
    const apiKey = context?.auth?.apiKey;

    return apiKey;
  },
);

startFlowServer({
  flows: [
    withFlowOptions(protectedFlowByAPIKey, {
      contextProvider: apiKey(FLOW_API_KEY),
    }),
  ],
  port: 3000,
  cors: true,
});

Custom Authentication Provider Example

While the API key method is straightforward, if you want to grant access based on user roles or permissions, you can create a custom authentication provider. Here’s an example of how to implement a custom authentication provider:

For this example, we’ll create a mock authentication service that simulates user authentication and role checking.

/**
 * MOCK AUTH SYSTEM
 *
 * This class simulates a real authentication provider (like Better Auth, Auth0, etc.)
 * for the purpose of this recipe. It manages an in-memory database of users and sessions.
 */
export class MockAuthService {
  // Simulating a database of users with tokens
  private users = [
    {
      id: 'usr_admin123',
      name: 'Alice Admin',
      email: 'alice@example.com',
      role: 'admin',
      token: 'admin-secret-token',
    },
    {
      id: 'usr_traveler456',
      name: 'Bob Traveler',
      email: 'bob@example.com',
      role: 'user',
      token: 'user-secret-token',
    },
  ];

  async verifyToken(token: string) {
    // Simulate network/database latency
    await new Promise((resolve) => setTimeout(resolve, 50));

    const user = this.users.find((u) => u.token === token);
    return user || null;
  }
}

Think of this MockAuthService as a placeholder for a real authentication such as Better Auth, Auth0, Firebase Auth, or any other service you might use in a production application. How you implement the authentication logic will depend on the specific service you choose, but the overall structure will be similar.

Now, let’s integrate this mock authentication service into our Genkit flow to secure it based on user roles.

Similar to the previous examples, we will use the ContextProvider interface to create a custom context provider that checks the user’s role before allowing access to the flow.

import { genkit, UserFacingError, z } from 'genkit';
import { ContextProvider, RequestData } from 'genkit/context';

const customContextProvider: ContextProvider<{
  auth: { user: string; role: string; email?: string };
}> = async (req: RequestData) => {
  const authHeader =
    req.headers['Authorization'] || req.headers['authorization'];

  const token =
    typeof authHeader === 'string'
      ? authHeader.replace(/^Bearer\s+/, '')
      : null;

  if (!token) {
    throw new UserFacingError(
      'UNAUTHENTICATED',
      'Missing authentication token.',
    );
  }

  // Verify against our mock service
  const user = await mockAuth.verifyToken(token);

  if (user) {
    console.log(`Authenticated user: ${user.email} (${user.role})`);
    return {
      auth: {
        user: user.id,
        role: user.role,
        email: user.email,
      },
    };
  }

  // If we reach here, no valid session was found
  throw new UserFacingError(
    'UNAUTHENTICATED',
    'Invalid authentication credentials.',
  );
};

In this custom context provider, we extract the token from the Authorization header and verify it using our MockAuthService. If the token is valid, we return the user’s information, including their role. If the token is missing or invalid, we throw a UserFacingError indicating that authentication has failed.

Next, we can modify our flow to check the user’s role before proceeding with the operation, and instead of apiKey, we will use our customContextProvider to secure the flow:

startFlowServer({
  flows: [
    withFlowOptions(protectedFlowByAuthContext, {
      contextProvider: customContextProvider,
    }),
  ],
  port: 3000,
  cors: true,
});

With this setup, only authenticated users with valid tokens can access the destinationExplorer flow. You can further enhance the security by checking the user’s role within the flow implementation to restrict access to certain features based on their permissions. For example, you might want to allow only users with the admin role to access specific destinations or perform certain actions.

You can test this custom authentication provider by making requests with valid and invalid tokens in the Authorization header, similar to the previous API key examples. Here’s an example of a request with a valid token:

curl -X POST http://localhost:3000/protectedFlowByAuthContext \
  -H "Authorization: Bearer admin-secret-token" \
  -H "Content-Type: application/json" \
  -d '{ "data": {}}'

And here is an example of a request with an invalid token:

curl -X POST http://localhost:3000/protectedFlowByAuthContext \
  -H "Authorization: Bearer invalid-token" \
  -H "Content-Type: application/json" \
  -d '{ "data": {}}'

Here is a complete example of a secured Genkit flow using a custom authentication provider:

import { genkit, UserFacingError, z } from 'genkit';
import { startFlowServer, withFlowOptions } from '@genkit-ai/express';
import { ContextProvider, RequestData } from 'genkit/context';

// Mock authentication service
class MockAuthService {
  private users = [
    {
      id: 'usr_admin123',
      name: 'Alice Admin',
      email: 'alice.admin@example.com',
      role: 'admin',
      token: 'admin-secret-token',
    },
    {
      id: 'usr_traveler456',
      name: 'Bob Traveler',
      email: 'bob.traveler@example.com',
      role: 'traveler',
      token: 'traveler-secret-token',
    },
  ];

  async verifyToken(token: string) {
    await new Promise((resolve) => setTimeout(resolve, 50));
    const user = this.users.find((u) => u.token === token);
    return user || null;
  }
}

const mockAuth = new MockAuthService();

const customContextProvider: ContextProvider<{
  auth: { user: string; role: string; email?: string };
}> = async (req: RequestData) => {
  const authHeader =
    req.headers['Authorization'] || req.headers['authorization'];

  const token =
    typeof authHeader === 'string'
      ? authHeader.replace(/^Bearer\s+/, '')
      : null;

  if (!token) {
    throw new UserFacingError(
      'UNAUTHENTICATED',
      'Missing authentication token.',
    );
  }

  const user = await mockAuth.verifyToken(token);

  if (user) {
    console.log(`Authenticated user: ${user.email} (${user.role})`);
    return {
      auth: {
        user: user.id,
        role: user.role,
        email: user.email,
      },
    };
  }

  throw new UserFacingError(
    'UNAUTHENTICATED',
    'Invalid authentication credentials.',
  );
};

const protectedFlowByAuthContext = ai.defineFlow(
  {
    name: 'protectedFlowByAuthContext',
    inputSchema: z.null(),
    outputSchema: z.string(), // We return raw Markdown text
  },
  async (input, { context }) => {
    console.log('Auth context:', context?.auth);

    return `# Welcome, ${context?.auth?.email || 'User'}!`;
  },
);

startFlowServer({
  flows: [
    withFlowOptions(protectedFlowByAuthContext, {
      contextProvider: customContextProvider,
    }),
  ],
  port: 3000,
  cors: true,
});

Accessing Context in Flow Implementation

If you are doing role-based access control, you might want to have access to the current users’ context within the flow implementation. This is important for making decisions based on the user’s role or other attributes.

For example, let’s say you want to generate different travel suggestions based on whether the user travel history or preferences. It’s not ideal to have that information as part of the flows input. There is where accessing the context becomes useful.

Another possible scenario, say you have a saas application where users submit their documents for analysis. You want to ensure that only the owner of the document can access the analysis results. By accessing the user context within the flow, you can restrict your retrieval tool to only fetch documents that belong to the authenticated user or organization.

Genkit flow provides you with the authenticated context. The second argument of the flow implementation function, is an object that contains the context property. And inside the context, you will find the authentication information provided by your context provider.

const protectedFlowByAPIKey = ai.defineFlow(
  {
    name: 'destinationExplorer',
  },
  async (input, { context }) => {
    const apiKey = context?.auth?.apiKey;

    // use something for the API key validation
  },
);

Whatever information you return from your context provider will be available inside the flow implementation via the context property, under context.auth. This allows you to make informed decisions based on the user’s authentication status and attributes, enhancing the security and functionality of your Genkit application.

Conclusion

This recipe one question that I often get is, “How do I secure my Genkit flows?” and hopefully, this recipe has provided you with a clear understanding of how to implement authentication and authorization mechanisms in your Genkit flows. By leveraging Genkit’s context provider system, you can easily integrate various authentication methods, such as API keys or custom providers, to ensure that only authorized users can access your flows.

Remember, security is an ongoing process, and it’s essential to regularly review and update your security measures to stay ahead of potential threats. Always follow best practices for authentication and authorization to protect your applications and user data.

🚀 Get the Complete Code

This recipe includes a full GitHub repository with working code, setup instructions, and examples. Unlock access to build faster and support this cookbook!

Complete, tested code repositories

Copy-paste ready examples

Lifetime access with updates

Support future recipes & content

Unlock All Recipes - 50% Off Already purchased? Sign in

Happy coding and stay secure!