/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */ import * as http from 'http' ; import * as net from 'net' ; import WebSocket from 'ws' ; import { v4 as uuidv4 } from 'uuid' ; import jwt from 'jsonwebtoken' ; import cookie from 'cookie' ; import Redis from 'ioredis' ; import StatsD from 'hot-shots' ; import { createLogger } from './logger' ; import { buildConfig } from './config' ; export type StreamResult = [ recordId : string , record : [ label : 'data' , data : string ] ] ; export type ErrorLevel = 'info' | 'warning' | 'error' ; export type SupersetError < ExtraType = Record < string , any > | null > = { error_type : string ; extra : ExtraType ; level : ErrorLevel ; message : string ; } ; type ListenerFunction = ( results : StreamResult [ ] ) => void ; interface EventValue { id : string ; channel_id : string ; job_id : string ; user_id ? : string ; status : string ; errors ? : SupersetError [ ] ; result_url ? : string ; } interface JwtPayload { channel : string ; } interface FetchRangeFromStreamParams { sessionId : string ; startId : string ; endId : string ; listener : ListenerFunction ; } export interface SocketInstance { ws : WebSocket ; channel : string ; pongTs : number ; } interface RedisConfig { port : number ; host : string ; password ? : string | null ; db : number ; ssl : boolean ; } interface ChannelValue { sockets : Array < string > ; } const environment = process . env . NODE_ENV ; const startServer = process . argv [ 2 ] === 'start' ; export const opts = buildConfig ( ) ; const logger = createLogger ( { silent : environment === 'test' , logLevel : opts . logLevel , logToFile : opts . logToFile , logFilename : opts . logFilename , } ) ; export const statsd = new StatsD ( { ... opts . statsd , errorHandler : ( e ) => { logger . error ( e ) ; } , } ) ; if ( startServer && opts . jwtSecret . length < 32 ) throw new Error ( 'Please provide a JWT secret at least 32 bytes long' ) ; export const redisUrlFromConfig = ( redisConfig ) : string => { let url = redisConfig . ssl ? 'rediss://' : 'redis://' ; if ( redisConfig . password ) url += `:${ redisConfig . password }@` ; url += `${ redisConfig . host }:${ redisConfig . port }/${ redisConfig . db }` ; return url ; } ; const redis = new Redis ( redisUrlFromConfig ( opts . redis ) ) ; const httpServer = http . createServer ( ) ; export const wss = new WebSocket . Server ( { noServer : true , clientTracking : false , } ) ; const SOCKET_ACTIVE_STATES = [ WebSocket . OPEN , WebSocket . CONNECTING ] ; const GLOBAL_EVENT_STREAM_NAME = `${ opts . redisStreamPrefix }full` ; const DEFAULT_STREAM_LAST_ID = '$' ; export let channels = { } ; export let sockets = { } ; let lastFirehoseId = DEFAULT_STREAM_LAST_ID ; export const setLastFirehoseId = ( id ) : void => { lastFirehoseId = id ; } ; /**
 * Adds the passed channel and socket instance to the internal registries.
 */ export const trackClient = ( channel , socketInstance ) : string => { statsd . increment ( 'ws_connected_client' ) ; const socketId = uuidv4 ( ) ; sockets [ socketId ] = socketInstance ; if ( channel in channels ) { channels [ channel ] . sockets . push ( socketId ) ; } else { channels [ channel ] = { sockets : [ socketId ] } ; } return socketId ; } ; /**
 * Sends a single async event payload to a single channel.
 * A channel may have multiple connected sockets, this emits
 * the event to all connected sockets within a channel.
 */ export const sendToChannel = ( channel , value ) : void => { const strData = JSON . stringify ( value ) ; if ( ! channels [ channel ] ) { logger . debug ( `channel ${ channel } is unknown, skipping` ) ; return ; } channels [ channel ] . sockets . forEach ( socketId => { const socketInstance = sockets [ socketId ] ; if ( ! socketInstance ) return cleanChannel ( channel ) ; try { socketInstance . ws . send ( strData ) ; } catch ( err ) { statsd . increment ( 'ws_client_send_error' ) ; logger . debug ( `Error sending to socket: ${ err }` ) ; cleanChannel ( channel ) ; } } ) ; } ; /**
 * Reads a range of events from a channel-specific Redis event stream.
 * Invoked in the client re-connection flow.
 */ export const fetchRangeFromStream = async ( { sessionId , startId , endId , listener , } ) => { const streamName = `${ opts . redisStreamPrefix }${ sessionId }` ; try { const reply = await redis . xrange ( streamName , startId , endId ) ; if ( ! reply || ! reply . length ) return ; listener ( reply as StreamResult [ ] ) ; } catch ( e ) { logger . error ( e ) ; } } ; /**
 * Reads from the global Redis event stream continuously.
 * Utilizes a blocking connection to Redis to wait for data to
 * be returned from the stream.
 */ export const subscribeToGlobalStream = async ( stream , listener ) => { while ( true ) { try { const reply = await redis . xread ( 'BLOCK' , opts . redisStreamReadBlockMs , 'COUNT' , opts . redisStreamReadCount , 'STREAMS' , stream , lastFirehoseId ) ; if ( ! reply ) { continue ; } const results = reply [ 0 ] [ 1 ] ; const { length } = results ; if ( ! results . length ) { continue ; } listener ( results as StreamResult [ ] ) ; setLastFirehoseId ( results [ length - 1 ] [ 0 ] ) ; } catch ( e ) { logger . error ( e ) ; continue ; } } } ; /**
 * Callback function to process events received from a Redis Stream
 */ export const processStreamResults = ( results ) : void => { logger . debug ( `events received: ${ results }` ) ; results . forEach ( item => { try { const id = item [ 0 ] ; const data = JSON . parse ( item [ 1 ] [ 1 ] ) ; sendToChannel ( data . channel_id , { id , ... data } ) ; } catch ( err ) { logger . error ( err ) ; } } ) ; } ; /**
 * Verify and parse a JWT cookie from an HTTP request.
 * Returns the JWT payload or throws an error on invalid token.
 */ const getJwtPayload = ( request ) : JwtPayload => { const cookies = cookie . parse ( request . headers . cookie || '' ) ; const token = cookies [ opts . jwtCookieName ] ; if ( ! token ) throw new Error ( 'JWT not present' ) ; return jwt . verify ( token , opts . jwtSecret ) as JwtPayload ; } ; /**
 * Extracts the `last_id` query param value from an HTTP request
 */ const getLastId = ( request ) : string | null => { const url = new URL ( String ( request . url ) , 'http://0.0.0.0' ) ; const queryParams = url . searchParams ; return queryParams . get ( 'last_id' ) ; } ; /**
 * Increments a Redis Stream ID
 */ export const incrementId = ( id ) : string => { const parts = id . split ( '-' ) ; if ( parts . length < 2 ) return id ; return parts [ 0 ] + '-' + ( Number ( parts [ 1 ] ) + 1 ) ; } ; /**
 * WebSocket `connection` event handler, called via wss
 */ export const wsConnection = ( ws , request ) => { const jwtPayload = getJwtPayload ( request ) ; const channel = jwtPayload . channel ; const socketInstance = { ws , channel , pongTs : Date . now ( ) } ; const socketId = trackClient ( channel , socketInstance ) ; logger . debug ( `socket ${ socketId } connected on channel ${ channel }` ) ; const lastId = getLastId ( request ) ; if ( lastId ) { const endId = lastFirehoseId === DEFAULT_STREAM_LAST_ID ? '+' : lastFirehoseId ; fetchRangeFromStream ( { sessionId : channel , startId : incrementId ( lastId ) , endId , listener : processStreamResults , } ) ; } ws . on ( 'pong' , function pong ( data ) { const socketId = data . toString ( ) ; const socketInstance = sockets [ socketId ] ; if ( ! socketInstance ) { logger . warn ( `pong received for nonexistent socket ${ socketId }` ) ; } else { socketInstance . pongTs = Date . now ( ) ; } } ) ; } ; /**
 * HTTP `request` event handler, called via httpServer
 */ export const httpRequest = ( request , response ) => { const rawUrl = request . url as string ; const method = request . method as string ; const headers = request . headers || { } ; const url = new URL ( rawUrl as string , `http://${ headers . host }` ) ; if ( url . pathname === '/health' && [ 'GET' , 'HEAD' ] . includes ( method ) ) { response . writeHead ( 200 ) ; response . end ( 'OK' ) ; } else { logger . info ( `Received unexpected request: ${ method } ${ rawUrl }` ) ; response . writeHead ( 404 ) ; response . end ( 'Not Found' ) ; } } ; /**
 * HTTP `upgrade` event handler, called via httpServer
 */ export const httpUpgrade = ( request , socket , head ) => { try { const jwtPayload = getJwtPayload ( request ) ; if ( ! jwtPayload . channel ) throw new Error ( 'Channel ID not present' ) ; } catch ( err ) { logger . error ( err ) ; socket . destroy ( ) ; return ; } wss . handleUpgrade ( request , socket , head , function cb ( ws , request ) { wss . emit ( 'connection' , ws , request ) ; } ) ; } ; /**
 * Iterate over all tracked sockets, terminating and removing references to
 * connections that have not responded with a _pong_ within the timeout window.
 * Sends a _ping_ to all active connections.
 */ export const checkSockets = ( ) => { logger . debug ( `channel count: ${ Object . keys ( channels ) . length }` ) ; logger . debug ( `socket count: ${ Object . keys ( sockets ) . length }` ) ; for ( const socketId in sockets ) { const socketInstance = sockets [ socketId ] ; const timeout = Date . now ( ) - socketInstance . pongTs ; let isActive = true ; if ( timeout >= opts . socketResponseTimeoutMs ) { logger . debug ( `terminating unresponsive socket: ${ socketId }, channel: ${ socketInstance . channel }` ) ; socketInstance . ws . terminate ( ) ; isActive = false ; } else if ( ! SOCKET_ACTIVE_STATES . includes ( socketInstance . ws . readyState ) ) { isActive = false ; } if ( isActive ) { socketInstance . ws . ping ( socketId ) ; } else { delete sockets [ socketId ] ; logger . debug ( `forgetting socket ${ socketId }` ) ; } } } ; /**
 * Iterate over all sockets within a channel, removing references to
 * inactive connections, ultimately removing the channel from the
 * _channels_ registry if no active connections remain.
 */ export const cleanChannel = ( channel ) => { const activeSockets = channels [ channel ] ?. sockets . filter ( socketId => { const socketInstance = sockets [ socketId ] ; if ( ! socketInstance ) return false ; if ( SOCKET_ACTIVE_STATES . includes ( socketInstance . ws . readyState ) ) return true ; return false ; } ) || [ ] ; if ( activeSockets . length === 0 ) { delete channels [ channel ] ; } else { channels [ channel ] . sockets = activeSockets ; } } ; if ( startServer ) { wss . on ( 'connection' , wsConnection ) ; httpServer . on ( 'request' , httpRequest ) ; httpServer . on ( 'upgrade' , httpUpgrade ) ; httpServer . listen ( opts . port ) ; logger . info ( `Server started on port ${ opts . port }` ) ; subscribeToGlobalStream ( GLOBAL_EVENT_STREAM_NAME , processStreamResults ) ; setInterval ( checkSockets , opts . pingSocketsIntervalMs ) ; setInterval ( function gc ( ) { for ( const channel in channels ) { cleanChannel ( channel ) ; } } , opts . gcChannelsIntervalMs ) ; } export const resetState = ( ) => { channels = { } ; sockets = { } ; lastFirehoseId = DEFAULT_STREAM_LAST_ID ; } ;