Skip to content

Commit

Permalink
upload: allow encrypting file in-place with UploadReader
Browse files Browse the repository at this point in the history
  • Loading branch information
tulir committed Sep 11, 2024
1 parent 85e67c5 commit bb3364a
Show file tree
Hide file tree
Showing 2 changed files with 16 additions and 2 deletions.
2 changes: 2 additions & 0 deletions upload.go
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,8 @@ func (cli *Client) Upload(ctx context.Context, plaintext []byte, appInfo MediaTy
// This is otherwise identical to [Upload], but it reads the plaintext from an [io.Reader] instead of a byte slice.
// A temporary file is required for the encryption process. If tempFile is nil, a temporary file will be created
// and deleted after the upload.
//
// To use only one file, pass the same file as both plaintext and tempFile. This will cause the file to be overwritten with encrypted data.
func (cli *Client) UploadReader(ctx context.Context, plaintext io.Reader, tempFile io.ReadWriteSeeker, appInfo MediaType) (resp UploadResponse, err error) {
resp.MediaKey = random.Bytes(32)
iv, cipherKey, macKey, _ := getMediaKeys(resp.MediaKey, appInfo)
Expand Down
16 changes: 14 additions & 2 deletions util/cbcutil/cbc.go
Original file line number Diff line number Diff line change
Expand Up @@ -168,8 +168,11 @@ func EncryptStream(key, iv, macKey []byte, plaintext io.Reader, ciphertext io.Wr
cipherMAC := hmac.New(sha256.New, macKey)
cipherMAC.Write(iv)

writerAt, hasWriterAt := ciphertext.(io.WriterAt)

buf := make([]byte, 32*1024)
var size int
var writePtr int64
hasMore := true
for hasMore {
var n int
Expand All @@ -186,14 +189,23 @@ func EncryptStream(key, iv, macKey []byte, plaintext io.Reader, ciphertext io.Wr
cbc.CryptBlocks(buf, buf)
cipherMAC.Write(buf)
cipherHasher.Write(buf)
_, err = ciphertext.Write(buf)
if hasWriterAt {
_, err = writerAt.WriteAt(buf, writePtr)
writePtr += int64(len(buf))
} else {
_, err = ciphertext.Write(buf)
}
if err != nil {
return nil, nil, 0, fmt.Errorf("failed to write file: %w", err)
}
}
mac := cipherMAC.Sum(nil)[:10]
cipherHasher.Write(mac)
_, err = ciphertext.Write(mac)
if hasWriterAt {
_, err = writerAt.WriteAt(mac, writePtr)
} else {
_, err = ciphertext.Write(mac)
}
if err != nil {
return nil, nil, 0, fmt.Errorf("failed to write checksum to file: %w", err)
}
Expand Down

0 comments on commit bb3364a

Please sign in to comment.