nix_actor/src/nix_actor/sockets.nim

# SPDX-FileCopyrightText: ☭ Emery Hemingway
# SPDX-License-Identifier: Unlicense

## Common module for communicating with Nix clients and daemons.

import std/[asyncdispatch, asyncnet, sets, strtabs, strutils, tables]
from std/nativesockets import AF_UNIX, SOCK_STREAM, Protocol

import eris
import preserves, syndicate

import ./protocol

{.pragma: workerProtocol, importc, header: "worker-protocol.hh".}

type Word* = uint64

proc `[]=`*[T](attrs: var AttrSet; key: string; val: T) =
  attrs[Symbol key] = val.toPreserve

const
  WORKER_MAGIC_1* = 0x6E697863
  WORKER_MAGIC_2* = 0x6478696F
  PROTOCOL_VERSION* = 0x100 or 35

  STDERR_NEXT* = 0x6F6C6d67
  STDERR_READ* = 0x64617461
  STDERR_WRITE* = 0x64617416
  STDERR_LAST* = 0x616C7473
  STDERR_ERROR* = 0x63787470
  STDERR_START_ACTIVITY* = 0x53545254
  STDERR_STOP_ACTIVITY* = 0x53544F50
  STDERR_RESULT* = 0x52534C54

  wopIsValidPath* = 1
  wopHasSubstitutes* = 3
  wopQueryReferrers* = 6
  wopAddToStore* = 7
  wopBuildPaths* = 9
  wopEnsurePath* = 10
  wopAddTempRoot* = 11
  wopAddIndirectRoot* = 12
  wopSyncWithGC* = 13
  wopFindRoots* = 14
  wopSetOptions* = 19
  wopCollectGarbage* = 20
  wopQuerySubstitutablePathInfo* = 21
  wopQueryAllValidPaths* = 23
  wopQueryFailedPaths* = 24
  wopClearFailedPaths* = 25
  wopQueryPathInfo* = 26
  wopQueryPathFromHashPart* = 29
  wopQuerySubstitutablePathInfos* = 30
  wopQueryValidPaths* = 31
  wopQuerySubstitutablePaths* = 32
  wopQueryValidDerivers* = 33
  wopOptimiseStore* = 34
  wopVerifyStore* = 35
  wopBuildDerivation* = 36
  wopAddSignatures* = 37
  wopNarFromPath* = 38
  wopAddToStoreNar* = 39
  wopQueryMissing* = 40
  wopQueryDerivationOutputMap* = 41
  wopRegisterDrvOutput* = 42
  wopQueryRealisation* = 43
  wopAddMultipleToStore* = 44
  wopAddBuildLog* = 45
  wopBuildPathsWithResults* = 46

type
  ProtocolError* = object of IOError
  Version* = uint16

  Session* = ref object
    socket*: AsyncSocket
    buffer*: seq[Word]
    version*: Version

func major*(version: Version): uint16 = version and 0xff00
func minor*(version: Version): uint16 = version and 0x00ff

proc close*(session: Session) =
  close(session.socket)
  reset(session.buffer)

proc send*(session: Session; words: varargs[Word]): Future[void] =
  if session.buffer.len < words.len:
    session.buffer.setLen(words.len)
  for i, word in words: session.buffer[i] = word
  send(session.socket, addr session.buffer[0], words.len shl 3)

proc send*(session: Session; s: string): Future[void] =
  let wordCount = 1 + ((s.len + 7) shr 3)
  if session.buffer.len < wordCount: setLen(session.buffer, wordCount)
  session.buffer[0] = Word s.len
  if s != "":
    session.buffer[pred wordCount] = 0x00
    copyMem(addr session.buffer[1], unsafeAddr s[0], s.len)
  send(session.socket, addr session.buffer[0], wordCount shl 3)

proc send*(session: Session; ss: StringSeq|StringSet): Future[void] =
  ## Send a set of strings. The set is sent as a contiguous buffer.
  session.buffer[0] = Word ss.len
  var off = 1
  for s in ss:
    let
      stringWordLen = (s.len + 7) shr 3
      bufferWordLen = off+1+stringWordLen
    if session.buffer.len < bufferWordLen:
      setLen(session.buffer, bufferWordLen)
    session.buffer[off] = Word s.len
    session.buffer[off+stringWordLen] = 0 # clear the aligning bits
    inc(off)
    copyMem(addr session.buffer[off], unsafeAddr s[0], s.len)
    inc(off, stringWordLen)
  send(session.socket, addr session.buffer[0], off shl 3)

proc recvWord*(sock: AsyncSocket): Future[Word] {.async.} =
  var w: Word
  let n = await recvInto(sock, addr w, sizeof(Word))
  if n != sizeof(Word): raise newException(ProtocolError, "short read")
  return w

proc recvWord*(session: Session): Future[Word] =
  recvWord(session.socket)

proc discardWords*(session: Session; n: int): Future[void] {.async.} =
  if session.buffer.len < n: setLen(session.buffer, n)
  let byteCount = n shl 3
  let n = await recvInto(session.socket, addr session.buffer[0], byteCount)
  if n != byteCount:
      raise newException(ProtocolError, "short read")

proc recvString*(socket: AsyncSocket): Future[string] {.async.} =
  let stringLen = int (await recvWord(socket))
  if stringLen > 0:
    var s = newString((stringLen + 7) and (not 7))
    let n = await recvInto(socket, addr s[0], s.len)
    if n != s.len:
      raise newException(ProtocolError, "short read")
    setLen(s, stringLen)
    return s
  return ""

proc recvString*(session: Session): Future[string] =
  recvString(session.socket)

proc recvStringSeq*(session: Session): Future[StringSeq] {.async.} =
  let count = int(await recvWord(session.socket))
  var strings = newSeq[string](count)
  for i in 0..<count: strings[i] = await recvString(session)
  return strings

proc recvStringSet*(session: Session): Future[StringSet] {.async.} =
  let count = int(await recvWord(session.socket))
  var strings = initHashSet[string](count)
  for i in 0..<count: incl(strings, await recvString(session))
  return strings

proc newUnixSocket*(): AsyncSocket =
  newAsyncSocket(
      domain = AF_UNIX,
      sockType = SOCK_STREAM,
      protocol = cast[Protocol](0),
      buffered = false,
    )

proc newSession*(socket: AsyncSocket): Session =
  Session(socket: socket, buffer: newSeq[Word](512))

proc newSession*(): Session =
  newUnixSocket().newSession()

proc ingestChunks*(session: Session; store: ErisStore): Future[ErisCap] {.async.} =
  var ingest: ErisIngest
  while true:
    let chunkLen = int await recvWord(session)
    if ingest.isNil:
      ingest = newErisIngest(
        store, recommendedChunkSize(chunkLen), convergentMode)
    if chunkLen == 0:
      break
    else:
      let wordLen = (chunkLen + 7) shr 3
      if session.buffer.len < wordLen: setLen(session.buffer, wordLen)
      let recvLen = await recvInto(session.socket, addr session.buffer[0], chunkLen)
        # each chunk must be received contiguously
      if recvLen != chunkLen:
        raise newException(ProtocolError, "invalid chunk read")
      await append(ingest, addr session.buffer[0], chunkLen)
  var cap = await cap(ingest)
  return cap

proc recoverChunks*(session: Session; store: ErisStore; cap: ErisCap) {.async.} =
  let stream = newErisStream(store, cap)
  session.buffer.setLen(succ(cap.chunkSize.int shr 3))
  while true:
    let n = await stream.readBuffer(addr session.buffer[1], cap.chunkSize.int)
    session.buffer[0] = Word n
    await send(session.socket, addr session.buffer[0], 8+n)
    if n == 0: break
  close(stream)