From 5801d9f42f055adadf4e9078f7e3d0877178b3d3 Mon Sep 17 00:00:00 2001
From: Alexander Neumann <alexander@bumpern.de>
Date: Sun, 15 Feb 2015 18:13:42 +0100
Subject: [PATCH] Add EncryptTo() methods

---
 key.go      | 92 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 key_test.go | 38 ++++++++++++++++++++++
 2 files changed, 130 insertions(+)

diff --git a/key.go b/key.go
index e306ec1bc..8ab1b20c0 100644
--- a/key.go
+++ b/key.go
@@ -10,6 +10,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"hash"
 	"io"
 	"io/ioutil"
 	"os"
@@ -375,6 +376,97 @@ func (k *Key) EncryptUserFrom(rd io.Reader) io.Reader {
 	return k.encryptFrom(k.user, rd)
 }
 
+type encryptWriter struct {
+	iv      []byte
+	wroteIV bool
+	h       hash.Hash
+	w       io.Writer
+	origWr  io.Writer
+	err     error // remember error writing iv
+}
+
+func (e *encryptWriter) Close() error {
+	// write hmac
+	_, err := e.origWr.Write(e.h.Sum(nil))
+	if err != nil {
+		return err
+	}
+
+	if w, ok := e.origWr.(io.Closer); ok {
+		err := w.Close()
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func (e *encryptWriter) Write(p []byte) (int, error) {
+	// write iv first
+	if !e.wroteIV {
+		_, e.err = e.origWr.Write(e.iv)
+		e.wroteIV = true
+	}
+
+	if e.err != nil {
+		return 0, e.err
+	}
+
+	n, err := e.w.Write(p)
+	if err != nil {
+		e.err = err
+		return n, err
+	}
+
+	return n, nil
+}
+
+func (k *Key) encryptTo(ks *keys, wr io.Writer) io.WriteCloser {
+	ew := &encryptWriter{
+		iv:     make([]byte, ivSize),
+		h:      hmac.New(sha256.New, ks.Sign),
+		origWr: wr,
+	}
+
+	_, err := io.ReadFull(rand.Reader, ew.iv)
+	if err != nil {
+		panic(fmt.Sprintf("unable to generate new random iv: %v", err))
+	}
+
+	// write iv to hmac
+	_, err = ew.h.Write(ew.iv)
+	if err != nil {
+		panic(err)
+	}
+
+	c, err := aes.NewCipher(ks.Encrypt)
+	if err != nil {
+		panic(fmt.Sprintf("unable to create cipher: %v", err))
+	}
+
+	ew.w = cipher.StreamWriter{
+		S: cipher.NewCTR(c, ew.iv),
+		W: io.MultiWriter(ew.h, wr),
+	}
+
+	return ew
+}
+
+// EncryptTo encrypts and signs data with the master key. The returned
+// io.Writer writes IV || Ciphertext || HMAC. For the hash function, SHA256 is
+// used.
+func (k *Key) EncryptTo(wr io.Writer) io.WriteCloser {
+	return k.encryptTo(k.master, wr)
+}
+
+// EncryptUserTo encrypts and signs data with the user key. The returned
+// io.Writer writes IV || Ciphertext || HMAC. For the hash function, SHA256 is
+// used.
+func (k *Key) EncryptUserTo(wr io.Writer) io.WriteCloser {
+	return k.encryptTo(k.user, wr)
+}
+
 // Decrypt verifes and decrypts the ciphertext. Ciphertext must be in the form
 // IV || Ciphertext || HMAC.
 func (k *Key) decrypt(ks *keys, ciphertext []byte) ([]byte, error) {
diff --git a/key_test.go b/key_test.go
index d303ed695..93a6395f4 100644
--- a/key_test.go
+++ b/key_test.go
@@ -302,3 +302,41 @@ func TestDecryptStreamReader(t *testing.T) {
 			data, plaintext)
 	}
 }
+
+func TestEncryptWriter(t *testing.T) {
+	s := setupBackend(t)
+	defer teardownBackend(t, s)
+	k := setupKey(t, s, testPassword)
+
+	tests := []int{5, 23, 2<<18 + 23, 1 << 20}
+	if *testLargeCrypto {
+		tests = append(tests, 7<<20+123)
+	}
+
+	for _, size := range tests {
+		data := make([]byte, size)
+		_, err := io.ReadFull(randomReader(42, size), data)
+		ok(t, err)
+
+		buf := bytes.NewBuffer(nil)
+		wr := k.EncryptTo(buf)
+
+		_, err = io.Copy(wr, bytes.NewReader(data))
+		ok(t, err)
+		ok(t, wr.Close())
+
+		ciphertext := buf.Bytes()
+
+		l := len(data) + restic.CiphertextExtension
+		assert(t, len(ciphertext) == l,
+			"wrong ciphertext length: expected %d, got %d",
+			l, len(ciphertext))
+
+		// decrypt with default function
+		plaintext, err := k.Decrypt(ciphertext)
+		ok(t, err)
+		assert(t, bytes.Equal(data, plaintext),
+			"wrong plaintext after decryption: expected %02x, got %02x",
+			data, plaintext)
+	}
+}