;;; remote-server.scm -- Remote build server.
;;; Copyright © 2020 Mathieu Othacehe <othacehe@gnu.org>
;;;
;;; This file is part of Cuirass.
;;;
;;; GNU Guix is free software; you can redistribute it and/or modify it
;;; under the terms of the GNU General Public License as published by
;;; the Free Software Foundation; either version 3 of the License, or (at
;;; your option) any later version.
;;;
;;; GNU Guix is distributed in the hope that it will be useful, but
;;; WITHOUT ANY WARRANTY; without even the implied warranty of
;;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
;;; GNU General Public License for more details.
;;;
;;; You should have received a copy of the GNU General Public License
;;; along with GNU Guix.  If not, see <http://www.gnu.org/licenses/>.

(define-module (cuirass remote-server)
  #:use-module (cuirass base)
  #:use-module (cuirass remote)
  #:use-module (gcrypt pk-crypto)
  #:use-module (guix avahi)
  #:use-module (guix base32)
  #:use-module (guix base64)
  #:use-module (guix config)
  #:use-module (guix derivations)
  #:use-module (guix records)
  #:use-module (guix packages)
  #:use-module (guix pki)
  #:use-module (guix scripts)
  #:use-module (guix store)
  #:use-module (guix ui)
  #:use-module (guix workers)
  #:use-module (guix build download)
  #:use-module (guix build syscalls)
  #:use-module (gcrypt hash)
  #:use-module (gcrypt pk-crypto)
  #:use-module (simple-zmq)
  #:use-module (rnrs bytevectors)
  #:use-module (srfi srfi-1)
  #:use-module (srfi srfi-26)
  #:use-module (srfi srfi-37)
  #:use-module (ice-9 atomic)
  #:use-module (ice-9 match)
  #:use-module (ice-9 q)
  #:use-module (ice-9 rdelim)
  #:use-module (ice-9 regex)
  #:use-module (ice-9 threads)

  #:export (remote-server))

;; Indicate if the process has to be stopped.
(define %stop-process?
  (make-atomic-box #f))

;; Whether to add build items to the store.
(define %add-to-store?
  (make-parameter #f))

(define %cache-directory
  (make-parameter #f))

(define %private-key
  (make-parameter #f))

(define %public-key
  (make-parameter #f))

(define service-name
  "Cuirass remote server")

(define (show-help)
  (format #t (G_ "Usage: remote-server [OPTION]...
Start a remote build server.\n"))
  (display (G_ "
  -a, --add-to-store        register built items to the store"))
  (display (G_ "
  -b, --backend-port=PORT   listen worker connections on PORT"))
  (display (G_ "
  -p, --publish-port=PORT   publish substitutes on PORT"))
  (display (G_ "
  -c, --cache=DIRECTORY     cache built items to DIRECTORY"))
  (display (G_ "
      --public-key=FILE     use FILE as the public key for signatures"))
  (display (G_ "
      --private-key=FILE    use FILE as the private key for signatures"))
  (newline)
  (display (G_ "
  -h, --help                display this help and exit"))
  (display (G_ "
  -V, --version             display version information and exit"))
  (newline)
  (show-bug-report-information))

(define %options
  (list (option '(#\h "help") #f #f
                (lambda _
                  (show-help)
                  (exit 0)))
        (option '(#\V "version") #f #f
                (lambda _
                  (show-version-and-exit "guix publish")))
        (option '(#\a "add-to-store") #t #f
                (lambda (opt name arg result)
                  (alist-cons 'add-to-store? arg result)))
        (option '(#\b "backend-port") #t #f
                (lambda (opt name arg result)
                  (alist-cons 'backend-port (string->number* arg) result)))
        (option '(#\p "publish-port") #t #f
                (lambda (opt name arg result)
                  (alist-cons 'publish-port (string->number* arg) result)))
        (option '(#\c "cache") #t #f
                (lambda (opt name arg result)
                  (alist-cons 'cache arg result)))
        (option '("public-key") #t #f
                (lambda (opt name arg result)
                  (alist-cons 'public-key-file arg result)))
        (option '("private-key") #t #f
                (lambda (opt name arg result)
                  (alist-cons 'private-key-file arg result)))))

(define %default-options
  `((backend-port . 5555)
    (publish-port . 5556)
    (public-key-file . ,%public-key-file)
    (private-key-file . ,%private-key-file)))

\f
;;;
;;; Build workers.
;;;

(define %workers
  ;; Set of connected workers.
  (make-hash-table))

(define %build-queues
  ;; Builds request queue.
  (map (lambda (system)
         (cons system (make-q)))
       %supported-systems))

(define (find-system-queues systems)
  "Return the list of build queues for SYSTEMS that are not empty."
  (filter-map (match-lambda
                ((system . queue)
                 (and (member system systems)
                      (not (q-empty? queue))
                      queue)))
              %build-queues))

(define (build-available? name)
  "Return #t if there is some available work for the worker with the given
NAME and #f otherwise."
  (let* ((worker (hash-ref %workers name))
         (systems (worker-systems worker))
         (queues (find-system-queues systems)))
    (not (null? queues))))

(define (pop-random-build name)
  "Pop randomly and return a build from all the build queues with available
work for the worker with the given NAME."
  (define (random-queue queues)
    (list-ref queues (random (length queues))))

  (let* ((worker (hash-ref %workers name))
         (systems (worker-systems worker))
         (queues (find-system-queues systems)))
    (q-pop! (random-queue queues))))

(define* (read-client-exp client exp)
  "Read the given EXP sent by CLIENT."
  (catch 'system-error
    (lambda ()
      (match (zmq-read-message exp)
        (('build ('drv drv) ('system system))
         (let ((system (or system
                           (derivation-system
                            (read-derivation-from-file drv)))))
           ;; Push the derivation to the matching queue according to the
           ;; targeted system. Also save the client ID in the queue to be able
           ;; to send it build events later on.
           (q-push! (assoc-ref %build-queues system)
                    (list client drv))))))
    (const #f)))

(define* (read-worker-exp exp #:key reply-worker)
  "Read the given EXP sent by a worker.  REPLY-WORKER is a procedure that can
be used to reply to the worker."
  (match (zmq-read-message exp)
    (('worker-ready worker)
     (let* ((worker* (sexp->worker worker))
            (name (worker-name worker*)))
       (info (G_ "Worker `~a' is ready.~%") name)
       (hash-set! %workers name worker*)))
    (('worker-request-work name)
     (if (build-available? name)
         (match (pop-random-build name)
           ((client drv)
            (reply-worker client (zmq-build-request-message drv))))
         (reply-worker
          (zmq-empty-delimiter)
          (zmq-no-build-message))))))

\f
;;;
;;; Fetch workers.
;;;

(define (zmq-fetch-workers-endpoint)
  "inproc://fetch-workers")

(define (zmq-fetch-worker-socket)
  "Return a socket used to communicate with the fetch workers."
  (let ((socket (zmq-create-socket %zmq-context ZMQ_DEALER))
        (endpoint (zmq-fetch-workers-endpoint)))
    (zmq-connect socket endpoint)
    socket))

(define (strip-store-prefix file)
  ; Given a file name like "/gnu/store/…-foo-1.2/bin/foo", return
  ;; "/bin/foo".
  (let* ((len  (string-length %store-directory))
         (base (string-drop file (+ 1 len))))
    (match (string-index base #\/)
      (#f    base)
      (index (string-drop base index)))))

(define (publish-nar-url publish-url store-hash)
  "Return the URL of STORE-HASH nar substitute on PUBLISH-URL."
  (format #f "~a/nar/gzip/~a" publish-url store-hash))

(define (publish-narinfo-url publish-url store-hash)
  "Return the URL of STORE-HASH narinfo file on PUBLISH-URL."
  (let ((hash (and=> (string-index store-hash #\-)
                     (cut string-take store-hash <>))))
    (format #f "~a/~a.narinfo" publish-url hash)))

(define (nar-path cache-directory output)
  "Return the path of the NAR file for OUTPUT in CACHE-DIRECTORY."
  (string-append cache-directory "/" (basename output) ".nar"))

(define (narinfo-path cache-directory output)
  "Return the path of the NARINFO file for OUTPUT in CACHE-DIRECTORY."
  (string-append cache-directory "/" (basename output) ".narinfo"))

(define* (sign-narinfo! narinfo)
  "Edit the given NARINFO file to replace the worker signature by the remote
build server signature."
  (define (signed-string s)
    (let* ((hash (bytevector->hash-data (sha256 (string->utf8 s))
                                        #:key-type (key-type (%public-key)))))
      (signature-sexp hash (%private-key) (%public-key))))

  (define base64-encode-string
    (compose base64-encode string->utf8))

  (define lines
    (call-with-input-file narinfo
      (lambda (port)
        (let loop ((line (read-line port))
                   (lines '()))
          (if (eof-object? line)
              (reverse lines)
              (loop (read-line port)
                    (cons line lines)))))))
  (let* ((lines
          (filter (lambda (line)
                    (not (string-match "^Signature:" line)))
                  lines))
         (info (format #f "~a~%" (string-join lines "\n")))
         (signature (base64-encode-string
                     (canonical-sexp->string (signed-string info)))))
    (call-with-output-file narinfo
      (lambda (port)
        (format port "~aSignature: 1;~a;~a~%"
                info (gethostname) signature)))))

(define (download-nar cache-directory outputs url)
  "Download in CACHE-DIRECTORY the OUTPUTS from the substitute server at URL."
  (for-each
   (lambda (output)
     (let* ((path (derivation-output-path output))
            (store-hash (strip-store-prefix path))
            (nar-file (nar-path cache-directory store-hash))
            (narinfo-file (narinfo-path cache-directory store-hash))
            (nar-url (publish-nar-url url store-hash))
            (narinfo-url (publish-narinfo-url url store-hash)))
       (unless (file-exists? nar-file)
         (url-fetch nar-url nar-file))

       (unless (file-exists? narinfo-file)
         (url-fetch narinfo-url narinfo-file)
         (sign-narinfo! narinfo-file))))
   outputs))

(define (add-to-store outputs url)
  "Add the OUTPUTS that are available from the substitute server at URL to the
store."
  (with-store store
    (for-each (lambda (output)
                (add-substitute-url store url)
                (ensure-path store output))
              (map derivation-output-path outputs))))

(define (need-fetching? message)
  "Return #t if the received MESSAGE implies that some output fetching is
required and #f otherwise."
  (match (zmq-read-message message)
    (('build-succeeded ('drv drv) ('url url))
     #t)
    (else #f)))

(define* (run-fetch message #:key reply)
  "Read MESSAGE and download the corresponding build outputs.  If
%CACHE-DIRECTORY is set, download the matching NAR and NARINFO files in this
directory.  If %ADD-TO-STORE? is set, add the build outputs to the store.

REPLY is procedure used to forward MESSAGE to the client once the build
outputs are downloaded."
  (define (build-outputs drv)
    (catch 'system-error
      (lambda ()
        (map (match-lambda
               ((output-name . output)
                output))
             (derivation-outputs
              (read-derivation-from-file drv))))
      (const '())))

  (match (zmq-read-message message)
    (('build-succeeded ('drv drv) ('url url))
     (let ((outputs (build-outputs drv)))
       (when (%add-to-store?)
         (add-to-store outputs url))
       (when (%cache-directory)
         (download-nar (%cache-directory) outputs url))
       (reply message)))))

(define (start-fetch-worker name)
  "Start a fetch worker thread with the given NAME.  This worker takes care of
downloading build outputs.  It communicates with the remote server using a ZMQ
socket."
  (define (reply socket client)
    (lambda (message)
      (zmq-send-msg-parts-bytevector
       socket
       (list client (zmq-empty-delimiter) (string->bv message)))))

  (call-with-new-thread
   (lambda ()
     (set-thread-name name)
     (let ((socket (zmq-fetch-worker-socket)))
       (let loop ()
         (match (zmq-get-msg-parts-bytevector socket '())
           ((client empty rest)
            (let ((message (bv->string rest)))
              (run-fetch (bv->string rest)
                         #:reply (reply socket client)))))
         (loop))))))

\f
;;;
;;; ZMQ connection.
;;;

(define %zmq-context
  (zmq-create-context))

(define (zmq-backend-endpoint backend-port)
  "Return a ZMQ endpoint string allowing TCP connections on BACKEND-PORT from
all network interfaces."
  (string-append "tcp://*:" (number->string backend-port)))

(define (zmq-start-proxy backend-port)
  "This procedure starts a proxy between client connections from the IPC
frontend to the workers connected through the TCP backend."
  (define (socket-ready? items socket)
    (find (lambda (item)
            (eq? (poll-item-socket item) socket))
          items))

  (let* ((client-socket
          (zmq-create-socket %zmq-context ZMQ_ROUTER))
         (build-socket
          (zmq-create-socket %zmq-context ZMQ_ROUTER))
         (fetch-socket
          (zmq-create-socket %zmq-context ZMQ_DEALER))
         (poll-items (list
                      (poll-item client-socket ZMQ_POLLIN)
                      (poll-item build-socket ZMQ_POLLIN)
                      (poll-item fetch-socket ZMQ_POLLIN))))

    (zmq-bind-socket client-socket (zmq-frontend-endpoint))
    (zmq-bind-socket build-socket (zmq-backend-endpoint backend-port))
    (zmq-bind-socket fetch-socket (zmq-fetch-workers-endpoint))

    ;; Change frontend socket permissions.
    (chmod (zmq-frontend-socket-name) #o666)

    ;; Do not use the built-in zmq-proxy as we want to edit the envelope of
    ;; frontend messages before forwarding them to the backend.
    (let loop ()
      (let ((items (zmq-poll* poll-items)))
        ;; CLIENT -> REMOTE-SERVER.
        (when (zmq-socket-ready? items client-socket)
          (match (zmq-get-msg-parts-bytevector client-socket)
            ((client empty rest)
             (read-client-exp client (bv->string rest)))))
        ;; BUILD-WORKER -> REMOTE-SERVER.
        (when (zmq-socket-ready? items build-socket)
          (match (zmq-get-msg-parts-bytevector build-socket)
            ((worker empty rest)
             (let ((reply-worker
                    (lambda (client message)
                      (zmq-send-msg-parts-bytevector
                       build-socket
                       (list worker
                             (zmq-empty-delimiter)
                             client
                             (zmq-empty-delimiter)
                             (string->bv message))))))
               (read-worker-exp (bv->string rest)
                                #:reply-worker reply-worker)))
            ((worker empty client empty rest)
             (let ((message (list client (zmq-empty-delimiter) rest)))
               (if (need-fetching? (bv->string rest))
                   (zmq-send-msg-parts-bytevector fetch-socket message)
                   (zmq-send-msg-parts-bytevector client-socket message))))))
        ;; FETCH-WORKER -> REMOTE-SERVER.
        (when (zmq-socket-ready? items fetch-socket)
          (let ((msg (zmq-get-msg-parts-bytevector fetch-socket)))
            (zmq-send-msg-parts-bytevector client-socket msg)))

        (loop)))))

\f
;;;
;;; Entry point.
;;;

;; The PID of the publish process.
(define %publish-pid
  (make-atomic-box #f))

;; The thread running the Avahi publish service.
(define %avahi-thread
  (make-atomic-box #f))

(define (signal-handler)
  "Catch SIGINT to stop the Avahi event loop and the publish process before
exiting."
  (sigaction SIGINT
    (lambda (signum)
      (let ((publish-pid (atomic-box-ref %publish-pid))
            (avahi-thread (atomic-box-ref %avahi-thread)))
        (atomic-box-set! %stop-process? #t)

        (and publish-pid
             (begin
               (kill publish-pid SIGHUP)
               (waitpid publish-pid)))

        (and avahi-thread
             (join-thread avahi-thread))

        (exit 1)))))

(define (remote-server args)
  (signal-handler)

  (with-error-handling
    (let* ((opts (args-fold* args %options
                             (lambda (opt name arg result)
                               (leave (G_ "~A: unrecognized option~%") name))
                             (lambda (arg result)
                               (leave (G_ "~A: extraneous argument~%") arg))
                             %default-options))
           (add-to-store? (assoc-ref opts 'add-to-store?))
           (backend-port (assoc-ref opts 'backend-port))
           (publish-port (assoc-ref opts 'publish-port))
           (cache (assoc-ref opts 'cache))
           (public-key
            (read-file-sexp
             (assoc-ref opts 'public-key-file)))
           (private-key
            (read-file-sexp
             (assoc-ref opts 'private-key-file))))

      (parameterize ((%add-to-store? add-to-store?)
                     (%cache-directory cache)
                     (%public-key public-key)
                     (%private-key private-key))

        (atomic-box-set!
         %publish-pid
         (publish-server publish-port
                         #:public-key public-key
                         #:private-key private-key))

        (atomic-box-set!
         %avahi-thread
         (avahi-publish-service-thread
          service-name
          #:type remote-server-service-type
          #:port backend-port
          #:stop-loop? (lambda ()
                         (atomic-box-ref %stop-process?))
          #:txt (list (string-append "publish="
                                     (number->string publish-port)))))

        (for-each (lambda (number)
                    (start-fetch-worker
                     (string-append "fetch-worker-" (number->string number))))
                  (iota 4))

        (zmq-start-proxy backend-port)))))