import { AsyncQueue , messages , PreparedObjectType } from '@pgtyped/wire' ; import crypto from 'crypto' ; import debugBase from 'debug' ; import * as tls from 'tls' ; import { IInterpolatedQuery , QueryParam } from './preprocessor' ; import { checkServerFinalMessage , createClientSASLContinueResponse , createInitialSASLResponse , } from './sasl-helpers' ; import { DatabaseTypeKind , isEnum , MappableType } from './type' ; const debugQuery = debugBase ( 'client:query' ) ; export const generateHash = ( username , password , salt ) => { const hash = ( str ) => crypto . createHash ( 'md5' ) . update ( str ) . digest ( 'hex' ) ; const shadow = hash ( password + username ) ; const result = crypto . createHash ( 'md5' ) ; result . update ( shadow ) ; result . update ( salt ) ; return 'md5' + result . digest ( 'hex' ) ; } ; export async function startup ( options , queue ) { try { await queue . connect ( options ) ; const startupParams = { user : options . user , database : options . dbName , client_encoding : "'utf-8'" , } ; await queue . send ( messages . startupMessage , { params : startupParams } ) ; const result = await queue . reply ( messages . readyForQuery , messages . authenticationCleartextPassword , messages . authenticationMD5Password , messages . authenticationSASL ) ; if ( 'trxStatus' in result ) { return ; } if ( ! options . password ) { throw new Error ( 'password required for hash auth' ) ; } let password = options . password ; if ( 'SASLMechanisms' in result ) { if ( result . SASLMechanisms ?. indexOf ( 'SCRAM-SHA-256' ) === - 1 ) { throw new Error ( 'SASL: Only mechanism SCRAM-SHA-256 is currently supported' ) ; } const { clientNonce , response : initialSASLResponse } = createInitialSASLResponse ( ) ; await queue . send ( messages . SASLInitialResponse , { mechanism : 'SCRAM-SHA-256' , responseLength : Buffer . byteLength ( initialSASLResponse ) , response : initialSASLResponse , } ) ; const SASLContinueResult = await queue . reply ( messages . AuthenticationSASLContinue ) ; const { response : SASLContinueResponse , calculatedServerSignature } = createClientSASLContinueResponse ( password , clientNonce , SASLContinueResult . SASLData ) ; await queue . send ( messages . SASLResponse , { response : SASLContinueResponse , } ) ; const finalMessage = await queue . multiMessageReply ( messages . authenticationSASLFinal , messages . authenticationOk , messages . parameterStatus , messages . backendKeyData , messages . readyForQuery ) ; const finalSASL = finalMessage . AuthenticationSASLFinal ; if ( 'SASLData' in finalSASL ) { checkServerFinalMessage ( finalSASL . SASLData , calculatedServerSignature ) ; } else { throw new Error ( 'SASL: No final SASL data returned' ) ; } } else if ( 'salt' in result ) { password = generateHash ( options . user , password , result . salt ) ; await queue . send ( messages . passwordMessage , { password } ) ; await queue . reply ( messages . authenticationOk ) ; await queue . reply ( messages . readyForQuery ) ; } } catch ( e ) { console . error ( `Connection failed: ${ ( e as any ) . message }` ) ; process . exit ( 1 ) ; } } export async function runQuery ( query , queue ) { const resultRows = [ ] ; await queue . send ( messages . query , { query } ) ; debugQuery ( 'sent query %o' , query ) ; { const result = await queue . reply ( messages . rowDescription ) ; debugQuery ( 'received row description: %o' , result . fields . map ( ( c ) => c . name . toString ( ) ) ) ; } { while ( true ) { const result = await queue . reply ( messages . dataRow , messages . commandComplete ) ; if ( 'commandTag' in result ) { break ; } const row = result . columns . map ( ( c ) => c . value . toString ( ) ) ; resultRows . push ( row ) ; debugQuery ( 'received row data: %o' , row ) ; } } return resultRows ; } export interface IQueryTypes { paramMetadata : { mapping : QueryParam [ ] ; params : MappableType [ ] ; } ; returnTypes : Array < { returnName : string ; columnName : string ; type : MappableType ; nullable ? : boolean ; } > ; } export interface IParseError { errorCode : string ; hint ? : string ; message : string ; position ? : string ; } type TypeData = { fields : Array < { name : string ; tableOID : number ; columnAttrNumber : number ; typeOID : number ; typeSize : number ; typeModifier : number ; formatCode : number ; } > ; params : Array < { oid : number ; } > ; } | IParseError ; param query param query param queue export async function getTypeData ( query , queue ) { const uniqueName = crypto . createHash ( 'md5' ) . update ( query ) . digest ( 'hex' ) ; await queue . send ( messages . parse , { name : uniqueName , query , dataTypes : [ ] , } ) ; await queue . send ( messages . describe , { name : uniqueName , type : PreparedObjectType . Statement , } ) ; await queue . send ( messages . close , { target : PreparedObjectType . Statement , targetName : uniqueName , } ) ; await queue . send ( messages . flush , { } ) ; const parseResult = await queue . reply ( messages . errorResponse , messages . parseComplete ) ; await queue . send ( messages . sync , { } ) ; if ( 'fields' in parseResult ) { const { fields : errorFields } = parseResult ; return { errorCode : errorFields . R , hint : errorFields . H , message : errorFields . M , position : errorFields . P , } ; } const paramsResult = await queue . reply ( messages . parameterDescription , messages . noData ) ; const params = 'params' in paramsResult ? paramsResult . params : [ ] ; const fieldsResult = await queue . reply ( messages . rowDescription , messages . noData ) ; const fields = 'fields' in fieldsResult ? fieldsResult . fields : [ ] ; await queue . reply ( messages . closeComplete ) ; return { params , fields } ; } enum TypeCategory { ARRAY = 'A' , BOOLEAN = 'B' , COMPOSITE = 'C' , DATE_TIME = 'D' , ENUM = 'E' , GEOMETRIC = 'G' , NETWORK_ADDRESS = 'I' , NUMERIC = 'N' , PSEUDO = 'P' , STRING = 'S' , TIMESPAN = 'T' , USERDEFINED = 'U' , BITSTRING = 'V' , UNKNOWN = 'X' } interface TypeRow { oid : string ; typeName : string ; typeKind : string ; enumLabel : string ; typeCategory ? : TypeCategory ; elementTypeOid ? : string ; } export function reduceTypeRows ( typeRows ) { const enumTypes = typeRows . filter ( ( r ) => r . typeKind === DatabaseTypeKind . Enum ) . reduce ( ( typeMap , { oid , typeName , enumLabel } ) => { const typ = typeMap [ oid ] ?? typeName ; return { ... typeMap , [ oid ] : { name : typeName , enumValues : [ ... ( isEnum ( typ ) ? typ . enumValues : [ ] ) , enumLabel ] , } , } ; } , { } as Record < string , MappableType > ) ; return typeRows . reduce ( ( typeMap , { oid , typeName , typeCategory , elementTypeOid } ) => { const typ = typeMap [ oid ] ?? typeName ; if ( oid in enumTypes ) { return { ... typeMap , [ oid ] : enumTypes [ oid ] } ; } if ( typeCategory === TypeCategory . ARRAY && elementTypeOid && elementTypeOid in enumTypes ) { return { ... typeMap , [ oid ] : { name : typeName , elementType : enumTypes [ elementTypeOid ] , } , } ; } return { ... typeMap , [ oid ] : typ } ; } , { } as Record < string , MappableType > ) ; } async function runTypesCatalogQuery ( typeOIDs , queue ) { let rows ; if ( typeOIDs . length > 0 ) { const concatenatedTypeOids = typeOIDs . join ( ',' ) ; rows = await runQuery ( `
SELECT pt.oid, pt.typname, pt.typtype, pe.enumlabel, pt.typelem, pt.typcategory
FROM pg_type pt
LEFT JOIN pg_enum pe ON pt.oid = pe.enumtypid
WHERE pt.oid IN (${ concatenatedTypeOids })
OR pt.oid IN (SELECT typelem FROM pg_type ptn WHERE ptn.oid IN (${ concatenatedTypeOids }));
` , queue ) ; } else { rows = [ ] ; } return rows . map ( ( [ oid , typeName , typeKind , enumLabel , elementTypeOid , typeCategory ] ) => ( { oid , typeName , typeKind , enumLabel , elementTypeOid , typeCategory , } ) ) ; } export async function getTypes ( queryData , queue ) { const typeData = await getTypeData ( queryData . query , queue ) ; if ( 'errorCode' in typeData ) { return typeData ; } const { params , fields } = typeData ; const paramTypeOIDs = params . map ( ( p ) => p . oid ) ; const returnTypesOIDs = fields . map ( ( f ) => f . typeOID ) ; const usedTypesOIDs = paramTypeOIDs . concat ( returnTypesOIDs ) ; const typeRows = await runTypesCatalogQuery ( usedTypesOIDs , queue ) ; const typeMap = reduceTypeRows ( typeRows ) ; const attrMatcher = ( { tableOID , columnAttrNumber , } ) => `(attrelid = ${ tableOID } and attnum = ${ columnAttrNumber })` ; const attrSelection = fields . length > 0 ? fields . map ( attrMatcher ) . join ( ' or ' ) : false ; const attributeRows = await runQuery ( `SELECT
      (attrelid || ':' || attnum) AS attid, attname, attnotnull
     FROM pg_attribute WHERE ${ attrSelection };` , queue ) ; const attrMap = attributeRows . reduce ( ( acc , [ attid , attname , attnotnull ] ) => ( { ... acc , [ attid ] : { columnName : attname , nullable : attnotnull !== 't' , } , } ) , { } ) ; const returnTypes = fields . map ( ( f ) => ( { ... attrMap [ `${ f . tableOID }:${ f . columnAttrNumber }` ] , returnName : f . name , type : typeMap [ f . typeOID ] , } ) ) ; const paramMetadata = { params : params . map ( ( { oid } ) => typeMap [ oid ] ) , mapping : queryData . mapping , } ; return { paramMetadata , returnTypes } ; }