Add support for short-lived tokens (#224)

The setup-gradle action tries to get a short-lived access token given the supplied Develocity access key.
This key can be passed either with the `DEVELOCITY_ACCESS_KEY` env var or via the  `develocity-access-key` input parameter.
If a token can be retrieved, then the `DEVELOCITY_ACCESS_KEY` env var will be set to the token. 
Otherwise the `DEVELOCITY_ACCESS_KEY` will be set to a blank string, to avoid a leak.

---------

Co-authored-by: daz <daz@gradle.com>

sources/src/configuration.ts

@@ -200,6 +200,14 @@ export class BuildScanConfig {
         return this.getTermsOfUseProp('build-scan-terms-of-use-agree', 'build-scan-terms-of-service-agree')
     }
 
+    getDevelocityAccessKey(): string {
+        return core.getInput('develocity-access-key') || process.env['DEVELOCITY_ACCESS_KEY'] || ''
+    }
+
+    getDevelocityTokenExpiry(): string {
+        return core.getInput('develocity-token-expiry')
+    }
+
     private verifyTermsOfUseAgreement(): boolean {
         if (
             (this.getBuildScanTermsOfUseUrl() !== 'https://gradle.com/terms-of-service' &&

sources/src/develocity/build-scan.ts

@@ -1,7 +1,8 @@
 import * as core from '@actions/core'
-import {BuildScanConfig} from './configuration'
+import {BuildScanConfig} from '../configuration'
+import {setupToken} from './short-lived-token'
 
-export function setup(config: BuildScanConfig): void {
+export async function setup(config: BuildScanConfig): Promise<void> {
     maybeExportVariable('DEVELOCITY_INJECTION_INIT_SCRIPT_NAME', 'gradle-actions.inject-develocity.init.gradle')
     maybeExportVariable('DEVELOCITY_AUTO_INJECTION_CUSTOM_VALUE', 'gradle-actions')
     if (config.getBuildScanPublishEnabled()) {
@@ -11,6 +12,16 @@ export function setup(config: BuildScanConfig): void {
         maybeExportVariable('DEVELOCITY_TERMS_OF_USE_URL', config.getBuildScanTermsOfUseUrl())
         maybeExportVariable('DEVELOCITY_TERMS_OF_USE_AGREE', config.getBuildScanTermsOfUseAgree())
     }
+    setupToken(
+        config.getDevelocityAccessKey(),
+        config.getDevelocityTokenExpiry(),
+        getEnv('DEVELOCITY_ENFORCE_URL'),
+        getEnv('DEVELOCITY_URL')
+    )
+}
+
+function getEnv(variableName: string): string | undefined {
+    return process.env[variableName]
 }
 
 function maybeExportVariable(variableName: string, value: unknown): void {

sources/src/develocity/short-lived-token.ts

@@ -0,0 +1,191 @@
+import * as httpm from 'typed-rest-client/HttpClient'
+import * as core from '@actions/core'
+
+export async function setupToken(
+    develocityAccessKey: string,
+    develocityTokenExpiry: string,
+    enforceUrl: string | undefined,
+    develocityUrl: string | undefined
+): Promise<void> {
+    const develocityAccesskeyEnvVar = 'DEVELOCITY_ACCESS_KEY'
+    if (develocityAccessKey) {
+        try {
+            core.debug('Fetching short-lived token...')
+            const tokens = await getToken(enforceUrl, develocityUrl, develocityAccessKey, develocityTokenExpiry)
+            if (tokens != null && !tokens.isEmpty()) {
+                core.debug(`Got token(s), setting the ${develocityAccesskeyEnvVar} env var`)
+                const token = tokens.raw()
+                core.setSecret(token)
+                core.exportVariable(develocityAccesskeyEnvVar, token)
+            } else {
+                // In case of not being able to generate a token we set the env variable to empty to avoid leaks
+                core.exportVariable(develocityAccesskeyEnvVar, '')
+            }
+        } catch (e) {
+            core.exportVariable(develocityAccesskeyEnvVar, '')
+            core.warning(`Failed to fetch short-lived token, reason: ${e}`)
+        }
+    }
+}
+
+export async function getToken(
+    enforceUrl: string | undefined,
+    serverUrl: string | undefined,
+    accessKey: string,
+    expiry: string
+): Promise<DevelocityAccessCredentials | null> {
+    const empty: Promise<DevelocityAccessCredentials | null> = new Promise(r => r(null))
+    const develocityAccessKey = DevelocityAccessCredentials.parse(accessKey)
+    const shortLivedTokenClient = new ShortLivedTokenClient()
+
+    async function promiseError(message: string): Promise<DevelocityAccessCredentials | null> {
+        return new Promise((resolve, reject) => reject(new Error(message)))
+    }
+
+    if (develocityAccessKey == null) {
+        return empty
+    }
+    if (enforceUrl === 'true' || develocityAccessKey.isSingleKey()) {
+        if (!serverUrl) {
+            return promiseError('Develocity Server URL not configured')
+        }
+        const hostname = extractHostname(serverUrl)
+        if (hostname == null) {
+            return promiseError('Could not extract hostname from Develocity server URL')
+        }
+        const hostAccessKey = develocityAccessKey.forHostname(hostname)
+        if (!hostAccessKey) {
+            return promiseError(`Could not find corresponding key for hostname ${hostname}`)
+        }
+        try {
+            const token = await shortLivedTokenClient.fetchToken(serverUrl, hostAccessKey, expiry)
+            return DevelocityAccessCredentials.of([token])
+        } catch (e) {
+            return new Promise((resolve, reject) => reject(e))
+        }
+    }
+
+    const tokens = new Array<HostnameAccessKey>()
+    for (const k of develocityAccessKey.keys) {
+        try {
+            const token = await shortLivedTokenClient.fetchToken(`https://${k.hostname}`, k, expiry)
+            tokens.push(token)
+        } catch (e) {
+            // Ignoring failed token, TODO: log this ?
+        }
+    }
+    if (tokens.length > 0) {
+        return DevelocityAccessCredentials.of(tokens)
+    }
+    return empty
+}
+
+function extractHostname(serverUrl: string): string | null {
+    try {
+        const parsedUrl = new URL(serverUrl)
+        return parsedUrl.hostname
+    } catch (error) {
+        return null
+    }
+}
+
+class ShortLivedTokenClient {
+    httpc = new httpm.HttpClient('gradle/setup-gradle')
+    maxRetries = 3
+    retryInterval = 1000
+
+    async fetchToken(serverUrl: string, accessKey: HostnameAccessKey, expiry: string): Promise<HostnameAccessKey> {
+        const queryParams = expiry ? `?expiresInHours${expiry}` : ''
+        const sanitizedServerUrl = !serverUrl.endsWith('/') ? `${serverUrl}/` : serverUrl
+        const headers = {
+            'Content-Type': 'application/json',
+            Authorization: `Bearer ${accessKey.key}`
+        }
+
+        let attempts = 0
+        while (attempts < this.maxRetries) {
+            try {
+                const requestUrl = `${sanitizedServerUrl}api/auth/token${queryParams}`
+                core.debug(`Attempt ${attempts} to fetch short lived token at ${requestUrl}`)
+                const response = await this.httpc.post(requestUrl, '', headers)
+                if (response.message.statusCode === 200) {
+                    const text = await response.readBody()
+                    return new Promise<HostnameAccessKey>(resolve => resolve({hostname: accessKey.hostname, key: text}))
+                }
+                // This should be only 404
+                attempts++
+                if (attempts === this.maxRetries) {
+                    return new Promise((resolve, reject) =>
+                        reject(
+                            new Error(
+                                `Develocity short lived token request failed ${serverUrl} with status code ${response.message.statusCode}`
+                            )
+                        )
+                    )
+                }
+            } catch (error) {
+                attempts++
+                if (attempts === this.maxRetries) {
+                    return new Promise((resolve, reject) => reject(error))
+                }
+            }
+            await new Promise(resolve => setTimeout(resolve, this.retryInterval))
+        }
+        return new Promise((resolve, reject) => reject(new Error('Illegal state')))
+    }
+}
+
+type HostnameAccessKey = {
+    hostname: string
+    key: string
+}
+
+export class DevelocityAccessCredentials {
+    static readonly accessKeyRegexp = /^(\S+=\w+)(;\S+=\w+)*$/
+    readonly keys: HostnameAccessKey[]
+
+    private constructor(allKeys: HostnameAccessKey[]) {
+        this.keys = allKeys
+    }
+
+    static of(allKeys: HostnameAccessKey[]): DevelocityAccessCredentials {
+        return new DevelocityAccessCredentials(allKeys)
+    }
+
+    private static readonly keyDelimiter = ';'
+    private static readonly hostDelimiter = '='
+
+    static parse(rawKey: string): DevelocityAccessCredentials | null {
+        if (!this.isValid(rawKey)) {
+            return null
+        }
+        return new DevelocityAccessCredentials(
+            rawKey.split(this.keyDelimiter).map(hostKey => {
+                const pair = hostKey.split(this.hostDelimiter)
+                return {hostname: pair[0], key: pair[1]}
+            })
+        )
+    }
+
+    isEmpty(): boolean {
+        return this.keys.length === 0
+    }
+
+    isSingleKey(): boolean {
+        return this.keys.length === 1
+    }
+
+    forHostname(hostname: string): HostnameAccessKey | undefined {
+        return this.keys.find(hostKey => hostKey.hostname === hostname)
+    }
+
+    raw(): string {
+        return this.keys
+            .map(k => `${k.hostname}${DevelocityAccessCredentials.hostDelimiter}${k.key}`)
+            .join(DevelocityAccessCredentials.keyDelimiter)
+    }
+
+    private static isValid(allKeys: string): boolean {
+        return this.accessKeyRegexp.test(allKeys)
+    }
+}

sources/src/setup-gradle.ts

@@ -4,7 +4,7 @@ import * as path from 'path'
 import * as os from 'os'
 import * as caches from './caching/caches'
 import * as jobSummary from './job-summary'
-import * as buildScan from './build-scan'
+import * as buildScan from './develocity/build-scan'
 
 import {loadBuildResults, markBuildResultsProcessed} from './build-results'
 import {CacheListener, generateCachingReport} from './caching/cache-reporting'
@@ -41,7 +41,7 @@ export async function setup(cacheConfig: CacheConfig, buildScanConfig: BuildScan
 
     core.saveState(CACHE_LISTENER, cacheListener.stringify())
 
-    buildScan.setup(buildScanConfig)
+    await buildScan.setup(buildScanConfig)
 
     return true
 }

sources/test/jest/short-lived-token.test.ts

@@ -0,0 +1,137 @@
+import {DevelocityAccessCredentials, getToken} from "../../src/develocity/short-lived-token";
+import nock from "nock";
+
+describe('short lived tokens', () => {
+    it('parse valid access key should return an object', async () => {
+        let develocityAccessCredentials = DevelocityAccessCredentials.parse('some-host.local=key1;host2=key2');
+
+        expect(develocityAccessCredentials).toStrictEqual(DevelocityAccessCredentials.of([
+            {hostname: 'some-host.local', key: 'key1'},
+            {hostname: 'host2', key: 'key2'}])
+        )
+    })
+
+    it('parse wrong access key should return null', async () => {
+        let develocityAccessCredentials = DevelocityAccessCredentials.parse('random;foo');
+
+        expect(develocityAccessCredentials).toBeNull()
+    })
+
+    it('parse empty access key should return null', async () => {
+        let develocityAccessCredentials = DevelocityAccessCredentials.parse('');
+
+        expect(develocityAccessCredentials).toBeNull()
+    })
+
+    it('access key as raw string', async () => {
+        let develocityAccessCredentials = DevelocityAccessCredentials.parse('host1=key1;host2=key2');
+
+        expect(develocityAccessCredentials?.raw()).toBe('host1=key1;host2=key2')
+    })
+
+    it('get short lived token fails when cannot connect', async () => {
+        nock('http://localhost:3333')
+            .post('/api/auth/token')
+            .times(3)
+            .replyWithError({
+                message: 'connect ECONNREFUSED 127.0.0.1:3333',
+                code: 'ECONNREFUSED'
+            })
+        try {
+            await getToken('true', 'http://localhost:3333', 'localhost=xyz;host1=key1', '')
+            expect('should have thrown').toBeUndefined()
+        } catch (e) {
+            // @ts-ignore
+            expect(e.code).toBe('ECONNREFUSED')
+        }
+    })
+
+    it('get short lived token fails when request fails', async () => {
+        nock('http://dev:3333')
+            .post('/api/auth/token')
+            .times(3)
+            .reply(500, 'Internal error')
+        expect.assertions(1)
+        await expect(getToken('true', 'http://dev:3333', 'dev=xyz;host1=key1', ''))
+            .rejects
+            .toThrow('Develocity short lived token request failed http://dev:3333 with status code 500')
+    })
+
+    it('get short lived token fails when server url is not set', async () => {
+        expect.assertions(1)
+        await expect(getToken('true', undefined, 'localhost=xyz;host1=key1', ''))
+            .rejects
+            .toThrow('Develocity Server URL not configured')
+    })
+
+    it('get short lived token returns null when access key is empty', async () => {
+        expect.assertions(1)
+        await expect(getToken('true', 'http://dev:3333', '', ''))
+            .resolves
+            .toBeNull()
+    })
+
+    it('get short lived token fails when host cannot be extracted from server url', async () => {
+        expect.assertions(1)
+        await expect(getToken('true', 'not_a_url', 'localhost=xyz;host1=key1', ''))
+            .rejects
+            .toThrow('Could not extract hostname from Develocity server URL')
+    })
+
+    it('get short lived token fails when access key does not contain corresponding host', async () => {
+        expect.assertions(1)
+        await expect(getToken('true', 'http://dev', 'host1=xyz;host2=key2', ''))
+            .rejects
+            .toThrow('Could not find corresponding key for hostname dev')
+    })
+
+    it('get short lived token succeeds when enforce url is true', async () => {
+        nock('https://dev')
+            .post('/api/auth/token')
+            .reply(200, 'token')
+        expect.assertions(1)
+        await expect(getToken('true', 'https://dev', 'dev=key1;host1=key2', ''))
+            .resolves
+            .toEqual({"keys": [{"hostname": "dev", "key": "token"}]})
+    })
+
+    it('get short lived token succeeds when enforce url is false and single key is set', async () => {
+        nock('https://dev')
+            .post('/api/auth/token')
+            .reply(200, 'token')
+        expect.assertions(1)
+        await expect(getToken('false', 'https://dev', 'dev=key1', ''))
+            .resolves
+            .toEqual({"keys": [{"hostname": "dev", "key": "token"}]})
+    })
+
+    it('get short lived token succeeds when enforce url is false and multiple keys are set', async () => {
+        nock('https://dev')
+            .post('/api/auth/token')
+            .reply(200, 'token1')
+        nock('https://prod')
+            .post('/api/auth/token')
+            .reply(200, 'token2')
+        expect.assertions(1)
+        await expect(getToken('false', 'https://dev', 'dev=key1;prod=key2', ''))
+            .resolves
+            .toEqual({"keys": [{"hostname": "dev", "key": "token1"}, {"hostname": "prod", "key": "token2"}]})
+    })
+
+    it('get short lived token succeeds when enforce url is false and multiple keys are set and one is failing', async () => {
+        nock('https://dev')
+            .post('/api/auth/token')
+            .reply(200, 'token1')
+        nock('https://bogus')
+            .post('/api/auth/token')
+            .times(3)
+            .reply(500, 'Internal Error')
+        nock('https://prod')
+            .post('/api/auth/token')
+            .reply(200, 'token2')
+        expect.assertions(1)
+        await expect(getToken('false', 'https://dev', 'dev=key1;bogus=key0;prod=key2', ''))
+            .resolves
+            .toEqual({"keys": [{"hostname": "dev", "key": "token1"}, {"hostname": "prod", "key": "token2"}]})
+    })
+})