Server data limit (#542)

This commit is contained in:
alalamav 2019-12-09 13:05:15 -05:00 • committed by GitHub
parent c09bb40e57
commit a8ffa2abae
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
9 changed files with 201 additions and 292 deletions

View file

@ -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

View file

@ -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<void>;
// Clears the transfer limit for the specified access key. Throws on failure.
removeAccessKeyDataLimit(id: AccessKeyId): Promise<void>;
// Sets a data transfer limit for all access keys.
setAccessKeyDataLimit(limit: DataLimit): Promise<void>;
// Removes the access key data transfer limit.
removeAccessKeyDataLimit(): Promise<void>;
}

View file

@ -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"

View file

@ -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) => {

View file

@ -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();

View file

@ -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();

View file

@ -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<AccessKeyConfigJson>({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_);
}
}

View file

@ -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<AccessKeyConfigJson>,
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<void> {
get dataLimit(): DataLimit|undefined {
return this.accessKeyDataLimit;
}
setAccessKeyDataLimit(limit: DataLimit): Promise<void> {
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<void> {
const accessKey = this.getAccessKey(id);
const wasOverDataLimit = accessKey.isOverDataLimit();
delete accessKey.dataLimit;
this.saveAccessKeys();
if (wasOverDataLimit) {
return this.updateServer();
}
return Promise.resolve();
removeAccessKeyDataLimit(): Promise<void> {
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<void> {
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,

View file

@ -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.