s3: stream cat/get decrypted chunks to output port
ober
f4a26f8d03a496069e393c7a78d8a9edf8d7850e
--- a/protonstorage/cli.ss +++ b/protonstorage/cli.ss @@ -1360,16 +1360,15 @@ (die 2 "usage: jdrive s3 cat REMOTE_PATH")) (let* ([password (s3-vault-password-from-opts opts)] [material (s3-yubikey-material-for-profile opts)]) - (put-bytevector - (current-output-port) - (jdrive-s3-cat-file-range-bytes - (profile-option opts) - (state-root-option opts) - password - material - (car pos) - (option-number opts "--offset" 0) - (option-number opts "--size" -1))))] + (jdrive-s3-stream-file-range-to-port + (profile-option opts) + (state-root-option opts) + password + material + (car pos) + (option-number opts "--offset" 0) + (option-number opts "--size" -1) + (current-output-port)))] [(string=? action "get") (unless (= (length pos) 2) (die 2 "usage: jdrive s3 get REMOTE_PATH LOCAL_PATH")) --- a/protonstorage/s3/drive.ss +++ b/protonstorage/s3/drive.ss @@ -33,6 +33,7 @@ jdrive-s3-move! jdrive-s3-cat-file-bytes jdrive-s3-cat-file-range-bytes + jdrive-s3-stream-file-range-to-port jdrive-s3-get-file! jdrive-s3-remove!) @@ -2324,18 +2325,56 @@ "CopiedChunks" (hashtable-ref copy-summary "CopiedChunks" 0) "RemovedFiles" (hashtable-ref remove-summary "RemovedFiles" 0)))))) - (define (jdrive-s3-cat-file-bytes profile state-root vault-password yubikey-material remote-path) - (jdrive-s3-cat-file-range-bytes - profile - state-root - vault-password - yubikey-material - remote-path - 0 - -1)) - - (define (jdrive-s3-cat-file-range-bytes - profile state-root vault-password yubikey-material remote-path offset size) + ;; Stream the [offset, want-end) byte range of one entry to `out`, decrypting + ;; a single chunk at a time and writing it before the next is fetched, so + ;; memory stays bounded by one chunk instead of the whole file. Returns the + ;; number of plaintext bytes written. + (define (write-entry-range-to-port! client bucket drive-key path entry offset want-end out) + (if (entry-chunked? entry) + (let ([object-id (jdrive-s3-object-id drive-key path)]) + (let loop ([chunks (entry-chunks entry)] + [chunk-start 0] + [written 0]) + (cond + [(or (null? chunks) (>= chunk-start want-end)) written] + [else + (let* ([chunk (car chunks)] + [chunk-size (chunk-field chunk "Size" 0)] + [chunk-end (+ chunk-start chunk-size)]) + (if (or (<= chunk-end offset) + (>= chunk-start want-end)) + (loop (cdr chunks) chunk-end written) + (let* ([plain + (decrypt-chunk-entry + client + bucket + drive-key + path + object-id + chunk)] + [selected + (slice-for-range + plain + chunk-start + offset + want-end)]) + (unless (bytevector-empty? selected) + (put-bytevector out selected)) + (loop + (cdr chunks) + chunk-end + (+ written (bytevector-length selected))))))]))) + (let* ([plain (decrypt-single-entry client bucket drive-key path entry)] + [selected (slice-for-range plain 0 offset want-end)]) + (unless (bytevector-empty? selected) + (put-bytevector out selected)) + (bytevector-length selected)))) + + ;; Stream a ranged remote file straight to an output port without buffering + ;; the whole file. Used by `cat` (stdout) and `get` (secure output + ;; transaction). Returns the number of plaintext bytes written. + (define (jdrive-s3-stream-file-range-to-port + profile state-root vault-password yubikey-material remote-path offset size out) (with-unlocked-profile profile state-root @@ -2350,46 +2389,39 @@ (let* ([total (jmaybe entry "Size" 0)] [offset (max 0 offset)] [want-end (range-finish offset size total)]) - (if (entry-chunked? entry) - (let ([object-id (jdrive-s3-object-id drive-key path)]) - (let loop ([chunks (entry-chunks entry)] - [chunk-start 0] - [out '()]) - (cond - [(or (null? chunks) (>= chunk-start want-end)) - (bytevectors-concat (reverse out))] - [else - (let* ([chunk (car chunks)] - [chunk-size (chunk-field chunk "Size" 0)] - [chunk-end (+ chunk-start chunk-size)]) - (if (or (<= chunk-end offset) - (>= chunk-start want-end)) - (loop (cdr chunks) chunk-end out) - (let* ([plain - (decrypt-chunk-entry - client - bucket - drive-key - path - object-id - chunk)] - [selected - (slice-for-range - plain - chunk-start - offset - want-end)]) - (loop - (cdr chunks) - chunk-end - (if (bytevector-empty? selected) - out - (cons selected out))))))]))) - (slice-for-range - (decrypt-single-entry client bucket drive-key path entry) - 0 - offset - want-end))))))) + (write-entry-range-to-port! + client + bucket + drive-key + path + entry + offset + want-end + out)))))) + + (define (jdrive-s3-cat-file-bytes profile state-root vault-password yubikey-material remote-path) + (jdrive-s3-cat-file-range-bytes + profile + state-root + vault-password + yubikey-material + remote-path + 0 + -1)) + + (define (jdrive-s3-cat-file-range-bytes + profile state-root vault-password yubikey-material remote-path offset size) + (call-with-bytevector-output-port + (lambda (out) + (jdrive-s3-stream-file-range-to-port + profile + state-root + vault-password + yubikey-material + remote-path + offset + size + out)))) (define (jdrive-s3-get-file! profile state-root vault-password yubikey-material remote-path local-path @@ -2398,28 +2430,32 @@ [size (if (and (pair? maybe-range) (pair? (cdr maybe-range))) (cadr maybe-range) -1)] - [data - (jdrive-s3-cat-file-range-bytes - profile - state-root - vault-password - yubikey-material - remote-path - offset - size)]) - ;; Write via the descriptor-relative secure output transaction (the same - ;; no-follow path recursive downloads use) so a symlink at local-path - ;; cannot redirect the decrypted bytes elsewhere. - (let ([root-handle (secure-directory-open (parent-directory local-path) #t)] - [name (path-basename local-path)]) - (dynamic-wind - (lambda () (void)) - (lambda () - (call-with-secure-output-file - root-handle - name - (lambda (out) (put-bytevector out data)))) - (lambda () (secure-directory-close root-handle)))) + [root-handle (secure-directory-open (parent-directory local-path) #t)] + [name (path-basename local-path)] + [written 0]) + ;; Stream decrypted chunks straight into the descriptor-relative secure + ;; output transaction (the same no-follow path recursive downloads use) so + ;; a symlink at local-path cannot redirect the bytes and a large file is + ;; never fully buffered. The transaction writes to a temp file and commits + ;; with an atomic rename; on error it aborts and leaves no partial file. + (dynamic-wind + (lambda () (void)) + (lambda () + (call-with-secure-output-file + root-handle + name + (lambda (out) + (set! written + (jdrive-s3-stream-file-range-to-port + profile + state-root + vault-password + yubikey-material + remote-path + offset + size + out))))) + (lambda () (secure-directory-close root-handle))) (json-object "Backend" "s3" "Profile" (jdrive-s3-profile-name profile) @@ -2427,7 +2463,7 @@ "LocalPath" local-path "Offset" offset "RequestedSize" size - "Bytes" (bytevector-length data)))) + "Bytes" written))) (define (remove-dir-path dirs prefix recursive?) (let loop ([xs dirs] [out '()]) --- a/test/integration-s3.ss +++ b/test/integration-s3.ss @@ -80,6 +80,43 @@ (put-bytevector p chunk 0 n)) (loop (- remaining n))))))))) +(define (write-varying-file path size) + (ensure-directory! (let loop ([i (- (string-length path) 1)]) + (cond + [(< i 0) "."] + [(char=? (string-ref path i) #\/) + (if (= i 0) "/" (substring path 0 i))] + [else (loop (- i 1))]))) + (let ([chunk (make-bytevector (* 64 1024))]) + (call-with-port + (open-file-output-port path (file-options no-fail) (buffer-mode block)) + (lambda (p) + (let loop ([written 0]) + (when (< written size) + (let* ([n (min (- size written) (bytevector-length chunk))] + [buf (if (= n (bytevector-length chunk)) + chunk + (make-bytevector n))]) + (let fill ([i 0]) + (when (< i n) + (bytevector-u8-set! buf i (modulo (+ written i) 251)) + (fill (+ i 1)))) + (put-bytevector p buf 0 n) + (loop (+ written n))))))))) + +(define (read-file-bytes path) + (call-with-port + (open-file-input-port path (file-options) (buffer-mode block)) + (lambda (p) (get-bytevector-all p)))) + +(define (bv-slice bv start end) + (let* ([n (bytevector-length bv)] + [s (max 0 (min start n))] + [e (max s (min end n))] + [out (make-bytevector (- e s))]) + (bytevector-copy! bv s out 0 (- e s)) + out)) + (define (bytevector=? a b) (and (= (bytevector-length a) (bytevector-length b)) (let loop ([i 0]) @@ -200,6 +237,34 @@ "/nested/b.txt" restored) (check "get writes restored file" (file-exists? restored)) +(check "get streams restored chunked file bytes" + (bytevector=? + (string->utf8 "nested plaintext for chunking") + (read-file-bytes restored))) + +(let* ([stream-src (path-join local-root "stream-src.bin")] + [stream-size (+ (* 6 chunk-size) 5)] + [stream-remote "/stream.bin"] + [stream-restored (string-append restored ".stream")] + [stream-range (string-append restored ".stream.range")] + [range-offset chunk-size] + [range-size (+ (* 2 chunk-size) 3)]) + (write-varying-file stream-src stream-size) + (jdrive-s3-put-local-path! + profile state-root vault-password #vu8() + stream-src stream-remote #f chunk-size) + (jdrive-s3-get-file! + profile state-root vault-password #vu8() + stream-remote stream-restored) + (check "get streams large multi-chunk file byte-for-byte" + (bytevector=? (read-file-bytes stream-src) (read-file-bytes stream-restored))) + (jdrive-s3-get-file! + profile state-root vault-password #vu8() + stream-remote stream-range range-offset range-size) + (check "get streams multi-chunk range byte-for-byte" + (bytevector=? + (bv-slice (read-file-bytes stream-src) range-offset (+ range-offset range-size)) + (read-file-bytes stream-range)))) (jdrive-s3-copy! profile state-root vault-password #vu8()