orca/src/main/ssh/relay-protocol.test.ts

import { describe, expect, it, vi } from 'vitest'
import {
  HEADER_LENGTH,
  MessageType,
  encodeFrame,
  encodeJsonRpcFrame,
  encodeKeepAliveFrame,
  FrameDecoder,
  parseJsonRpcMessage,
  parseUnameToRelayPlatform,
  isGitResponseStreamMarker,
  type JsonRpcRequest,
  type DecodedFrame
} from './relay-protocol'

describe('git response stream marker', () => {
  it('accepts only complete non-negative integer metadata', () => {
    expect(
      isGitResponseStreamMarker({
        __orcaGitResponseStream: { streamId: 1, totalBytes: 1024, chunkCount: 2 }
      })
    ).toBe(true)
    expect(isGitResponseStreamMarker({ __orcaGitResponseStream: {} })).toBe(false)
    expect(
      isGitResponseStreamMarker({
        __orcaGitResponseStream: { streamId: -1, totalBytes: 1024, chunkCount: 2 }
      })
    ).toBe(false)
  })
})

describe('frame encoding', () => {
  it('encodes a frame with 13-byte header', () => {
    const payload = Buffer.from('hello')
    const frame = encodeFrame(MessageType.Regular, 1, 0, payload)

    expect(frame.length).toBe(HEADER_LENGTH + payload.length)
    expect(frame[0]).toBe(MessageType.Regular)
    expect(frame.readUInt32BE(1)).toBe(1) // ID
    expect(frame.readUInt32BE(5)).toBe(0) // ACK
    expect(frame.readUInt32BE(9)).toBe(5) // LENGTH
    expect(frame.subarray(HEADER_LENGTH).toString()).toBe('hello')
  })

  it('encodes keepalive frame with empty payload', () => {
    const frame = encodeKeepAliveFrame(42, 10)

    expect(frame.length).toBe(HEADER_LENGTH)
    expect(frame[0]).toBe(MessageType.KeepAlive)
    expect(frame.readUInt32BE(1)).toBe(42) // ID
    expect(frame.readUInt32BE(5)).toBe(10) // ACK
    expect(frame.readUInt32BE(9)).toBe(0) // LENGTH
  })

  it('encodes JSON-RPC frame', () => {
    const msg: JsonRpcRequest = {
      jsonrpc: '2.0',
      id: 1,
      method: 'pty.spawn',
      params: { cols: 80, rows: 24 }
    }
    const frame = encodeJsonRpcFrame(msg, 5, 3)

    expect(frame[0]).toBe(MessageType.Regular)
    expect(frame.readUInt32BE(1)).toBe(5)
    expect(frame.readUInt32BE(5)).toBe(3)

    const payloadLen = frame.readUInt32BE(9)
    const payload = frame.subarray(HEADER_LENGTH, HEADER_LENGTH + payloadLen)
    const decoded = JSON.parse(payload.toString('utf-8'))
    expect(decoded.method).toBe('pty.spawn')
    expect(decoded.params.cols).toBe(80)
  })

  it('rejects messages larger than MAX_MESSAGE_SIZE', () => {
    const bigPayload = {
      jsonrpc: '2.0' as const,
      id: 1,
      method: 'x',
      params: { data: 'a'.repeat(17 * 1024 * 1024) }
    }
    expect(() => encodeJsonRpcFrame(bigPayload, 1, 0)).toThrow('Message too large')
  })
})

describe('FrameDecoder', () => {
  // Framing correctness should not depend on wall-clock turn budgeting; the
  // default 4ms maxTurnMs can defer later frames via setImmediate under CI load.
  // Time-sliced drain is covered by relay-protocol-backpressure.test.ts.
  const frozenClock = { now: () => 0 }

  it('decodes a complete frame', () => {
    const frames: DecodedFrame[] = []
    const decoder = new FrameDecoder((f) => frames.push(f), undefined, frozenClock)

    const payload = Buffer.from('test')
    const encoded = encodeFrame(MessageType.Regular, 1, 0, payload)
    decoder.feed(encoded)

    expect(frames).toHaveLength(1)
    expect(frames[0].type).toBe(MessageType.Regular)
    expect(frames[0].id).toBe(1)
    expect(frames[0].ack).toBe(0)
    expect(frames[0].payload.toString()).toBe('test')
  })

  it('handles partial frames across multiple feeds', () => {
    const frames: DecodedFrame[] = []
    const decoder = new FrameDecoder((f) => frames.push(f), undefined, frozenClock)

    const payload = Buffer.from('hello world')
    const encoded = encodeFrame(MessageType.Regular, 2, 1, payload)

    // Feed in two parts
    decoder.feed(encoded.subarray(0, 10))
    expect(frames).toHaveLength(0)

    decoder.feed(encoded.subarray(10))
    expect(frames).toHaveLength(1)
    expect(frames[0].payload.toString()).toBe('hello world')
  })

  it('decodes multiple frames from a single chunk', () => {
    const frames: DecodedFrame[] = []
    const decoder = new FrameDecoder((f) => frames.push(f), undefined, frozenClock)

    const frame1 = encodeFrame(MessageType.Regular, 1, 0, Buffer.from('a'))
    const frame2 = encodeFrame(MessageType.Regular, 2, 1, Buffer.from('b'))
    const combined = Buffer.concat([frame1, frame2])

    decoder.feed(combined)
    expect(frames).toHaveLength(2)
    expect(frames[0].payload.toString()).toBe('a')
    expect(frames[1].payload.toString()).toBe('b')
  })

  it('decodes keepalive frames', () => {
    const frames: DecodedFrame[] = []
    const decoder = new FrameDecoder((f) => frames.push(f), undefined, frozenClock)

    decoder.feed(encodeKeepAliveFrame(5, 3))
    expect(frames).toHaveLength(1)
    expect(frames[0].type).toBe(MessageType.KeepAlive)
    expect(frames[0].payload.length).toBe(0)
  })

  it('skips oversized frames and calls onError instead of throwing', () => {
    const errors: Error[] = []
    const frames: DecodedFrame[] = []
    const decoder = new FrameDecoder(
      (f) => frames.push(f),
      (err) => errors.push(err),
      frozenClock
    )

    const oversizedLength = 17 * 1024 * 1024
    const header = Buffer.alloc(HEADER_LENGTH)
    header[0] = MessageType.Regular
    header.writeUInt32BE(1, 1)
    header.writeUInt32BE(0, 5)
    header.writeUInt32BE(oversizedLength, 9)

    const fakePayload = Buffer.alloc(oversizedLength)
    const fullFrame = Buffer.concat([header, fakePayload])

    decoder.feed(fullFrame)
    expect(frames).toHaveLength(0)
    expect(errors).toHaveLength(1)
    expect(errors[0].message).toContain('discarded')
  })

  it('reset clears internal buffer', () => {
    const frames: DecodedFrame[] = []
    const decoder = new FrameDecoder((f) => frames.push(f), undefined, frozenClock)

    // Feed a partial frame
    const encoded = encodeFrame(MessageType.Regular, 1, 0, Buffer.from('test'))
    decoder.feed(encoded.subarray(0, 5))
    decoder.reset()

    // Feed a new complete frame
    decoder.feed(encodeFrame(MessageType.Regular, 2, 0, Buffer.from('new')))
    expect(frames).toHaveLength(1)
    expect(frames[0].id).toBe(2)
  })

  it('decodes frames fed one byte at a time (worst-case boundary straddling)', () => {
    const frames: DecodedFrame[] = []
    const decoder = new FrameDecoder((f) => frames.push(f), undefined, frozenClock)

    const frame1 = encodeFrame(MessageType.Regular, 1, 0, Buffer.from('first payload'))
    const frame2 = encodeFrame(MessageType.Regular, 2, 1, Buffer.from('second'))
    const combined = Buffer.concat([frame1, frame2])

    for (let i = 0; i < combined.length; i += 1) {
      decoder.feed(combined.subarray(i, i + 1))
    }

    expect(frames).toHaveLength(2)
    expect(frames[0].payload.toString()).toBe('first payload')
    expect(frames[1].id).toBe(2)
    expect(frames[1].payload.toString()).toBe('second')
  })

  it('skips an oversized frame fed in odd-sized chunks and resynchronizes', () => {
    const errors: Error[] = []
    const frames: DecodedFrame[] = []
    const decoder = new FrameDecoder(
      (f) => frames.push(f),
      (err) => errors.push(err),
      frozenClock
    )

    const oversizedLength = 17 * 1024 * 1024
    const header = Buffer.alloc(HEADER_LENGTH)
    header[0] = MessageType.Regular
    header.writeUInt32BE(1, 1)
    header.writeUInt32BE(0, 5)
    header.writeUInt32BE(oversizedLength, 9)
    const oversized = Buffer.concat([header, Buffer.alloc(oversizedLength)])
    const valid = encodeFrame(MessageType.Regular, 2, 0, Buffer.from('after'))
    const combined = Buffer.concat([oversized, valid])

    const chunkSize = 1024 * 1024 - 7
    for (let i = 0; i < combined.length; i += chunkSize) {
      decoder.feed(combined.subarray(i, i + chunkSize))
    }

    expect(errors).toHaveLength(1)
    expect(frames).toHaveLength(1)
    expect(frames[0].payload.toString()).toBe('after')
  })

  it('never rebuilds the buffered stream per feed while assembling a large frame', () => {
    // Regression: feed() used Buffer.concat([buffered, chunk]) per data event,
    // re-copying the whole backlog for every TCP chunk — O(n²) memcpy on the
    // Electron main thread while fs.streamChunk frames arrive (SSH typing lag).
    const frames: DecodedFrame[] = []
    const decoder = new FrameDecoder((f) => frames.push(f), undefined, frozenClock)
    const frame = encodeFrame(MessageType.Regular, 1, 0, Buffer.alloc(512 * 1024, 0x61))

    const concatSpy = vi.spyOn(Buffer, 'concat')
    try {
      for (let i = 0; i < frame.length; i += 32 * 1024) {
        decoder.feed(frame.subarray(i, i + 32 * 1024))
      }
    } finally {
      concatSpy.mockRestore()
    }

    expect(frames).toHaveLength(1)
    expect(frames[0].payload.length).toBe(512 * 1024)
    expect(concatSpy).not.toHaveBeenCalled()
  })
})

describe('parseJsonRpcMessage', () => {
  it('parses a valid JSON-RPC request', () => {
    const payload = Buffer.from(
      JSON.stringify({
        jsonrpc: '2.0',
        id: 1,
        method: 'pty.spawn',
        params: { cols: 80 }
      })
    )
    const msg = parseJsonRpcMessage(payload)
    expect('method' in msg && msg.method).toBe('pty.spawn')
  })

  it('throws on invalid jsonrpc version', () => {
    const payload = Buffer.from(JSON.stringify({ jsonrpc: '1.0', id: 1, method: 'x' }))
    expect(() => parseJsonRpcMessage(payload)).toThrow('Invalid JSON-RPC version')
  })

  it('throws on malformed JSON', () => {
    const payload = Buffer.from('not json')
    expect(() => parseJsonRpcMessage(payload)).toThrow()
  })
})

describe('parseUnameToRelayPlatform', () => {
  it('maps Linux x86_64', () => {
    expect(parseUnameToRelayPlatform('Linux', 'x86_64')).toBe('linux-x64')
  })

  it('maps Linux aarch64', () => {
    expect(parseUnameToRelayPlatform('Linux', 'aarch64')).toBe('linux-arm64')
  })

  it('maps Darwin x86_64', () => {
    expect(parseUnameToRelayPlatform('Darwin', 'x86_64')).toBe('darwin-x64')
  })

  it('maps Darwin arm64', () => {
    expect(parseUnameToRelayPlatform('Darwin', 'arm64')).toBe('darwin-arm64')
  })

  it('handles amd64 alias', () => {
    expect(parseUnameToRelayPlatform('Linux', 'amd64')).toBe('linux-x64')
  })

  it('maps Windows amd64', () => {
    expect(parseUnameToRelayPlatform('Windows', 'AMD64')).toBe('win32-x64')
  })

  it('maps Windows X64 runtime architecture', () => {
    expect(parseUnameToRelayPlatform('Windows', 'X64')).toBe('win32-x64')
  })

  it('maps Windows arm64', () => {
    expect(parseUnameToRelayPlatform('win32', 'ARM64')).toBe('win32-arm64')
  })

  it('maps MSYS/MINGW uname output as Windows', () => {
    expect(parseUnameToRelayPlatform('MINGW64_NT-10.0', 'x86_64')).toBe('win32-x64')
  })

  it('returns null for unsupported OS', () => {
    expect(parseUnameToRelayPlatform('FreeBSD', 'x86_64')).toBeNull()
  })

  it('returns null for unsupported arch', () => {
    expect(parseUnameToRelayPlatform('Linux', 'mips')).toBeNull()
  })

  it('is case-insensitive', () => {
    expect(parseUnameToRelayPlatform('LINUX', 'X86_64')).toBe('linux-x64')
  })
})