From a8ffa2abaef45bf7e2b84abc6f59884c27020744 Mon Sep 17 00:00:00 2001 From: alalamav <2132122+alalamav@users.noreply.github.com> Date: Mon, 9 Dec 2019 13:05:15 -0500 Subject: [PATCH] Server data limit (#542) --- src/shadowbox/README.md | 10 +- src/shadowbox/model/access_key.ts | 22 +-- src/shadowbox/server/api.yml | 80 +++----- src/shadowbox/server/main.ts | 2 +- src/shadowbox/server/manager_service.spec.ts | 84 ++++----- src/shadowbox/server/manager_service.ts | 53 +++--- .../server/server_access_key.spec.ts | 174 +++++++----------- src/shadowbox/server/server_access_key.ts | 65 +++---- src/shadowbox/server/server_config.ts | 3 + 9 files changed, 201 insertions(+), 292 deletions(-) diff --git a/src/shadowbox/README.md b/src/shadowbox/README.md index 4f5a27c7..560074ea 100644 --- a/src/shadowbox/README.md +++ b/src/shadowbox/README.md @@ -119,15 +119,15 @@ Remove an access key curl --insecure -X DELETE $API_URL/access-keys/2 ``` -Set an access key data limit -(e.g. limit outbound data transfer for access key 2 to 1MB over 30 days) +Set a data limit for all access keys +(e.g. limit outbound data transfer access keys to 1MB over 30 days) ``` -curl -v --insecure -X PUT -H "Content-Type: application/json" -d '{"limit": {"bytes": 1000}}' $API_URL/access-keys/2/data-limit +curl -v --insecure -X PUT -H "Content-Type: application/json" -d '{"limit": {"bytes": 1000}}' $API_URL/experimental/access-key-data-limit ``` -Remove an access key data limit +Remove the access key data limit ``` -curl -v --insecure -X DELETE $API_URL/access-keys/2/data-limit +curl -v --insecure -X DELETE $API_URL/experimental/access-key-data-limit ``` ## Testing diff --git a/src/shadowbox/model/access_key.ts b/src/shadowbox/model/access_key.ts index 1db50663..8ebbb021 100644 --- a/src/shadowbox/model/access_key.ts +++ b/src/shadowbox/model/access_key.ts @@ -27,8 +27,10 @@ export interface ProxyParams { readonly password: string; } -// Data transfer measured in bytes. -export interface DataUsage { readonly bytes: number; } +// Data transfer allowance, measured in bytes. Must be a serializable JSON object. +export interface DataLimit { + readonly bytes: number; +} // AccessKey is what admins work with. It gives ProxyParams a name and identity. export interface AccessKey { @@ -40,12 +42,8 @@ export interface AccessKey { readonly metricsId: AccessKeyMetricsId; // Parameters to access the proxy readonly proxyParams: ProxyParams; - // Admin-controlled, data transfer limit for this access key. Unlimited if unset. - readonly dataLimit?: DataUsage; - // Data transferred by this access key over a 30 day sliding timeframe. - readonly dataUsage: DataUsage; - // Returns whether the access key has exceeded its data transfer limit. - isOverDataLimit(): boolean; + // Whether the access key has exceeded the data transfer limit. + readonly isOverDataLimit: boolean; } export interface AccessKeyRepository { @@ -61,8 +59,8 @@ export interface AccessKeyRepository { renameAccessKey(id: AccessKeyId, name: string): void; // Gets the metrics id for a given Access Key. getMetricsId(id: AccessKeyId): AccessKeyMetricsId|undefined; - // Sets the transfer limit for the specified access key. Throws on failure. - setAccessKeyDataLimit(id: AccessKeyId, limit: DataUsage): Promise; - // Clears the transfer limit for the specified access key. Throws on failure. - removeAccessKeyDataLimit(id: AccessKeyId): Promise; + // Sets a data transfer limit for all access keys. + setAccessKeyDataLimit(limit: DataLimit): Promise; + // Removes the access key data transfer limit. + removeAccessKeyDataLimit(): Promise; } diff --git a/src/shadowbox/server/api.yml b/src/shadowbox/server/api.yml index 840945fa..2ca72de5 100644 --- a/src/shadowbox/server/api.yml +++ b/src/shadowbox/server/api.yml @@ -165,7 +165,7 @@ paths: content: application/json: schema: - $ref: "#/components/schemas/DataUsage" + $ref: "#/components/schemas/DataLimit" examples: '0': value: "{bytes: 10000}" @@ -174,52 +174,6 @@ paths: description: Access key renamed successfully '404': description: Access key inexistent - /access-keys/{id}/data-limit: - put: - description: Sets an access key data transfer limit - tags: - - Access Key - - Limit - parameters: - - name: id - in: path - required: true - description: The id of the access key to set a limit - schema: - type: string - requestBody: - required: true - content: - application/json: - schema: - $ref: "#/components/schemas/DataUsage" - examples: - '0': - value: "{bytes: 10000}" - responses: - '204': - description: Access key limit set successfully - '400': - description: Invalid data limit - '404': - description: Access key inexistent - delete: - description: Removes an access key data transfer limit, lifting data transfer restrictions on the key. - tags: - - Access Key - - Limit - parameters: - - name: id - in: path - required: true - description: The id of the access key to delete a limit - schema: - type: string - responses: - '204': - description: Access key limit deleted successfully. - '404': - description: Access key inexistent /metrics/transfer: get: description: Returns the data transferred per access key @@ -279,6 +233,34 @@ paths: description: Setting successful '400': description: Invalid request + /experimental/access-key-data-limit: + put: + description: Sets a data transfer limit for all access keys + tags: + - Access Key + - Limit + requestBody: + required: true + content: + application/json: + schema: + $ref: "#/components/schemas/DataLimit" + examples: + '0': + value: "limit: {bytes: 10000}" + responses: + '204': + description: Access key data limit set successfully + '400': + description: Invalid data limit + delete: + description: Removes the access key data limit, lifting data transfer restrictions on all access keys. + tags: + - Access Key + - Limit + responses: + '204': + description: Access key limit deleted successfully. components: schemas: @@ -295,7 +277,7 @@ components: portForNewAccessKeys: type: integer - DataUsage: + DataLimit: properties: bytes: type: integer @@ -317,5 +299,3 @@ components: type: string accessUrl: type: string - limit: - $ref: "#/components/schemas/DataUsage" diff --git a/src/shadowbox/server/main.ts b/src/shadowbox/server/main.ts index a794e4cb..67d9acd6 100644 --- a/src/shadowbox/server/main.ts +++ b/src/shadowbox/server/main.ts @@ -144,7 +144,7 @@ async function main() { } const accessKeyRepository = new ServerAccessKeyRepository( serverConfig.data().portForNewAccessKeys, proxyHostname, accessKeyConfig, shadowsocksServer, - prometheusClient); + prometheusClient, serverConfig.data().accessKeyDataLimit); const metricsReader = new PrometheusUsageMetrics(prometheusClient); const toMetricsId = (id: AccessKeyId) => { diff --git a/src/shadowbox/server/manager_service.spec.ts b/src/shadowbox/server/manager_service.spec.ts index 7260b24b..f5a3b3a2 100644 --- a/src/shadowbox/server/manager_service.spec.ts +++ b/src/shadowbox/server/manager_service.spec.ts @@ -15,7 +15,7 @@ import * as net from 'net'; import {InMemoryConfig} from '../infrastructure/json_config'; -import {AccessKey, AccessKeyRepository, DataUsage} from '../model/access_key'; +import {AccessKey, AccessKeyRepository, DataLimit} from '../model/access_key'; import {ShadowsocksManagerService} from './manager_service'; import {FakePrometheusClient, FakeShadowsocksServer} from './mocks/mocks'; @@ -25,12 +25,13 @@ import {SharedMetricsPublisher} from './shared_metrics'; interface ServerInfo { name: string; + accessKeyDataLimit?: DataLimit; } const NEW_PORT = 12345; const OLD_PORT = 54321; const EXPECTED_ACCESS_KEY_PROPERTIES = - ['id', 'name', 'password', 'port', 'method', 'accessUrl', 'dataLimit'].sort(); + ['id', 'name', 'password', 'port', 'method', 'accessUrl'].sort(); describe('ShadowsocksManagerService', () => { // After processing the response callback, we should set @@ -59,15 +60,18 @@ describe('ShadowsocksManagerService', () => { }, done); }); - it('Return saved name', (done) => { + it('Returns persisted properties', (done) => { const repo = getAccessKeyRepository(); - const serverConfig = new InMemoryConfig({name: 'Server'} as ServerConfigJson); + const accessKeyDataLimit = {bytes: 999}; + const serverConfig = + new InMemoryConfig({name: 'Server', accessKeyDataLimit} as ServerConfigJson); const service = new ShadowsocksManagerService('default name', serverConfig, repo, null, null); service.getServer( {params: {}}, { send: (httpCode, data: ServerInfo) => { expect(httpCode).toEqual(200); expect(data.name).toEqual('Server'); + expect(data.accessKeyDataLimit).toEqual(accessKeyDataLimit); responseProcessed = true; } }, @@ -118,8 +122,6 @@ describe('ShadowsocksManagerService', () => { const service = new ShadowsocksManagerService('', null, repo, null, null); const accessKey = await repo.createNewAccessKey(); await repo.createNewAccessKey(); - const limit = {bytes: 10000}; - await repo.setAccessKeyDataLimit(accessKey.id, limit); const accessKeyName = 'new name'; await repo.renameAccessKey(accessKey.id, accessKeyName); const res = { @@ -131,7 +133,6 @@ describe('ShadowsocksManagerService', () => { expect(Object.keys(serviceAccessKey1).sort()).toEqual(EXPECTED_ACCESS_KEY_PROPERTIES); expect(Object.keys(serviceAccessKey2).sort()).toEqual(EXPECTED_ACCESS_KEY_PROPERTIES); expect(serviceAccessKey1.name).toEqual(accessKeyName); - expect(serviceAccessKey1.dataLimit).toEqual(limit); responseProcessed = true; // required for afterEach to pass. } }; @@ -384,29 +385,36 @@ describe('ShadowsocksManagerService', () => { describe('setAccessKeyDataLimit', () => { it('sets access key limit', async (done) => { + const serverConfig = new InMemoryConfig({} as ServerConfigJson); const repo = getAccessKeyRepository(); - const service = new ShadowsocksManagerService('default name', null, repo, null, null); - const accessKey = await repo.createNewAccessKey(); - expect(accessKey.dataLimit).toBeUndefined(); - expect(accessKey.isOverDataLimit()).toBeFalsy(); + spyOn(repo, 'setAccessKeyDataLimit'); + const service = new ShadowsocksManagerService('default name', serverConfig, repo, null, null); const limit = {bytes: 10000}; const res = { send: (httpCode, data) => { expect(httpCode).toEqual(204); - expect(accessKey.dataLimit).toEqual(limit); - expect(accessKey.isOverDataLimit()).toBeFalsy(); - responseProcessed = true; // required for afterEach to pass. + expect(serverConfig.data().accessKeyDataLimit).toEqual(limit); + expect(repo.setAccessKeyDataLimit).toHaveBeenCalledWith(limit); + service.getServer( + {params: {}}, { + send: (httpCode, data: ServerInfo) => { + expect(httpCode).toEqual(200); + expect(data.accessKeyDataLimit).toEqual(limit); + responseProcessed = true; // required for afterEach to pass. + } + }, + done); } }; - service.setAccessKeyDataLimit({params: {id: accessKey.id, limit}}, res, done); + service.setAccessKeyDataLimit({params: {limit}}, res, done); }); it('returns 400 when limit is missing values', async (done) => { const repo = getAccessKeyRepository(); const service = new ShadowsocksManagerService('default name', null, repo, null, null); const accessKey = await repo.createNewAccessKey(); - const limit = {} as DataUsage; + const limit = {} as DataLimit; const res = {send: (httpCode, data) => {}}; - service.setAccessKeyDataLimit({params: {id: accessKey.id, limit}}, res, (error) => { + service.setAccessKeyDataLimit({params: {limit}}, res, (error) => { expect(error.statusCode).toEqual(400); responseProcessed = true; // required for afterEach to pass. done(); @@ -418,23 +426,12 @@ describe('ShadowsocksManagerService', () => { const accessKey = await repo.createNewAccessKey(); const limit = {bytes: -1}; const res = {send: (httpCode, data) => {}}; - service.setAccessKeyDataLimit({params: {id: accessKey.id, limit}}, res, (error) => { + service.setAccessKeyDataLimit({params: {limit}}, res, (error) => { expect(error.statusCode).toEqual(400); responseProcessed = true; // required for afterEach to pass. done(); }); }); - it('returns 404 when the access key is not found', async (done) => { - const repo = getAccessKeyRepository(); - const service = new ShadowsocksManagerService('default name', null, repo, null, null); - const limit = {bytes: 1000}; - const res = {send: (httpCode, data) => {}}; - service.setAccessKeyDataLimit({params: {id: 'doesnotexist', limit}}, res, (error) => { - expect(error.statusCode).toEqual(404); - responseProcessed = true; // required for afterEach to pass. - done(); - }); - }); it('returns 500 when the repository throws an exception', async (done) => { const repo = getAccessKeyRepository(); spyOn(repo, 'setAccessKeyDataLimit').and.throwError('cannot write to disk'); @@ -442,7 +439,7 @@ describe('ShadowsocksManagerService', () => { const accessKey = await repo.createNewAccessKey(); const limit = {bytes: 10000}; const res = {send: (httpCode, data) => {}}; - service.setAccessKeyDataLimit({params: {id: accessKey.id, limit}}, res, (error) => { + service.setAccessKeyDataLimit({params: {limit}}, res, (error) => { expect(error.statusCode).toEqual(500); responseProcessed = true; // required for afterEach to pass. done(); @@ -452,32 +449,21 @@ describe('ShadowsocksManagerService', () => { describe('removeAccessKeyDataLimit', () => { it('clears access key limit', async (done) => { - const repo = getAccessKeyRepository(); - const service = new ShadowsocksManagerService('default name', null, repo, null, null); const limit = {bytes: 10000}; - const accessKey = await repo.createNewAccessKey(); - await repo.setAccessKeyDataLimit(accessKey.id, limit); - expect(accessKey.dataLimit).toEqual(limit); - expect(accessKey.dataUsage.bytes).toEqual(0); + const serverConfig = new InMemoryConfig({'accessKeyDataLimit': limit} as ServerConfigJson); + const repo = getAccessKeyRepository(); + spyOn(repo, 'removeAccessKeyDataLimit').and.callThrough(); + const service = new ShadowsocksManagerService('default name', serverConfig, repo, null, null); + await repo.setAccessKeyDataLimit(limit); const res = { send: (httpCode, data) => { expect(httpCode).toEqual(204); - expect(accessKey.dataLimit).toBeUndefined(); - expect(accessKey.isOverDataLimit()).toBeFalsy(); + expect(serverConfig.data().accessKeyDataLimit).toBeUndefined(); + expect(repo.removeAccessKeyDataLimit).toHaveBeenCalled(); responseProcessed = true; // required for afterEach to pass. } }; - service.removeAccessKeyDataLimit({params: {id: accessKey.id}}, res, done); - }); - it('returns 404 when the access key is not found', async (done) => { - const repo = getAccessKeyRepository(); - const service = new ShadowsocksManagerService('default name', null, repo, null, null); - const res = {send: (httpCode, data) => {}}; - service.removeAccessKeyDataLimit({params: {id: 'doesnotexist'}}, res, (error) => { - expect(error.statusCode).toEqual(404); - responseProcessed = true; // required for afterEach to pass. - done(); - }); + service.removeAccessKeyDataLimit({params: {}}, res, done); }); it('returns 500 when the repository throws an exception', async (done) => { const repo = getAccessKeyRepository(); diff --git a/src/shadowbox/server/manager_service.ts b/src/shadowbox/server/manager_service.ts index acf57280..4f090c03 100644 --- a/src/shadowbox/server/manager_service.ts +++ b/src/shadowbox/server/manager_service.ts @@ -14,12 +14,12 @@ import * as restify from 'restify'; import {makeConfig, SIP002_URI} from 'ShadowsocksConfig/shadowsocks_config'; -import {version} from '../package.json'; import {JsonConfig} from '../infrastructure/json_config'; import * as logging from '../infrastructure/logging'; -import {AccessKey, AccessKeyRepository, DataUsage} from '../model/access_key'; +import {AccessKey, AccessKeyRepository, DataLimit} from '../model/access_key'; import * as errors from '../model/errors'; +import {version} from '../package.json'; import {ManagerMetrics} from './manager_metrics'; import {ServerConfigJson} from './server_config'; @@ -42,8 +42,7 @@ function accessKeyToJson(accessKey: AccessKey) { method: accessKey.proxyParams.encryptionMethod, password: accessKey.proxyParams.password, outline: 1, - })), - dataLimit: accessKey.dataLimit + })) }; } @@ -53,7 +52,7 @@ interface RequestParams { // id: string // name: string // metricsEnabled: boolean - // limit: DataUsage + // limit: DataLimit // port: number // hours: number [param: string]: unknown; @@ -85,14 +84,18 @@ export function bindService( apiServer.del(`${apiPrefix}/access-keys/:id`, service.removeAccessKey.bind(service)); apiServer.put(`${apiPrefix}/access-keys/:id/name`, service.renameAccessKey.bind(service)); - apiServer.put( - `${apiPrefix}/access-keys/:id/data-limit`, service.setAccessKeyDataLimit.bind(service)); - apiServer.del( - `${apiPrefix}/access-keys/:id/data-limit`, service.removeAccessKeyDataLimit.bind(service)); apiServer.get(`${apiPrefix}/metrics/transfer`, service.getDataUsage.bind(service)); apiServer.get(`${apiPrefix}/metrics/enabled`, service.getShareMetrics.bind(service)); apiServer.put(`${apiPrefix}/metrics/enabled`, service.setShareMetrics.bind(service)); + + // Experimental APIs + apiServer.put( + `${apiPrefix}/experimental/access-key-data-limit`, + service.setAccessKeyDataLimit.bind(service)); + apiServer.del( + `${apiPrefix}/experimental/access-key-data-limit`, + service.removeAccessKeyDataLimit.bind(service)); } function validateAccessKeyId(accessKeyId: unknown): string { @@ -138,7 +141,8 @@ export class ShadowsocksManagerService { serverId: this.serverConfig.data().serverId, metricsEnabled: this.serverConfig.data().metricsEnabled || false, createdTimestampMs: this.serverConfig.data().createdTimestampMs, - version + version, + accessKeyDataLimit: this.serverConfig.data().accessKeyDataLimit }); next(); } @@ -248,26 +252,23 @@ export class ShadowsocksManagerService { public async setAccessKeyDataLimit(req: RequestType, res: ResponseType, next: restify.Next) { try { logging.debug(`setAccessKeyDataLimit request ${JSON.stringify(req.params)}`); - const accessKeyId = validateAccessKeyId(req.params.id); - const limit = req.params.limit as DataUsage; + const limit = req.params.limit as DataLimit; if (!limit) { return next( - new restify.MissingParameterError({statusCode: 400}, 'Parameter `limit` is missing')); + new restify.MissingParameterError({statusCode: 400}, 'Missing `limit` parameter')); } else if (!Number.isInteger(limit.bytes)) { - return next(new restify.InvalidArgumentError( - {statusCode: 400}, 'Parameter `limit.bytes` must be an integer')); + return next( + new restify.InvalidArgumentError({statusCode: 400}, '`limit` must be an integer')); } - await this.accessKeys.setAccessKeyDataLimit(accessKeyId, limit); + this.accessKeys.setAccessKeyDataLimit(limit); + this.serverConfig.data().accessKeyDataLimit = limit; + this.serverConfig.write(); res.send(HttpSuccess.NO_CONTENT); return next(); } catch (error) { logging.error(error); if (error instanceof errors.InvalidAccessKeyDataLimit) { return next(new restify.InvalidArgumentError({statusCode: 400}, error.message)); - } else if (error instanceof errors.AccessKeyNotFound) { - return next(new restify.NotFoundError(error.message)); - } else if (error instanceof restify.HttpError) { - return next(error); } return next(new restify.InternalServerError()); } @@ -276,24 +277,18 @@ export class ShadowsocksManagerService { public async removeAccessKeyDataLimit(req: RequestType, res: ResponseType, next: restify.Next) { try { logging.debug(`removeAccessKeyDataLimit request ${JSON.stringify(req.params)}`); - const accessKeyId = validateAccessKeyId(req.params.id); - await this.accessKeys.removeAccessKeyDataLimit(accessKeyId); + await this.accessKeys.removeAccessKeyDataLimit(); + delete this.serverConfig.data().accessKeyDataLimit; + this.serverConfig.write(); res.send(HttpSuccess.NO_CONTENT); return next(); } catch (error) { logging.error(error); - if (error instanceof errors.AccessKeyNotFound) { - return next(new restify.NotFoundError(error.message)); - } else if (error instanceof restify.HttpError) { - return next(error); - } return next(new restify.InternalServerError()); } } public async getDataUsage(req: RequestType, res: ResponseType, next: restify.Next) { - // TODO(alalama): use AccessKey.dataUsage to avoid querying Prometheus. Deprecate this call in - // the manager in favor of `GET /access-keys`. try { res.send(HttpSuccess.OK, await this.managerMetrics.getOutboundByteTransfer({hours: 30 * 24})); return next(); diff --git a/src/shadowbox/server/server_access_key.spec.ts b/src/shadowbox/server/server_access_key.spec.ts index 58a89133..cf2f4ee2 100644 --- a/src/shadowbox/server/server_access_key.spec.ts +++ b/src/shadowbox/server/server_access_key.spec.ts @@ -17,7 +17,7 @@ import * as net from 'net'; import {ManualClock} from '../infrastructure/clock'; import {PortProvider} from '../infrastructure/get_port'; import {InMemoryConfig} from '../infrastructure/json_config'; -import {AccessKeyRepository, DataUsage} from '../model/access_key'; +import {AccessKeyRepository, DataLimit} from '../model/access_key'; import * as errors from '../model/errors'; import {FakePrometheusClient, FakeShadowsocksServer} from './mocks/mocks'; @@ -38,11 +38,10 @@ describe('ServerAccessKeyRepository', () => { }); }); - it('Creates access keys without limit and under limit', async (done) => { + it('Creates access keys under limit', async (done) => { const repo = new RepoBuilder().build(); const accessKey = await repo.createNewAccessKey(); - expect(accessKey.dataLimit).toBeUndefined(); - expect(accessKey.isOverDataLimit()).toBeFalsy(); + expect(accessKey.isOverDataLimit).toBeFalsy(); done(); }); @@ -166,22 +165,11 @@ describe('ServerAccessKeyRepository', () => { }); }); - it('Can set access key data limit', async (done) => { + it('can set access key data limit', async (done) => { const repo = new RepoBuilder().build(); - const accessKey = await repo.createNewAccessKey(); const limit = {bytes: 5000}; - await expectNoAsyncThrow(repo.setAccessKeyDataLimit.bind(repo, accessKey.id, limit)); - expect(accessKey.dataLimit).toEqual(limit); - expect(accessKey.dataUsage.bytes).toEqual(0); - done(); - }); - - it('setAccessKeyDataLimit throws for missing keys', async (done) => { - const repo = new RepoBuilder().build(); - await repo.createNewAccessKey(); - const limit = {bytes: 1000}; - await expectAsyncThrow( - repo.setAccessKeyDataLimit.bind(repo, 'doesnotexist', limit), errors.AccessKeyNotFound); + await expectNoAsyncThrow(repo.setAccessKeyDataLimit.bind(repo, limit)); + expect(repo.dataLimit).toEqual(limit); done(); }); @@ -191,17 +179,15 @@ describe('ServerAccessKeyRepository', () => { // Negative values const negativeBytesLimit = {bytes: -1000}; await expectAsyncThrow( - repo.setAccessKeyDataLimit.bind(repo, accessKey.id, negativeBytesLimit), + repo.setAccessKeyDataLimit.bind(repo, negativeBytesLimit), errors.InvalidAccessKeyDataLimit); // Missing properties - const missingDataLimit = {} as DataUsage; + const missingDataLimit = {} as DataLimit; await expectAsyncThrow( - repo.setAccessKeyDataLimit.bind(repo, accessKey.id, missingDataLimit), - errors.InvalidAccessKeyDataLimit); + repo.setAccessKeyDataLimit.bind(repo, missingDataLimit), errors.InvalidAccessKeyDataLimit); // Undefined limit await expectAsyncThrow( - repo.setAccessKeyDataLimit.bind(repo, accessKey.id, undefined), - errors.InvalidAccessKeyDataLimit); + repo.setAccessKeyDataLimit.bind(repo, undefined), errors.InvalidAccessKeyDataLimit); done(); }); @@ -214,21 +200,20 @@ describe('ServerAccessKeyRepository', () => { const accessKey2 = await repo.createNewAccessKey(); await repo.start(new ManualClock()); - await repo.setAccessKeyDataLimit(accessKey1.id, {bytes: 200}); - expect(accessKey1.isOverDataLimit()).toBeTruthy(); - expect(accessKey2.isOverDataLimit()).toBeFalsy(); + await repo.setAccessKeyDataLimit({bytes: 250}); + expect(accessKey1.isOverDataLimit).toBeTruthy(); + expect(accessKey2.isOverDataLimit).toBeFalsy(); // We determine which access keys have been enabled/disabled by accessing them from // the server's perspective, ensuring `server.update` has been called. let serverAccessKeys = server.getAccessKeys(); expect(serverAccessKeys.length).toEqual(1); expect(serverAccessKeys[0].id).toEqual(accessKey2.id); - // The over-limit access key should be re-enabled after increasing its limit, while the - // under-limit key should be disabled after setting its limit. - prometheusClient.bytesTransferredById = {'0': 800, '1': 199}; - await repo.setAccessKeyDataLimit(accessKey1.id, {bytes: 1000}); - await repo.setAccessKeyDataLimit(accessKey2.id, {bytes: 100}); - expect(accessKey1.isOverDataLimit()).toBeFalsy(); - expect(accessKey2.isOverDataLimit()).toBeTruthy(); + // The over-limit key should be re-enabled after increasing the data limit, while the other key + // should be disabled after its data usage increased. + prometheusClient.bytesTransferredById = {'0': 500, '1': 1000}; + await repo.setAccessKeyDataLimit({bytes: 700}); + expect(accessKey1.isOverDataLimit).toBeFalsy(); + expect(accessKey2.isOverDataLimit).toBeTruthy(); serverAccessKeys = server.getAccessKeys(); expect(serverAccessKeys.length).toEqual(1); expect(serverAccessKeys[0].id).toEqual(accessKey1.id); @@ -236,81 +221,57 @@ describe('ServerAccessKeyRepository', () => { }); it('can remove access key limits', async (done) => { - const repo = new RepoBuilder().build(); - const accessKey = await repo.createNewAccessKey(); const limit = {bytes: 100}; - await repo.setAccessKeyDataLimit(accessKey.id, limit); - expect(accessKey.dataLimit).toBeDefined(); - await expectNoAsyncThrow(repo.removeAccessKeyDataLimit.bind(repo, accessKey.id)); - expect(accessKey.dataLimit).toBeUndefined(); - done(); - }); - - it('removeAccessKeyDataLimit throws for missing keys', async (done) => { - const repo = new RepoBuilder().build(); - await repo.createNewAccessKey(); - await expectAsyncThrow( - repo.removeAccessKeyDataLimit.bind(repo, 'doesnotexist'), errors.AccessKeyNotFound); + const repo = new RepoBuilder().accessKeyDataLimit(limit).build(); + expect(repo.dataLimit).toEqual(limit); + await expectNoAsyncThrow(repo.removeAccessKeyDataLimit.bind(repo)); + expect(repo.dataLimit).toBeUndefined(); done(); }); it('removeAccessKeyDataLimit restores over-limit access keys', async (done) => { const server = new FakeShadowsocksServer(); const prometheusClient = new FakePrometheusClient({'0': 500, '1': 100}); - const repo = - new RepoBuilder().prometheusClient(prometheusClient).shadowsocksServer(server).build(); + const repo = new RepoBuilder() + .prometheusClient(prometheusClient) + .shadowsocksServer(server) + .accessKeyDataLimit({bytes: 200}) + .build(); const accessKey1 = await repo.createNewAccessKey(); const accessKey2 = await repo.createNewAccessKey(); await repo.start(new ManualClock()); - await repo.setAccessKeyDataLimit(accessKey1.id, {bytes: 100}); expect(server.getAccessKeys().length).toEqual(1); // Remove the limit; expect the key to be under limit and enabled. - await expectNoAsyncThrow(repo.removeAccessKeyDataLimit.bind(repo, accessKey1.id)); + await expectNoAsyncThrow(repo.removeAccessKeyDataLimit.bind(repo)); expect(server.getAccessKeys().length).toEqual(2); - expect(accessKey1.isOverDataLimit()).toBeFalsy(); - expect(accessKey2.isOverDataLimit()).toBeFalsy(); - expect(accessKey1.dataLimit).toBeUndefined(); - expect(accessKey2.dataLimit).toBeUndefined(); + expect(accessKey1.isOverDataLimit).toBeFalsy(); + expect(accessKey2.isOverDataLimit).toBeFalsy(); done(); }); it('enforceAccessKeyDataLimits updates keys limit status', async (done) => { - const NUM_ACCESS_KEYS = 20; - const bytesTransferredById = {}; - for (let i = 0; i < NUM_ACCESS_KEYS; ++i) { - bytesTransferredById[`${i}`] = i * 1000; - } - const prometheusClient = new FakePrometheusClient(bytesTransferredById); - const repo = new RepoBuilder().prometheusClient(prometheusClient).build(); - for (let i = 0; i < NUM_ACCESS_KEYS; ++i) { - const key = await repo.createNewAccessKey(); - if (i % 2 !== 0) { - // Set a limit on half of the keys. - await repo.setAccessKeyDataLimit(key.id, {bytes: i * 100}); - } + const prometheusClient = + new FakePrometheusClient({'0': 100, '1': 200, '2': 300, '3': 400, '4': 500}); + const limit = {bytes: 250}; + const repo = + new RepoBuilder().prometheusClient(prometheusClient).accessKeyDataLimit(limit).build(); + for (let i = 0; i < Object.keys(prometheusClient.bytesTransferredById).length; ++i) { + await repo.createNewAccessKey(); } await repo.enforceAccessKeyDataLimits(); for (const key of repo.listAccessKeys()) { - const hasDataLimit = !!key.dataLimit; - // Keys with data limits should be over the limit; keys without a limit shouldn't. - expect(key.isOverDataLimit()).toEqual(hasDataLimit); - if (hasDataLimit) { - expect(key.dataUsage.bytes).toEqual(bytesTransferredById[key.id]); - } + expect(key.isOverDataLimit) + .toEqual(prometheusClient.bytesTransferredById[key.id] > limit.bytes); } // Simulate a change in usage. - for (let i = 0; i < NUM_ACCESS_KEYS; ++i) { - bytesTransferredById[`${i}`] = i; - } - prometheusClient.bytesTransferredById = bytesTransferredById; + prometheusClient.bytesTransferredById = {'0': 500, '1': 400, '2': 300, '3': 200, '4': 100}; await repo.enforceAccessKeyDataLimits(); for (const key of repo.listAccessKeys()) { - // All keys should be under the data limit. - expect(key.isOverDataLimit()).toBeFalsy(); - expect(key.dataUsage.bytes).toEqual(bytesTransferredById[key.id]); + expect(key.isOverDataLimit) + .toEqual(prometheusClient.bytesTransferredById[key.id] > limit.bytes); } done(); }); @@ -318,12 +279,14 @@ describe('ServerAccessKeyRepository', () => { it('enforceAccessKeyDataLimits enables and disables keys', async (done) => { const server = new FakeShadowsocksServer(); const prometheusClient = new FakePrometheusClient({'0': 500, '1': 100}); - const repo = - new RepoBuilder().prometheusClient(prometheusClient).shadowsocksServer(server).build(); + const repo = new RepoBuilder() + .prometheusClient(prometheusClient) + .shadowsocksServer(server) + .accessKeyDataLimit({bytes: 200}) + .build(); const accessKey1 = await repo.createNewAccessKey(); const accessKey2 = await repo.createNewAccessKey(); - await repo.setAccessKeyDataLimit(accessKey1.id, {bytes: 200}); await repo.enforceAccessKeyDataLimits(); const accessKeys = await repo.listAccessKeys(); @@ -344,7 +307,6 @@ describe('ServerAccessKeyRepository', () => { // Create 2 new access keys await Promise.all([repo1.createNewAccessKey(), repo1.createNewAccessKey()]); // Modify properties - await repo1.setAccessKeyDataLimit('0', {bytes: 100}); repo1.renameAccessKey('1', 'name'); // Create a 2nd repo from the same config file. This simulates what @@ -393,40 +355,34 @@ describe('ServerAccessKeyRepository', () => { it('start periodically enforces access key data limits', async (done) => { const server = new FakeShadowsocksServer(); - const prometheusClient = new FakePrometheusClient({'0': 500, '1': 300, '2': 400}); + const prometheusClient = new FakePrometheusClient({'0': 500, '1': 200, '2': 400}); const repo = new RepoBuilder().prometheusClient(prometheusClient).shadowsocksServer(server).build(); const accessKey1 = await repo.createNewAccessKey(); const accessKey2 = await repo.createNewAccessKey(); const accessKey3 = await repo.createNewAccessKey(); - await repo.setAccessKeyDataLimit(accessKey1.id, {bytes: 300}); - await repo.setAccessKeyDataLimit(accessKey2.id, {bytes: 100}); + await repo.setAccessKeyDataLimit({bytes: 300}); const clock = new ManualClock(); - await repo.start(clock); await clock.runCallbacks(); - expect(accessKey1.isOverDataLimit()).toBeTruthy(); - expect(accessKey2.isOverDataLimit()).toBeTruthy(); - expect(accessKey3.isOverDataLimit()).toBeFalsy(); - expect(accessKey1.dataUsage.bytes).toEqual(500); - expect(accessKey2.dataUsage.bytes).toEqual(300); - expect(accessKey3.dataLimit).toBeUndefined(); + + expect(accessKey1.isOverDataLimit).toBeTruthy(); + expect(accessKey2.isOverDataLimit).toBeFalsy(); + expect(accessKey3.isOverDataLimit).toBeTruthy(); let serverAccessKeys = await server.getAccessKeys(); expect(serverAccessKeys.length).toEqual(1); - expect(serverAccessKeys[0].id).toEqual(accessKey3.id); + expect(serverAccessKeys[0].id).toEqual(accessKey2.id); + // Simulate a change in usage. - prometheusClient.bytesTransferredById = {'0': 100, '1': 300, '2': 1000}; + prometheusClient.bytesTransferredById = {'0': 100, '1': 200, '2': 1000}; await clock.runCallbacks(); - expect(accessKey1.isOverDataLimit()).toBeFalsy(); - expect(accessKey2.isOverDataLimit()).toBeTruthy(); - expect(accessKey3.isOverDataLimit()).toBeFalsy(); - expect(accessKey1.dataUsage.bytes).toEqual(100); - expect(accessKey2.dataUsage.bytes).toEqual(300); - expect(accessKey3.dataLimit).toBeUndefined(); + expect(accessKey1.isOverDataLimit).toBeFalsy(); + expect(accessKey2.isOverDataLimit).toBeFalsy(); + expect(accessKey3.isOverDataLimit).toBeTruthy(); serverAccessKeys = await server.getAccessKeys(); expect(serverAccessKeys.length).toEqual(2); expect(serverAccessKeys[0].id).toEqual(accessKey1.id); - expect(serverAccessKeys[1].id).toEqual(accessKey3.id); + expect(serverAccessKeys[1].id).toEqual(accessKey2.id); done(); }); }); @@ -464,6 +420,7 @@ class RepoBuilder { private keyConfig_ = new InMemoryConfig({accessKeys: [], nextId: 0}); private shadowsocksServer_ = new FakeShadowsocksServer(); private prometheusClient_ = new FakePrometheusClient({}); + private accessKeyDataLimit_; public port(port: number): RepoBuilder { this.port_ = port; @@ -481,9 +438,14 @@ class RepoBuilder { this.prometheusClient_ = prometheusClient; return this; } + public accessKeyDataLimit(limit: DataLimit): RepoBuilder { + this.accessKeyDataLimit_ = limit; + return this; + } public build(): ServerAccessKeyRepository { return new ServerAccessKeyRepository( - this.port_, 'hostname', this.keyConfig_, this.shadowsocksServer_, this.prometheusClient_); + this.port_, 'hostname', this.keyConfig_, this.shadowsocksServer_, this.prometheusClient_, + this.accessKeyDataLimit_); } } diff --git a/src/shadowbox/server/server_access_key.ts b/src/shadowbox/server/server_access_key.ts index 06c27eb3..38cd2a89 100644 --- a/src/shadowbox/server/server_access_key.ts +++ b/src/shadowbox/server/server_access_key.ts @@ -20,7 +20,7 @@ import {isPortUsed} from '../infrastructure/get_port'; import {JsonConfig} from '../infrastructure/json_config'; import * as logging from '../infrastructure/logging'; import {PrometheusClient} from '../infrastructure/prometheus_scraper'; -import {AccessKey, AccessKeyId, AccessKeyMetricsId, AccessKeyRepository, DataUsage, ProxyParams} from '../model/access_key'; +import {AccessKey, AccessKeyId, AccessKeyMetricsId, AccessKeyRepository, DataLimit, ProxyParams} from '../model/access_key'; import * as errors from '../model/errors'; import {ShadowsocksServer} from '../model/shadowsocks_server'; import {PrometheusManagerMetrics} from './manager_metrics'; @@ -33,7 +33,6 @@ interface AccessKeyJson { password: string; port: number; encryptionMethod?: string; - dataLimit?: DataUsage; } // The configuration file format as json. @@ -45,20 +44,13 @@ export interface AccessKeyConfigJson { // AccessKey implementation with write access enabled on properties that may change. class ServerAccessKey implements AccessKey { - public dataUsage: DataUsage = {bytes: 0}; + public isOverDataLimit = false; constructor( readonly id: AccessKeyId, public name: string, public metricsId: AccessKeyMetricsId, - readonly proxyParams: ProxyParams, public dataLimit?: DataUsage) {} - - isOverDataLimit(): boolean { - if (!this.dataLimit) { - return false; - } - return this.dataUsage.bytes > this.dataLimit.bytes; - } + readonly proxyParams: ProxyParams) {} } -function isValidAccessKeyDataLimit(limit: DataUsage): boolean { +function isValidAccessKeyDataLimit(limit: DataLimit): boolean { return limit && limit.bytes >= 0; } @@ -75,8 +67,7 @@ function makeAccessKey(hostname: string, accessKeyJson: AccessKeyJson): AccessKe password: accessKeyJson.password, }; return new ServerAccessKey( - accessKeyJson.id, accessKeyJson.name, accessKeyJson.metricsId, proxyParams, - accessKeyJson.dataLimit); + accessKeyJson.id, accessKeyJson.name, accessKeyJson.metricsId, proxyParams); } function makeAccessKeyJson(accessKey: AccessKey): AccessKeyJson { @@ -86,8 +77,7 @@ function makeAccessKeyJson(accessKey: AccessKey): AccessKeyJson { name: accessKey.name, password: accessKey.proxyParams.password, port: accessKey.proxyParams.portNumber, - encryptionMethod: accessKey.proxyParams.encryptionMethod, - dataLimit: accessKey.dataLimit + encryptionMethod: accessKey.proxyParams.encryptionMethod }; } @@ -102,7 +92,8 @@ export class ServerAccessKeyRepository implements AccessKeyRepository { constructor( private portForNewAccessKeys: number, private proxyHostname: string, private keyConfig: JsonConfig, - private shadowsocksServer: ShadowsocksServer, private prometheusClient: PrometheusClient) { + private shadowsocksServer: ShadowsocksServer, private prometheusClient: PrometheusClient, + private accessKeyDataLimit?: DataLimit) { if (this.keyConfig.data().accessKeys === undefined) { this.keyConfig.data().accessKeys = []; } @@ -155,7 +146,7 @@ export class ServerAccessKeyRepository implements AccessKeyRepository { encryptionMethod: this.NEW_USER_ENCRYPTION_METHOD, password, }; - const accessKey = new ServerAccessKey(id, '', metricsId, proxyParams, undefined); + const accessKey = new ServerAccessKey(id, '', metricsId, proxyParams); this.accessKeys.push(accessKey); this.saveAccessKeys(); await this.updateServer(); @@ -185,29 +176,21 @@ export class ServerAccessKeyRepository implements AccessKeyRepository { this.saveAccessKeys(); } - setAccessKeyDataLimit(id: AccessKeyId, limit: DataUsage): Promise { + get dataLimit(): DataLimit|undefined { + return this.accessKeyDataLimit; + } + + setAccessKeyDataLimit(limit: DataLimit): Promise { if (!isValidAccessKeyDataLimit(limit)) { throw new errors.InvalidAccessKeyDataLimit(); } - const accessKey = this.getAccessKey(id); - const wasOverDataLimit = accessKey.isOverDataLimit(); - accessKey.dataLimit = limit; - this.saveAccessKeys(); - if (accessKey.isOverDataLimit() !== wasOverDataLimit) { - return this.updateServer(); - } - return Promise.resolve(); + this.accessKeyDataLimit = limit; + return this.enforceAccessKeyDataLimits(); } - removeAccessKeyDataLimit(id: AccessKeyId): Promise { - const accessKey = this.getAccessKey(id); - const wasOverDataLimit = accessKey.isOverDataLimit(); - delete accessKey.dataLimit; - this.saveAccessKeys(); - if (wasOverDataLimit) { - return this.updateServer(); - } - return Promise.resolve(); + removeAccessKeyDataLimit(): Promise { + delete this.accessKeyDataLimit; + return this.enforceAccessKeyDataLimits(); } getMetricsId(id: AccessKeyId): AccessKeyMetricsId|undefined { @@ -223,9 +206,11 @@ export class ServerAccessKeyRepository implements AccessKeyRepository { (await metrics.getOutboundByteTransfer({hours: 30 * 24})).bytesTransferredByUserId; let limitStatusChanged = false; for (const accessKey of this.accessKeys) { - const wasOverDataLimit = accessKey.isOverDataLimit(); - accessKey.dataUsage = {bytes: bytesTransferredById[accessKey.id] || 0}; - limitStatusChanged = accessKey.isOverDataLimit() !== wasOverDataLimit || limitStatusChanged; + const usageBytes = bytesTransferredById[accessKey.id] || 0; + const wasOverDataLimit = accessKey.isOverDataLimit; + accessKey.isOverDataLimit = + this.accessKeyDataLimit ? usageBytes > this.accessKeyDataLimit.bytes : false; + limitStatusChanged = accessKey.isOverDataLimit !== wasOverDataLimit || limitStatusChanged; } if (limitStatusChanged) { await this.updateServer(); @@ -233,7 +218,7 @@ export class ServerAccessKeyRepository implements AccessKeyRepository { } private updateServer(): Promise { - const serverAccessKeys = this.accessKeys.filter(key => !key.isOverDataLimit()).map(key => { + const serverAccessKeys = this.accessKeys.filter(key => !key.isOverDataLimit).map(key => { return { id: key.id, port: key.proxyParams.portNumber, diff --git a/src/shadowbox/server/server_config.ts b/src/shadowbox/server/server_config.ts index 427de8ad..6cbd58b1 100644 --- a/src/shadowbox/server/server_config.ts +++ b/src/shadowbox/server/server_config.ts @@ -15,6 +15,7 @@ import * as uuidv4 from 'uuid/v4'; import * as json_config from '../infrastructure/json_config'; +import {DataLimit} from '../model/access_key'; // Serialized format for the server config. // WARNING: Renaming fields will break backwards-compatibility. @@ -34,6 +35,8 @@ export interface ServerConfigJson { // We don't serialize the shadowbox version, this is obtained dynamically from node. // Public proxy hostname. hostname?: string; + // Data transfer limit applied to all access keys. + accessKeyDataLimit?: DataLimit; } // Serialized format for rollouts.