Skip to content

Commit

Permalink
feat: Bring back demo regime of core working with IDP #150 (#153)
Browse files Browse the repository at this point in the history 
Co-authored-by: Aliaksandr Stsiapanay <[email protected]>
  • Loading branch information
@astsiapanay
astsiapanay and astsiapanay authored Jan 19, 2024
1 parent 8e4fcd9 commit 500b1ab
Show file tree
Hide file tree
Showing 3 changed files with 95 additions and 30 deletions.
64 changes: 38 additions & 26 deletions README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,58 +2,69 @@

## Overview

HTTP Proxy provides unified API to different chat completion and embedding models, assistants and applications.
HTTP Proxy provides unified API to different chat completion and embedding models, assistants and applications.
Written in Java 17 and built on top of [Eclipse Vert.x](https://vertx.io/).

## Build

Build the project with Gradle and Java 17:

```
./gradlew build
```

## Run

Run the project with Gradle:

```
./gradlew run
```

Or run com.epam.aidial.core.AIDial class from your favorite IDE.

## Configuration

### Static settings

Static settings are used on startup and cannot be changed while application is running. Priority order:

* Environment variables with extra "aidial." prefix. E.g. "aidial.server.port", "aidial.config.files".
* File specified in "AIDIAL_SETTINGS" environment variable.
* Default resource file: src/main/resources/aidial.settings.json.

| Setting | Default |Description
|-----------------------------------------|--------------------|-
| config.files | aidial.config.json |Config files with parts of the whole config.
| config.reload | 60000 |Config reload interval in milliseconds.
| identityProviders | - |List of identity providers
| identityProviders.*.jwksUrl | - |Url to jwks provider.
| identityProviders.*.rolePath | - |Path to the claim user roles in JWT token, e.g. `resource_access.chatbot-ui.roles` or just `roles`.
| identityProviders.*.loggingKey | - |User information to search in claims of JWT token.
| identityProviders.*.loggingSalt | - |Salt to hash user information for logging.
| identityProviders.*.cacheSize | 10 |How many JWT tokens to cache.
| identityProviders.*.cacheExpiration | 10 |How long to retain JWT token in cache.
| identityProviders.*.cacheExpirationUnit | MINUTES |Unit of cache expiration.
| identityProviders.*.issuerPattern | - |Regexp to match the claim "iss" to identity provider
| vertx.* | - |Vertx settings.
| server.* | - |Vertx HTTP server settings for incoming requests.
| client.* | - |Vertx HTTP client settings for outbound requests.
| storage.provider | - |Specifies blob storage provider. Supported providers: s3, aws-s3, azureblob, google-cloud-storage
| storage.endpoint | - |Optional. Specifies endpoint url for s3 compatible storages
| storage.identity | - |Blob storage access key
| storage.credential | - |Blob storage secret key
| storage.bucket | - |Blob storage bucket
| storage.createBucket | false |Indicates whether bucket should be created on start-up
| encryption.password | - |Password used for AES encryption
| encryption.salt | - |Salt used for AES encryption
| Setting | Default | Description
|--------------------------------------------|--------------------|-------------------------------------------------------------------------------------------------------------------
| config.files | aidial.config.json | Config files with parts of the whole config.
| config.reload | 60000 | Config reload interval in milliseconds.
| identityProviders | - | List of identity providers. **Note**. At least one identity provider must be provided.
| identityProviders.*.jwksUrl | - | Url to jwks provider. **Required** if `disabledVerifyJwt` is set to `false`
| identityProviders.*.rolePath | - | Path to the claim user roles in JWT token, e.g. `resource_access.chatbot-ui.roles` or just `roles`. **Required**.
| identityProviders.*.loggingKey | - | User information to search in claims of JWT token.
| identityProviders.*.loggingSalt | - | Salt to hash user information for logging.
| identityProviders.*.cacheSize | 10 | How many JWT tokens to cache.
| identityProviders.*.cacheExpiration | 10 | How long to retain JWT token in cache.
| identityProviders.*.cacheExpirationUnit | MINUTES | Unit of cache expiration.
| identityProviders.*.issuerPattern | - | Regexp to match the claim "iss" to identity provider
| identityProviders.*.disableJwtVerification | false | The flag disables JWT verification
| vertx.* | - | Vertx settings.
| server.* | - | Vertx HTTP server settings for incoming requests.
| client.* | - | Vertx HTTP client settings for outbound requests.
| storage.provider | - | Specifies blob storage provider. Supported providers: s3, aws-s3, azureblob, google-cloud-storage
| storage.endpoint | - | Optional. Specifies endpoint url for s3 compatible storages
| storage.identity | - | Blob storage access key
| storage.credential | - | Blob storage secret key
| storage.bucket | - | Blob storage bucket
| storage.createBucket | false | Indicates whether bucket should be created on start-up
| encryption.password | - | Password used for AES encryption
| encryption.salt | - | Salt used for AES encryption

### Dynamic settings
Dynamic settings are stored in JSON files, specified via "config.files" static setting, and reloaded at interval, specified via "config.reload" static setting.

Dynamic settings are stored in JSON files, specified via "config.files" static setting, and reloaded at interval,
specified via "config.reload" static setting.
Dynamic settings include:

* Models
* Assistant
* Applications
Expand All @@ -62,6 +73,7 @@ Dynamic settings include:
* Rate Limits

## License

Copyright (C) 2023 EPAM Systems

Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
32 changes: 28 additions & 4 deletions src/main/java/com/epam/aidial/core/security/IdentityProvider.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,35 +31,54 @@ public class IdentityProvider {

public static final ExtractedClaims CLAIMS_WITH_EMPTY_ROLES = new ExtractedClaims(null, Collections.emptyList(), null);

// path to the claim of user roles in JWT
private final String[] rolePath;

private final JwkProvider jwkProvider;

// in memory cache store results obtained from JWK provider
private final ConcurrentHashMap<String, Future<JwkResult>> cache = new ConcurrentHashMap<>();

// the name of the claim in JWT to extract user email
private final String loggingKey;
// random salt is used to digest user email
private final String loggingSalt;

private final MessageDigest sha256Digest;

// the flag determines if user email should be obfuscated
private final boolean obfuscateUserEmail;

private final Vertx vertx;

// the duration is how many milliseconds success JWK result should be stored in the cache
private final long positiveCacheExpirationMs;

// the duration is how many milliseconds failed JWK result should be stored in the cache
private final long negativeCacheExpirationMs;

// the pattern is used to match if the given JWT can be verified by the current provider
private final Pattern issuerPattern;

// the flag disables JWT verification
private final boolean disableJwtVerification;

public IdentityProvider(JsonObject settings, Vertx vertx, Function<String, JwkProvider> jwkProviderSupplier) {
if (settings == null) {
throw new IllegalArgumentException("Identity provider settings are missed");
}
this.vertx = vertx;
positiveCacheExpirationMs = settings.getLong("positiveCacheExpirationMs", TimeUnit.MINUTES.toMillis(10));
negativeCacheExpirationMs = settings.getLong("negativeCacheExpirationMs", TimeUnit.SECONDS.toMillis(10));
String jwksUrl = Objects.requireNonNull(settings.getString("jwksUrl"), "jwksUrl is missed");

disableJwtVerification = settings.getBoolean("disableJwtVerification", false);
if (disableJwtVerification) {
jwkProvider = null;
} else {
String jwksUrl = Objects.requireNonNull(settings.getString("jwksUrl"), "jwksUrl is missed");
jwkProvider = jwkProviderSupplier.apply(jwksUrl);
}

rolePath = Objects.requireNonNull(settings.getString("rolePath"), "rolePath is missed").split("\\.");

loggingKey = settings.getString("loggingKey");
Expand All @@ -69,8 +88,6 @@ public IdentityProvider(JsonObject settings, Vertx vertx, Function<String, JwkPr
loggingSalt = null;
}

jwkProvider = jwkProviderSupplier.apply(jwksUrl);

try {
sha256Digest = MessageDigest.getInstance("SHA-256");
} catch (NoSuchAlgorithmException e) {
Expand Down Expand Up @@ -191,7 +208,14 @@ Future<ExtractedClaims> extractClaims(DecodedJWT decodedJwt) {
if (decodedJwt == null) {
return Future.failedFuture(new IllegalArgumentException("decoded JWT must not be null"));
}
return verifyJwt(decodedJwt).map(jwt -> new ExtractedClaims(extractUserSub(jwt), extractUserRoles(jwt), extractUserHash(jwt)));
if (disableJwtVerification) {
return Future.succeededFuture(from(decodedJwt));
}
return verifyJwt(decodedJwt).map(this::from);
}

private ExtractedClaims from(DecodedJWT jwt) {
return new ExtractedClaims(extractUserSub(jwt), extractUserRoles(jwt), extractUserHash(jwt));
}

boolean match(DecodedJWT jwt) {
Expand Down
29 changes: 29 additions & 0 deletions src/test/java/com/epam/aidial/core/security/IdentityProviderTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verifyNoInteractions;
import static org.mockito.Mockito.when;

@ExtendWith(MockitoExtension.class)
Expand All @@ -61,6 +62,8 @@ public void beforeEach() {
settings.put("jwksUrl", "http://host/jwks");
settings.put("rolePath", "roles");
settings.put("issuerPattern", "issuer");
settings.put("loggingKey", "email");
settings.put("loggingSalt", "salt");
}

@Test
Expand Down Expand Up @@ -291,6 +294,32 @@ public void testExtractClaims_11() throws JwkException {
});
}

@Test
public void testExtractClaims_12() {
settings.put("disableJwtVerification", Boolean.TRUE);
IdentityProvider identityProvider = new IdentityProvider(settings, vertx, url -> jwkProvider);
Algorithm algorithm = Algorithm.RSA256((RSAPublicKey) keyPair.getPublic(), (RSAPrivateKey) keyPair.getPrivate());

String token = JWT.create().withHeader(Map.of("kid", "kid1"))
.withClaim("roles", List.of("role"))
.withClaim("email", "[email protected]")
.withClaim("sub", "sub").sign(algorithm);

Future<ExtractedClaims> result = identityProvider.extractClaims(JWT.decode(token));

verifyNoInteractions(jwkProvider);

assertNotNull(result);
result.onComplete(res -> {
assertTrue(res.succeeded());
ExtractedClaims claims = res.result();
assertNotNull(claims);
assertEquals(List.of("role"), claims.userRoles());
assertEquals("sub", claims.sub());
assertNotNull(claims.userHash());
});
}

@Test
public void testMatch_Failure() {
IdentityProvider identityProvider = new IdentityProvider(settings, vertx, url -> jwkProvider);
Expand Down

0 comments on commit 500b1ab

Please sign in to comment.