import * as AWS from "aws-sdk" ; import { freeTierHardLimit , freeTierSec , freeTierSoftLimit } from "./limits" ; import { customAlphabet } from "nanoid/async" ; const dynamoDb = new AWS . DynamoDB . DocumentClient ( ) ; const NETWORKING_TABLE = process . env . NETWORKING_TABLE as string ; const lambda = new AWS . Lambda ( ) ; const RunInstanceLambda = process . env . RUN_INSTANCE as string ; const nanoid = customAlphabet ( "0123456789abcdefghijklmnopqrstuvwxyz" , 6 ) ; interface Token { day : number ; region : string ; endTime : number ; ttlSec : number ; } export type TaskType = "ipx" ; export async function getFreeTierSoftCount ( namespace , id , day ) { const freeTierSoftKey = getFreeTierSoftKey ( namespace , id , day ) ; const response = await dynamoDb . get ( { TableName : NETWORKING_TABLE , Key : { key : freeTierSoftKey , } , } ) . promise ( ) ; return response . Item ?. count || 0 ; } export async function getFreeTierHardCount ( day ) { const freeTierHardKey = getFreeTierHardKey ( day ) ; const response = await dynamoDb . get ( { TableName : NETWORKING_TABLE , Key : { key : freeTierHardKey , } , } ) . promise ( ) ; return response . Item ?. count || 0 ; } export async function createFreeToken ( namespace , id , day , region ) { const freeTierSoftKey = getFreeTierSoftKey ( namespace , id , day ) ; const freeTierHardKey = getFreeTierHardKey ( day ) ; const inc = async ( key , inc ) => { const updateParams = { TableName : NETWORKING_TABLE , Key : { key , } , UpdateExpression : "ADD #count :inc" , ExpressionAttributeNames : { "#count" : "count" , } , ExpressionAttributeValues : { ":inc" : inc , } , ReturnValues : "ALL_NEW" , } ; const response = await dynamoDb . update ( updateParams ) . promise ( ) ; if ( typeof response . Attributes ?. count !== "number" ) { console . error ( "Can't get updated count value for key " + key , response ) ; throw new Error ( "Can't get updated count value for key " + key ) ; } return response . Attributes ?. count as number ; } ; const [ softCount , hardCount ] = await Promise . all ( [ inc ( freeTierSoftKey , 1 ) , inc ( freeTierHardKey , 1 ) ] ) ; try { if ( softCount > freeTierSoftLimit ) { throw new Error ( "free-soft-limit" ) ; } if ( hardCount > freeTierHardLimit ) { throw new Error ( "free-hard-limit" ) ; } } catch ( e ) { await Promise . all ( [ inc ( freeTierSoftKey , - 1 ) , inc ( freeTierHardKey , - 1 ) ] ) ; throw e ; } return createToken ( namespace , id , day , region , freeTierSec ) ; } export async function getToken ( token ) { const item = ( await dynamoDb . get ( { TableName : NETWORKING_TABLE , Key : { key : getTokenKey ( token ) , } , } ) . promise ( ) ) . Item as Token | undefined ; if ( item === undefined ) { throw new Error ( "not-found" ) ; } item . ttlSec = Math . round ( ( item . endTime - new Date ( ) . getTime ( ) ) / 1000 ) ; return item ; } export async function createToken ( namespace , id , day , region , durationSec ) { let token = null ; while ( token === null ) { token = await nanoid ( ) ; const createParams = { TableName : NETWORKING_TABLE , Item : { "key" : getTokenKey ( token ) , namespace , id , day , region , durationSec , "createdAt" : new Date ( ) . getTime ( ) , "endTime" : new Date ( ) . getTime ( ) + durationSec * 1000 , } , ConditionExpression : "attribute_not_exists(#key)" , ExpressionAttributeNames : { "#key" : "key" , } , } ; try { await dynamoDb . put ( createParams ) . promise ( ) ; } catch ( e ) { console . error ( e ) ; token = null ; } } return token ; } export async function newTask ( tokenId , task ) { const token = await getToken ( tokenId ) ; if ( token . ttlSec < 60 ) { throw new Error ( "ended" ) ; } const result = await lambda . invoke ( { FunctionName : RunInstanceLambda , Payload : JSON . stringify ( { token : tokenId , region : token . region , task , } ) , } ) . promise ( ) ; if ( ! result . Payload ) { throw new Error ( "unable-to-start-task" ) ; } const payload = JSON . parse ( JSON . parse ( result . Payload as string ) . body ) ; if ( payload . success === false ) { throw new Error ( payload . errorCode ) ; } const arn = payload . arn ; if ( arn === undefined ) { throw new Error ( "ARN missed in runInstanceV2 response" ) ; } const updateParams = { TableName : NETWORKING_TABLE , Key : { key : getTokenKey ( tokenId ) , } , UpdateExpression : `SET ${ task }Arn = :arn` , ExpressionAttributeValues : { ":arn" : arn , } , } ; await dynamoDb . update ( updateParams ) . promise ( ) ; return arn ; } export async function putIp ( token , arn , ip , task ) { const updateParams = { TableName : NETWORKING_TABLE , Key : { key : getTokenKey ( token ) , } , UpdateExpression : `SET ${ task }Ip = :ip` , ConditionExpression : `${ task }Arn = :arn` , ExpressionAttributeValues : { ":arn" : arn , ":ip" : ip , } , } ; await dynamoDb . update ( updateParams ) . promise ( ) ; } export async function stopTask ( token , arn , task ) { const updateParams = { TableName : NETWORKING_TABLE , Key : { key : getTokenKey ( token ) , } , UpdateExpression : `REMOVE ${ task }Arn, ${ task }Ip` , ConditionExpression : `${ task }Arn = :arn` , ExpressionAttributeValues : { ":arn" : arn , } , } ; await dynamoDb . update ( updateParams ) . promise ( ) ; } function getFreeTierSoftKey ( namespace , id , day ) { return "free-soft " + day + " " + id + "@" + namespace ; } function getFreeTierHardKey ( day ) { return "free-hard " + day ; } function getTokenKey ( token ) { return "token " + token ; }