Skip to content

Commit

Permalink
🔒✉️ feat: allow only certain domain (danny-avila#1562)
Browse files Browse the repository at this point in the history
* feat: allow only certain domain

* Update dotenv.md

* refactor( registrationController) & handle ALLOWED_REGISTRATION_DOMAINS not specified

* cleanup and moved to AuthService for better  error handling

* refactor: replace environment variable with librechat config item, add typedef for custom config, update docs for new registration object and allowedDomains values

* ci(AuthService): test for `isDomainAllowed`

---------

Co-authored-by: Danny Avila <messagedaniel@protonmail.com>
  • Loading branch information
berry-13 and danny-avila authored Feb 5, 2024
1 parent 7d9ea65 commit 58b542d
Show file tree
Hide file tree
Showing 8 changed files with 136 additions and 21 deletions.
4 changes: 3 additions & 1 deletion api/cache/getCustomConfig.js
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,9 @@ const getLogStores = require('./getLogStores');

/**
* Retrieves the configuration object
* @function getCustomConfig */
* @function getCustomConfig
* @returns {Promise<TCustomConfig | null>}
* */
async function getCustomConfig() {
const cache = getLogStores(CacheKeys.CONFIG_STORE);
let customConfig = await cache.get(CacheKeys.CUSTOM_CONFIG);
Expand Down
29 changes: 29 additions & 0 deletions api/server/services/AuthService.js
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
const crypto = require('crypto');
const bcrypt = require('bcryptjs');
const { registerSchema, errorsToString } = require('~/strategies/validators');
const getCustomConfig = require('~/cache/getCustomConfig');
const Token = require('~/models/schema/tokenSchema');
const { sendEmail } = require('~/server/utils');
const Session = require('~/models/Session');
Expand All @@ -12,6 +13,27 @@ const domains = {
server: process.env.DOMAIN_SERVER,
};

async function isDomainAllowed(email) {
if (!email) {
return false;
}

const domain = email.split('@')[1];

if (!domain) {
return false;
}

const customConfig = await getCustomConfig();
if (!customConfig) {
return true;
} else if (!customConfig?.registration?.allowedDomains) {
return true;
}

return customConfig.registration.allowedDomains.includes(domain);
}

const isProduction = process.env.NODE_ENV === 'production';

/**
Expand Down Expand Up @@ -80,6 +102,12 @@ const registerUser = async (user) => {
return { status: 500, message: 'Something went wrong' };
}

if (!(await isDomainAllowed(email))) {
const errorMessage = 'Registration from this domain is not allowed.';
logger.error(`[registerUser] [Registration not allowed] [Email: ${user.email}]`);
return { status: 403, message: errorMessage };
}

//determine if this is the first registered user (not counting anonymous_user)
const isFirstRegisteredUser = (await User.countDocuments({})) === 0;

Expand Down Expand Up @@ -239,6 +267,7 @@ const setAuthTokens = async (userId, res, sessionId = null) => {
module.exports = {
registerUser,
logoutUser,
isDomainAllowed,
requestPasswordReset,
resetPassword,
setAuthTokens,
Expand Down
39 changes: 39 additions & 0 deletions api/server/services/AuthService.spec.js
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
const getCustomConfig = require('~/cache/getCustomConfig');
const { isDomainAllowed } = require('./AuthService');

jest.mock('~/cache/getCustomConfig', () => jest.fn());

describe('isDomainAllowed', () => {
it('should allow domain when customConfig is not available', async () => {
getCustomConfig.mockResolvedValue(null);
await expect(isDomainAllowed('test@domain1.com')).resolves.toBe(true);
});

it('should allow domain when allowedDomains is not defined in customConfig', async () => {
getCustomConfig.mockResolvedValue({});
await expect(isDomainAllowed('test@domain1.com')).resolves.toBe(true);
});

it('should reject an email if it is falsy', async () => {
getCustomConfig.mockResolvedValue({});
await expect(isDomainAllowed('')).resolves.toBe(false);
});

it('should allow a domain if it is included in the allowedDomains', async () => {
getCustomConfig.mockResolvedValue({
registration: {
allowedDomains: ['domain1.com', 'domain2.com'],
},
});
await expect(isDomainAllowed('user@domain1.com')).resolves.toBe(true);
});

it('should reject a domain if it is not included in the allowedDomains', async () => {
getCustomConfig.mockResolvedValue({
registration: {
allowedDomains: ['domain1.com', 'domain2.com'],
},
});
await expect(isDomainAllowed('user@domain3.com')).resolves.toBe(false);
});
});
6 changes: 6 additions & 0 deletions api/typedefs.js
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,12 @@
* @memberof typedefs
*/

/**
* @exports TCustomConfig
* @typedef {import('librechat-data-provider').TCustomConfig} TCustomConfig
* @memberof typedefs
*/

/**
* @exports TMessage
* @typedef {import('librechat-data-provider').TMessage} TMessage
Expand Down
62 changes: 44 additions & 18 deletions docs/install/configuration/custom_config.md
Original file line number Diff line number Diff line change
Expand Up @@ -4,21 +4,6 @@ description: Comprehensive guide for configuring the `librechat.yaml` file AKA t
weight: -10
---

<!-- # Table of Contents
- [Intro](#librechat-configuration-guide)
- [Setup](#setup)
- [Docker Setup](#docker-setup)
- [Config Structure](#config-structure)
- [1. Version](#1-version)
- [2. Cache Settings](#2-cache-settings)
- [3. Endpoints](#3-endpoints)
- [Endpoint Object Structure](#endpoint-object-structure)
- [Additional Notes](#additional-notes)
- [Default Parameters](#default-parameters)
- [Breakdown of Default Params](#breakdown-of-default-params)
- [Example Config](#example-config) -->

# LibreChat Configuration Guide

Welcome to the guide for configuring the **librechat.yaml** file in LibreChat.
Expand All @@ -43,8 +28,11 @@ Stay tuned for ongoing enhancements to customize your LibreChat instance!
- [Version](#version)
- [Cache Settings](#cache-settings)
- [File Strategy](#file-strategy)
- [Registration](#registration)
- [Endpoints](#endpoints)
- [Endpoint Object Structure](#endpoint-object-structure)
- [Registration Object Structure](#registration-object-structure)
- [**allowedDomains**:](#allowedDomains)
- [Custom Endpoint Object Structure](#custom-endpoint-object-structure)
- [**name**:](#name)
- [**apiKey**:](#apikey)
- [**baseURL**:](#baseurl)
Expand Down Expand Up @@ -120,16 +108,50 @@ docker-compose up # no need to rebuild
- **Description**: Determines where to save user uploaded/generated files. Defaults to `"local"` if omitted.
- **Example**: `fileStrategy: "firebase"`

### Registration
- **Key**: `registration`
- **Type**: Object
- **Description**: Configures registration-related settings for the application.
- **Sub-Key**: `allowedDomains`
- **Type**: Array of Strings
- **Description**: Specifies a list of allowed email domains for user registration. Users attempting to register with email domains not listed here will be restricted from registering.
- [Registration Object Structure](#registration-object-structure)

### Endpoints
- **Key**: `endpoints`
- **Type**: Object
- **Description**: Defines custom API endpoints for the application.
- **Sub-Key**: `custom`
- **Type**: Array of Objects
- **Description**: Each object in the array represents a unique endpoint configuration.
- [Custom Endpoint Object Structure](#custom-endpoint-object-structure)
- **Required**

## Endpoint Object Structure
## Registration Object Structure

```yaml
# Example Registration Object Structure
registration:
allowedDomains:
- "gmail.com"
- "protonmail.com"
```
### **allowedDomains**:
> A list specifying allowed email domains for registration.
- Type: Array of Strings
- Example:
```yaml
allowedDomains:
- "gmail.com"
- "protonmail.com"
```
- **Required**
- **Note**: Users with email domains not listed will be restricted from registering.
## Custom Endpoint Object Structure
Each endpoint in the `custom` array should have the following structure:

```yaml
Expand Down Expand Up @@ -345,8 +367,12 @@ Custom endpoints share logic with the OpenAI endpoint, and thus have default par
## Example Config

```yaml
version: 1.0.1
version: 1.0.2
cache: true
# Example Registration Object Structure
registration:
allowedDomains:
- "gmail.com"
endpoints:
custom:
# Mistral AI API
Expand Down
3 changes: 2 additions & 1 deletion docs/install/configuration/dotenv.md
Original file line number Diff line number Diff line change
Expand Up @@ -665,7 +665,8 @@ see: **[User/Auth System](../configuration/user_auth_system.md)**
```bash
ALLOW_EMAIL_LOGIN=true
ALLOW_REGISTRATION=true
ALLOW_REGISTRATION=true
ALLOWED_REGISTRATION_DOMAINS=
ALLOW_SOCIAL_LOGIN=false
ALLOW_SOCIAL_REGISTRATION=false
```
Expand Down
7 changes: 6 additions & 1 deletion librechat.example.yaml
Original file line number Diff line number Diff line change
@@ -1,9 +1,14 @@
# Configuration version (required)
version: 1.0.1
version: 1.0.2

# Cache settings: Set to true to enable caching
cache: true

# Example Registration Object Structure (optional)
# registration:
# allowedDomains:
# - "gmail.com"

# Definition of custom endpoints
endpoints:
custom:
Expand Down
7 changes: 7 additions & 0 deletions packages/data-provider/src/config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -30,13 +30,20 @@ export const configSchema = z.object({
version: z.string(),
cache: z.boolean(),
fileStrategy: fileSourceSchema.optional(),
registration: z
.object({
allowedDomains: z.array(z.string()).optional(),
})
.optional(),
endpoints: z
.object({
custom: z.array(endpointSchema.partial()),
})
.strict(),
});

export type TCustomConfig = z.infer<typeof configSchema>;

export enum KnownEndpoints {
mistral = 'mistral',
openrouter = 'openrouter',
Expand Down

0 comments on commit 58b542d

Please sign in to comment.