Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(frontend): Implement aws-js-sdk crendentials to support IRSA for s3 #8651

Merged
merged 8 commits into from
Jan 25, 2023
Merged
94 changes: 1 addition & 93 deletions frontend/server/aws-helper.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -11,99 +11,7 @@
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
import fetch from 'node-fetch';
import { awsInstanceProfileCredentials, isS3Endpoint } from './aws-helper';

// mock node-fetch module
jest.mock('node-fetch');

function sleep(ms: number) {
return new Promise(resolve => setTimeout(resolve, ms));
}

beforeEach(() => {
awsInstanceProfileCredentials.reset();
jest.clearAllMocks();
});

describe('awsInstanceProfileCredentials', () => {
const mockedFetch: jest.Mock = fetch as any;

describe('getCredentials', () => {
it('retrieves, caches, and refreshes the AWS EC2 instance profile and session credentials everytime it is called.', async () => {
let count = 0;
const expectedCredentials = [
{
AccessKeyId: 'AccessKeyId',
Code: 'Success',
Expiration: new Date(Date.now() + 1000).toISOString(), // expires 1 sec later
LastUpdated: '2019-12-17T10:55:38Z',
SecretAccessKey: 'SecretAccessKey',
Token: 'SessionToken',
Type: 'AWS-HMAC',
},
{
AccessKeyId: 'AccessKeyId2',
Code: 'Success',
Expiration: new Date(Date.now() + 10000).toISOString(), // expires 10 sec later
LastUpdated: '2019-12-17T10:55:38Z',
SecretAccessKey: 'SecretAccessKey2',
Token: 'SessionToken2',
Type: 'AWS-HMAC',
},
];
const mockFetch = (url: string) => {
if (url === 'http://169.254.169.254/latest/meta-data/iam/security-credentials/') {
return Promise.resolve({ text: () => Promise.resolve('some_iam_role') });
}
return Promise.resolve({
json: () => Promise.resolve(expectedCredentials[count++]),
});
};
mockedFetch.mockImplementation(mockFetch);

// expect to get cred from ec2 instance metadata store
expect(await awsInstanceProfileCredentials.getCredentials()).toBe(expectedCredentials[0]);
// expect to call once for profile name, and once for credential
expect(mockedFetch.mock.calls.length).toBe(2);
// expect to get same credential as it has not expire
expect(await awsInstanceProfileCredentials.getCredentials()).toBe(expectedCredentials[0]);
// expect to not to have any more calls
expect(mockedFetch.mock.calls.length).toBe(2);
// let credential expire
await sleep(1500);
// expect to get new cred as old one expire
expect(await awsInstanceProfileCredentials.getCredentials()).toBe(expectedCredentials[1]);
// expect to get same cred as it has not expire
expect(await awsInstanceProfileCredentials.getCredentials()).toBe(expectedCredentials[1]);
});

it('fails gracefully if there is no instance profile.', async () => {
const mockFetch = (url: string) => {
if (url === 'http://169.254.169.254/latest/meta-data/iam/security-credentials/') {
return Promise.resolve({ text: () => Promise.resolve('') });
}
return Promise.reject('Unknown error');
};
mockedFetch.mockImplementation(mockFetch);

expect(await awsInstanceProfileCredentials.ok()).toBeFalsy();
expect(async () => await awsInstanceProfileCredentials.getCredentials()).not.toThrow();
expect(await awsInstanceProfileCredentials.getCredentials()).toBeUndefined();
});

it('fails gracefully if there is no metadata store.', async () => {
const mockFetch = (_: string) => {
return Promise.reject('Unknown error');
};
mockedFetch.mockImplementation(mockFetch);

expect(await awsInstanceProfileCredentials.ok()).toBeFalsy();
expect(async () => await awsInstanceProfileCredentials.getCredentials()).not.toThrow();
expect(await awsInstanceProfileCredentials.getCredentials()).toBeUndefined();
});
});
});
import { isS3Endpoint } from './aws-helper';

describe('isS3Endpoint', () => {
it('checks a valid s3 endpoint', () => {
Expand Down
107 changes: 0 additions & 107 deletions frontend/server/aws-helper.ts
Original file line number Diff line number Diff line change
Expand Up @@ -11,40 +11,6 @@
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
import fetch from 'node-fetch';

/** AWSMetadataCredentials describes the credentials provided by aws metadata store. */
export interface AWSMetadataCredentials {
Code: string;
LastUpdated: string;
Type: string;
AccessKeyId: string;
SecretAccessKey: string;
Token: string;
Expiration: string;
}

/** url for aws metadata store. */
const metadataUrl = 'http://169.254.169.254/latest/meta-data';

/**
* Get the AWS IAM instance profile.
*/
async function getIAMInstanceProfile(): Promise<string | undefined> {
try {
const resp = await fetch(`${metadataUrl}/iam/security-credentials/`);
const profiles = (await resp.text()).split('\n');
if (profiles.length > 0) {
return profiles[0].trim(); // return first profile
}
return;
} catch (error) {
console.error(
`Unable to fetch credentials from AWS metadata store (${metadataUrl}/iam/security-credentials/): ${error}`,
);
return;
}
}

/**
* Check if the provided string is an S3 endpoint (can be any region).
Expand All @@ -54,76 +20,3 @@ async function getIAMInstanceProfile(): Promise<string | undefined> {
export function isS3Endpoint(endpoint: string = ''): boolean {
return !!endpoint.match(/s3.{0,}\.amazonaws\.com\.?.{0,}/i);
}

/**
* Class to handle the session credentials for AWS ec2 instance profile.
*/
class AWSInstanceProfileCredentials {
_iamProfile?: string;
_credentials?: AWSMetadataCredentials;
_expiration: number = 0;

/** reset all caches */
reset() {
this._iamProfile = undefined;
this._credentials = undefined;
this._expiration = 0;
return this;
}

/**
* EC2 Instance profile
*/
async profile() {
if (!this._iamProfile) {
try {
this._iamProfile = await getIAMInstanceProfile();
} catch (err) {
console.error(err);
}
}
return this._iamProfile;
}

/**
* Return true only if there is a metadata store and instance profile.
*/
async ok() {
try {
const profile = await this.profile();
return profile && profile.length > 0;
} catch (_) {
return false;
}
}

async _fetchCredentials(): Promise<AWSMetadataCredentials | undefined> {
try {
const profile = await this.profile();
const resp = await fetch(`${metadataUrl}/iam/security-credentials/${profile}`);
const credentials = await resp.json();
return credentials;
} catch (error) {
console.error(`Unable to fetch credentials from AWS metadata store: ${error}`);
return;
}
}

/**
* Get the AWS metadata store session credentials.
*/
async getCredentials(): Promise<AWSMetadataCredentials | undefined> {
// query for credentials if going to expire or no credentials yet
if (Date.now() + 10 >= this._expiration || !this._credentials) {
this._credentials = await this._fetchCredentials();
if (this._credentials && this._credentials.Expiration) {
this._expiration = new Date(this._credentials.Expiration).getTime();
} else {
this._expiration = -1; // always retry
}
}
return this._credentials;
}
}

export const awsInstanceProfileCredentials = new AWSInstanceProfileCredentials();
1 change: 0 additions & 1 deletion frontend/server/handlers/artifacts.ts
Original file line number Diff line number Diff line change
Expand Up @@ -95,7 +95,6 @@ export function getArtifactsHandler({
peek,
)(req, res);
break;

case 's3':
getMinioArtifactHandler(
{
Expand Down
29 changes: 0 additions & 29 deletions frontend/server/minio-helper.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -14,12 +14,10 @@
import * as zlib from 'zlib';
import { PassThrough } from 'stream';
import { Client as MinioClient } from 'minio';
import { awsInstanceProfileCredentials } from './aws-helper';
import { createMinioClient, isTarball, maybeTarball, getObjectStream } from './minio-helper';
import { V1beta1JobTemplateSpec } from '@kubernetes/client-node';

jest.mock('minio');
jest.mock('./aws-helper');
ryansteakley marked this conversation as resolved.
Show resolved Hide resolved

describe('minio-helper', () => {
const MockedMinioClient: jest.Mock = MinioClient as any;
Expand Down Expand Up @@ -54,33 +52,6 @@ describe('minio-helper', () => {
endPoint: 'minio.kubeflow:80',
});
});

it('uses EC2 metadata credentials if access key are not provided.', async () => {
(awsInstanceProfileCredentials.getCredentials as jest.Mock).mockImplementation(() =>
Promise.resolve({
AccessKeyId: 'AccessKeyId',
Code: 'Success',
Expiration: new Date(Date.now() + 1000).toISOString(), // expires 1 sec later
LastUpdated: '2019-12-17T10:55:38Z',
SecretAccessKey: 'SecretAccessKey',
Token: 'SessionToken',
Type: 'AWS-HMAC',
}),
);
(awsInstanceProfileCredentials.ok as jest.Mock).mockImplementation(() =>
Promise.resolve(true),
);

const client = await createMinioClient({ endPoint: 's3.awsamazon.com' });

expect(client).toBeInstanceOf(MinioClient);
expect(MockedMinioClient).toHaveBeenCalledWith({
accessKey: 'AccessKeyId',
endPoint: 's3.awsamazon.com',
secretKey: 'SecretAccessKey',
sessionToken: 'SessionToken',
});
});
});

describe('isTarball', () => {
Expand Down
28 changes: 12 additions & 16 deletions frontend/server/minio-helper.ts
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ import * as tar from 'tar-stream';
import peek from 'peek-stream';
import gunzip from 'gunzip-maybe';
import { Client as MinioClient, ClientOptions as MinioClientOptions } from 'minio';
import { awsInstanceProfileCredentials } from './aws-helper';
const { fromNodeProviderChain } = require('@aws-sdk/credential-providers');

/** MinioRequestConfig describes the info required to retrieve an artifact. */
export interface MinioRequestConfig {
Expand All @@ -38,23 +38,19 @@ export interface MinioClientOptionsWithOptionalSecrets extends Partial<MinioClie
*/
export async function createMinioClient(config: MinioClientOptionsWithOptionalSecrets) {
if (!config.accessKey || !config.secretKey) {
try {
if (await awsInstanceProfileCredentials.ok()) {
const credentials = await awsInstanceProfileCredentials.getCredentials();
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In the design doc, you mention that we will fail back to existing solution if the IRSA doesn't work. It doesn't look like the case in this block. Can we retain the existing solution?

if (credentials) {
const {
AccessKeyId: accessKey,
SecretAccessKey: secretKey,
Token: sessionToken,
} = credentials;
return new MinioClient({ ...config, accessKey, secretKey, sessionToken });
}
console.error('unable to get credentials from AWS metadata store.');
}
} catch (err) {
console.error('Unable to get aws instance profile credentials: ', err);
const credentials = fromNodeProviderChain();
const aws_credentials = await credentials();
if (aws_credentials) {
const {
accessKeyId: accessKey,
secretAccessKey: secretKey,
sessionToken: sessionToken,
} = aws_credentials;
return new MinioClient({ ...config, accessKey, secretKey, sessionToken });
}
console.error('unable to get credentials from AWS credential provider chain.');
}

return new MinioClient(config as MinioClientOptions);
}

Expand Down
Loading