497 lines
14 KiB
Plaintext
497 lines
14 KiB
Plaintext
import { ReadableStream } from 'node:stream/web'
|
|
import { Readable } from 'node:stream'
|
|
import { TSR_SCRIPT_BARRIER_ID } from './constants'
|
|
import type { AnyRouter } from '../router'
|
|
|
|
export function transformReadableStreamWithRouter(
|
|
router: AnyRouter,
|
|
routerStream: ReadableStream,
|
|
) {
|
|
return transformStreamWithRouter(router, routerStream)
|
|
}
|
|
|
|
export function transformPipeableStreamWithRouter(
|
|
router: AnyRouter,
|
|
routerStream: Readable,
|
|
) {
|
|
return Readable.fromWeb(
|
|
transformStreamWithRouter(router, Readable.toWeb(routerStream)),
|
|
)
|
|
}
|
|
|
|
// Use string constants for simple indexOf matching
|
|
const BODY_END_TAG = '</body>'
|
|
const HTML_END_TAG = '</html>'
|
|
|
|
// Minimum length of a valid closing tag: </a> = 4 characters
|
|
const MIN_CLOSING_TAG_LENGTH = 4
|
|
|
|
// Default timeout values (in milliseconds)
|
|
const DEFAULT_SERIALIZATION_TIMEOUT_MS = 60000
|
|
const DEFAULT_LIFETIME_TIMEOUT_MS = 60000
|
|
|
|
// Module-level encoder (stateless, safe to reuse)
|
|
const textEncoder = new TextEncoder()
|
|
|
|
/**
|
|
* Finds the position just after the last valid HTML closing tag in the string.
|
|
*
|
|
* Valid closing tags match the pattern: </[a-zA-Z][\w:.-]*>
|
|
* Examples: </div>, </my-component>, </slot:name.nested>
|
|
*
|
|
* @returns Position after the last closing tag, or -1 if none found
|
|
*/
|
|
function findLastClosingTagEnd(str: string): number {
|
|
const len = str.length
|
|
if (len < MIN_CLOSING_TAG_LENGTH) return -1
|
|
|
|
let i = len - 1
|
|
|
|
while (i >= MIN_CLOSING_TAG_LENGTH - 1) {
|
|
// Look for > (charCode 62)
|
|
if (str.charCodeAt(i) === 62) {
|
|
// Look backwards for valid tag name characters
|
|
let j = i - 1
|
|
|
|
// Skip through valid tag name characters
|
|
while (j >= 1) {
|
|
const code = str.charCodeAt(j)
|
|
// Check if it's a valid tag name char: [a-zA-Z0-9_:.-]
|
|
if (
|
|
(code >= 97 && code <= 122) || // a-z
|
|
(code >= 65 && code <= 90) || // A-Z
|
|
(code >= 48 && code <= 57) || // 0-9
|
|
code === 95 || // _
|
|
code === 58 || // :
|
|
code === 46 || // .
|
|
code === 45 // -
|
|
) {
|
|
j--
|
|
} else {
|
|
break
|
|
}
|
|
}
|
|
|
|
// Check if the first char after </ is a valid start char (letter only)
|
|
const tagNameStart = j + 1
|
|
if (tagNameStart < i) {
|
|
const startCode = str.charCodeAt(tagNameStart)
|
|
// Tag name must start with a letter (a-z or A-Z)
|
|
if (
|
|
(startCode >= 97 && startCode <= 122) ||
|
|
(startCode >= 65 && startCode <= 90)
|
|
) {
|
|
// Check for </ (charCodes: < = 60, / = 47)
|
|
if (
|
|
j >= 1 &&
|
|
str.charCodeAt(j) === 47 &&
|
|
str.charCodeAt(j - 1) === 60
|
|
) {
|
|
return i + 1 // Return position after the closing >
|
|
}
|
|
}
|
|
}
|
|
}
|
|
i--
|
|
}
|
|
return -1
|
|
}
|
|
|
|
export function transformStreamWithRouter(
|
|
router: AnyRouter,
|
|
appStream: ReadableStream,
|
|
opts?: {
|
|
/** Timeout for serialization to complete after app render finishes (default: 60000ms) */
|
|
timeoutMs?: number
|
|
/** Maximum lifetime of the stream transform (default: 60000ms). Safety net for cleanup. */
|
|
lifetimeMs?: number
|
|
},
|
|
) {
|
|
// Check upfront if serialization already finished synchronously
|
|
// This is the fast path for routes with no deferred data
|
|
const serializationAlreadyFinished =
|
|
router.serverSsr?.isSerializationFinished() ?? false
|
|
|
|
// Take any HTML that was buffered before we started listening
|
|
const initialBufferedHtml = router.serverSsr?.takeBufferedHtml()
|
|
|
|
// True passthrough: if serialization already finished and nothing buffered,
|
|
// we can avoid any decoding/scanning while still honoring cleanup + setRenderFinished.
|
|
if (serializationAlreadyFinished && !initialBufferedHtml) {
|
|
let cleanedUp = false
|
|
let controller: ReadableStreamDefaultController<Uint8Array> | undefined
|
|
let isStreamClosed = false
|
|
let lifetimeTimeoutHandle: ReturnType<typeof setTimeout> | undefined
|
|
|
|
const cleanup = () => {
|
|
if (cleanedUp) return
|
|
cleanedUp = true
|
|
|
|
if (lifetimeTimeoutHandle !== undefined) {
|
|
clearTimeout(lifetimeTimeoutHandle)
|
|
lifetimeTimeoutHandle = undefined
|
|
}
|
|
|
|
router.serverSsr?.cleanup()
|
|
}
|
|
|
|
const safeClose = () => {
|
|
if (isStreamClosed) return
|
|
isStreamClosed = true
|
|
try {
|
|
controller?.close()
|
|
} catch {
|
|
// ignore
|
|
}
|
|
}
|
|
|
|
const safeError = (error: unknown) => {
|
|
if (isStreamClosed) return
|
|
isStreamClosed = true
|
|
try {
|
|
controller?.error(error)
|
|
} catch {
|
|
// ignore
|
|
}
|
|
}
|
|
|
|
const lifetimeMs = opts?.lifetimeMs ?? DEFAULT_LIFETIME_TIMEOUT_MS
|
|
lifetimeTimeoutHandle = setTimeout(() => {
|
|
if (!cleanedUp && !isStreamClosed) {
|
|
console.warn(
|
|
`SSR stream transform exceeded maximum lifetime (${lifetimeMs}ms), forcing cleanup`,
|
|
)
|
|
safeError(new Error('Stream lifetime exceeded'))
|
|
cleanup()
|
|
}
|
|
}, lifetimeMs)
|
|
|
|
const stream = new ReadableStream<Uint8Array>({
|
|
start(c: ReadableStreamDefaultController<Uint8Array>) {
|
|
controller = c
|
|
},
|
|
cancel() {
|
|
isStreamClosed = true
|
|
cleanup()
|
|
},
|
|
})
|
|
|
|
;(async () => {
|
|
const reader = appStream.getReader()
|
|
try {
|
|
while (true) {
|
|
const { done, value } = await reader.read()
|
|
if (done) break
|
|
if (cleanedUp || isStreamClosed) return
|
|
controller?.enqueue(value as unknown as Uint8Array)
|
|
}
|
|
|
|
if (cleanedUp || isStreamClosed) return
|
|
|
|
router.serverSsr?.setRenderFinished()
|
|
safeClose()
|
|
cleanup()
|
|
} catch (error) {
|
|
if (cleanedUp) return
|
|
console.error('Error reading appStream:', error)
|
|
router.serverSsr?.setRenderFinished()
|
|
safeError(error)
|
|
cleanup()
|
|
} finally {
|
|
reader.releaseLock()
|
|
}
|
|
})().catch((error) => {
|
|
if (cleanedUp) return
|
|
console.error('Error in stream transform:', error)
|
|
safeError(error)
|
|
cleanup()
|
|
})
|
|
|
|
return stream
|
|
}
|
|
|
|
let stopListeningToInjectedHtml: (() => void) | undefined
|
|
let stopListeningToSerializationFinished: (() => void) | undefined
|
|
let serializationTimeoutHandle: ReturnType<typeof setTimeout> | undefined
|
|
let lifetimeTimeoutHandle: ReturnType<typeof setTimeout> | undefined
|
|
let cleanedUp = false
|
|
|
|
let controller: ReadableStreamDefaultController<any>
|
|
let isStreamClosed = false
|
|
|
|
const textDecoder = new TextDecoder()
|
|
|
|
// concat'd router HTML; avoids array joins on each flush
|
|
let pendingRouterHtml = initialBufferedHtml ?? ''
|
|
|
|
// between-chunk text buffer; keep bounded to avoid unbounded memory
|
|
let leftover = ''
|
|
|
|
// captured closing tags from </body> onward
|
|
let pendingClosingTags = ''
|
|
|
|
// conservative cap: enough to hold any partial closing tag + a bit
|
|
const MAX_LEFTOVER_CHARS = 2048
|
|
|
|
let isAppRendering = true
|
|
let streamBarrierLifted = false
|
|
let serializationFinished = serializationAlreadyFinished
|
|
|
|
function safeEnqueue(chunk: string | Uint8Array) {
|
|
if (isStreamClosed) return
|
|
if (typeof chunk === 'string') {
|
|
controller.enqueue(textEncoder.encode(chunk))
|
|
} else {
|
|
controller.enqueue(chunk)
|
|
}
|
|
}
|
|
|
|
function safeClose() {
|
|
if (isStreamClosed) return
|
|
isStreamClosed = true
|
|
try {
|
|
controller.close()
|
|
} catch {
|
|
// ignore
|
|
}
|
|
}
|
|
|
|
function safeError(error: unknown) {
|
|
if (isStreamClosed) return
|
|
isStreamClosed = true
|
|
try {
|
|
controller.error(error)
|
|
} catch {
|
|
// ignore
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Cleanup with guards; must be idempotent.
|
|
*/
|
|
function cleanup() {
|
|
if (cleanedUp) return
|
|
cleanedUp = true
|
|
|
|
try {
|
|
stopListeningToInjectedHtml?.()
|
|
stopListeningToSerializationFinished?.()
|
|
} catch {
|
|
// ignore
|
|
}
|
|
stopListeningToInjectedHtml = undefined
|
|
stopListeningToSerializationFinished = undefined
|
|
|
|
if (serializationTimeoutHandle !== undefined) {
|
|
clearTimeout(serializationTimeoutHandle)
|
|
serializationTimeoutHandle = undefined
|
|
}
|
|
if (lifetimeTimeoutHandle !== undefined) {
|
|
clearTimeout(lifetimeTimeoutHandle)
|
|
lifetimeTimeoutHandle = undefined
|
|
}
|
|
|
|
pendingRouterHtml = ''
|
|
leftover = ''
|
|
pendingClosingTags = ''
|
|
|
|
router.serverSsr?.cleanup()
|
|
}
|
|
|
|
const stream = new ReadableStream({
|
|
start(c: ReadableStreamDefaultController<any>) {
|
|
controller = c
|
|
},
|
|
cancel() {
|
|
isStreamClosed = true
|
|
cleanup()
|
|
},
|
|
})
|
|
|
|
function flushPendingRouterHtml() {
|
|
if (!pendingRouterHtml) return
|
|
safeEnqueue(pendingRouterHtml)
|
|
pendingRouterHtml = ''
|
|
}
|
|
|
|
function appendRouterHtml(html: string) {
|
|
if (!html) return
|
|
pendingRouterHtml += html
|
|
}
|
|
|
|
/**
|
|
* Finish only when app done and serialization complete.
|
|
*/
|
|
function tryFinish() {
|
|
if (isAppRendering || !serializationFinished) return
|
|
if (cleanedUp || isStreamClosed) return
|
|
|
|
if (serializationTimeoutHandle !== undefined) {
|
|
clearTimeout(serializationTimeoutHandle)
|
|
serializationTimeoutHandle = undefined
|
|
}
|
|
|
|
// Flush any remaining bytes in the TextDecoder
|
|
const decoderRemainder = textDecoder.decode()
|
|
|
|
if (leftover) safeEnqueue(leftover)
|
|
if (decoderRemainder) safeEnqueue(decoderRemainder)
|
|
flushPendingRouterHtml()
|
|
if (pendingClosingTags) safeEnqueue(pendingClosingTags)
|
|
|
|
safeClose()
|
|
cleanup()
|
|
}
|
|
|
|
// Safety net: cleanup even if consumer never reads
|
|
const lifetimeMs = opts?.lifetimeMs ?? DEFAULT_LIFETIME_TIMEOUT_MS
|
|
lifetimeTimeoutHandle = setTimeout(() => {
|
|
if (!cleanedUp && !isStreamClosed) {
|
|
console.warn(
|
|
`SSR stream transform exceeded maximum lifetime (${lifetimeMs}ms), forcing cleanup`,
|
|
)
|
|
safeError(new Error('Stream lifetime exceeded'))
|
|
cleanup()
|
|
}
|
|
}, lifetimeMs)
|
|
|
|
if (!serializationAlreadyFinished) {
|
|
stopListeningToInjectedHtml = router.subscribe('onInjectedHtml', () => {
|
|
if (cleanedUp || isStreamClosed) return
|
|
const html = router.serverSsr?.takeBufferedHtml()
|
|
if (!html) return
|
|
|
|
// If we've already captured </body> (pendingClosingTags), we must keep appending
|
|
// so injection stays before the stored closing tags.
|
|
if (isAppRendering || leftover || pendingClosingTags) {
|
|
appendRouterHtml(html)
|
|
} else {
|
|
// App is done rendering - flush any pending buffer first to maintain order,
|
|
// then write the new HTML directly
|
|
flushPendingRouterHtml()
|
|
safeEnqueue(html)
|
|
}
|
|
})
|
|
|
|
stopListeningToSerializationFinished = router.subscribe(
|
|
'onSerializationFinished',
|
|
() => {
|
|
serializationFinished = true
|
|
tryFinish()
|
|
},
|
|
)
|
|
}
|
|
|
|
// Transform the appStream
|
|
;(async () => {
|
|
const reader = appStream.getReader()
|
|
try {
|
|
while (true) {
|
|
const { done, value } = await reader.read()
|
|
if (done) break
|
|
|
|
if (cleanedUp || isStreamClosed) return
|
|
|
|
const text =
|
|
value instanceof Uint8Array
|
|
? textDecoder.decode(value, { stream: true })
|
|
: String(value)
|
|
|
|
// Fast path: most chunks have no pending left-over.
|
|
const chunkString = leftover ? leftover + text : text
|
|
|
|
if (!streamBarrierLifted) {
|
|
if (chunkString.includes(TSR_SCRIPT_BARRIER_ID)) {
|
|
streamBarrierLifted = true
|
|
router.serverSsr?.liftScriptBarrier()
|
|
}
|
|
}
|
|
|
|
// If we already saw </body>, everything else is part of tail; buffer it.
|
|
if (pendingClosingTags) {
|
|
pendingClosingTags += chunkString
|
|
leftover = ''
|
|
continue
|
|
}
|
|
|
|
const bodyEndIndex = chunkString.indexOf(BODY_END_TAG)
|
|
const htmlEndIndex = chunkString.indexOf(HTML_END_TAG)
|
|
|
|
if (
|
|
bodyEndIndex !== -1 &&
|
|
htmlEndIndex !== -1 &&
|
|
bodyEndIndex < htmlEndIndex
|
|
) {
|
|
pendingClosingTags = chunkString.slice(bodyEndIndex)
|
|
safeEnqueue(chunkString.slice(0, bodyEndIndex))
|
|
flushPendingRouterHtml()
|
|
leftover = ''
|
|
continue
|
|
}
|
|
|
|
const lastClosingTagEnd = findLastClosingTagEnd(chunkString)
|
|
|
|
if (lastClosingTagEnd > 0) {
|
|
safeEnqueue(chunkString.slice(0, lastClosingTagEnd))
|
|
flushPendingRouterHtml()
|
|
|
|
leftover = chunkString.slice(lastClosingTagEnd)
|
|
if (leftover.length > MAX_LEFTOVER_CHARS) {
|
|
// Ensure bounded memory even if a consumer streams long text sequences
|
|
// without any closing tags. This may reduce injection granularity but is correct.
|
|
safeEnqueue(leftover.slice(0, leftover.length - MAX_LEFTOVER_CHARS))
|
|
leftover = leftover.slice(-MAX_LEFTOVER_CHARS)
|
|
}
|
|
} else {
|
|
// No closing tag found; keep small tail to handle split closing tags,
|
|
// but stream older bytes to prevent unbounded buffering.
|
|
const combined = chunkString
|
|
if (combined.length > MAX_LEFTOVER_CHARS) {
|
|
const flushUpto = combined.length - MAX_LEFTOVER_CHARS
|
|
safeEnqueue(combined.slice(0, flushUpto))
|
|
leftover = combined.slice(flushUpto)
|
|
} else {
|
|
leftover = combined
|
|
}
|
|
}
|
|
}
|
|
|
|
if (cleanedUp || isStreamClosed) return
|
|
|
|
isAppRendering = false
|
|
router.serverSsr?.setRenderFinished()
|
|
|
|
if (serializationFinished) {
|
|
tryFinish()
|
|
} else {
|
|
const timeoutMs = opts?.timeoutMs ?? DEFAULT_SERIALIZATION_TIMEOUT_MS
|
|
serializationTimeoutHandle = setTimeout(() => {
|
|
if (!cleanedUp && !isStreamClosed) {
|
|
console.error('Serialization timeout after app render finished')
|
|
safeError(
|
|
new Error('Serialization timeout after app render finished'),
|
|
)
|
|
cleanup()
|
|
}
|
|
}, timeoutMs)
|
|
}
|
|
} catch (error) {
|
|
if (cleanedUp) return
|
|
console.error('Error reading appStream:', error)
|
|
isAppRendering = false
|
|
router.serverSsr?.setRenderFinished()
|
|
safeError(error)
|
|
cleanup()
|
|
} finally {
|
|
reader.releaseLock()
|
|
}
|
|
})().catch((error) => {
|
|
if (cleanedUp) return
|
|
console.error('Error in stream transform:', error)
|
|
safeError(error)
|
|
cleanup()
|
|
})
|
|
|
|
return stream
|
|
}
|