forked from ProtonMail/gopenpgp
/
attachment.go
147 lines (124 loc) · 4.11 KB
/
attachment.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
package crypto
import (
"bytes"
"io"
"io/ioutil"
"runtime"
"sync"
"golang.org/x/crypto/openpgp"
"golang.org/x/crypto/openpgp/packet"
)
// AttachmentProcessor keeps track of the progress of encrypting an attachment
// (optimized for encrypting large files).
type AttachmentProcessor struct {
w *io.WriteCloser
pipe *io.PipeWriter
done sync.WaitGroup
split *PGPSplitMessage
garbageCollector int
err error
}
// Process writes attachment data to be encrypted.
func (ap *AttachmentProcessor) Process(plainData []byte) {
if _, err := (*ap.w).Write(plainData); err != nil {
panic(err)
}
}
// Finish closes the attachment and returns the encrypted data.
func (ap *AttachmentProcessor) Finish() (*PGPSplitMessage, error) {
if ap.err != nil {
return nil, ap.err
}
if err := (*ap.w).Close(); err != nil {
return nil, err
}
if err := (*ap.pipe).Close(); err != nil {
return nil, err
}
ap.done.Wait()
if ap.garbageCollector > 0 {
runtime.GC()
}
return ap.split, nil
}
// newAttachmentProcessor creates an AttachmentProcessor which can be used to encrypt
// a file. It takes an estimatedSize and fileName as hints about the file.
func (keyRing *KeyRing) newAttachmentProcessor(
estimatedSize int, fileName string, garbageCollector int,
) (*AttachmentProcessor, error) {
attachmentProc := &AttachmentProcessor{}
// You could also add these one at a time if needed.
attachmentProc.done.Add(1)
attachmentProc.garbageCollector = garbageCollector
hints := &openpgp.FileHints{
FileName: fileName,
}
config := &packet.Config{
DefaultCipher: packet.CipherAES256,
Time: getTimeGenerator(),
}
reader, writer := io.Pipe()
go func() {
defer attachmentProc.done.Done()
ciphertext, _ := ioutil.ReadAll(reader)
message := NewPGPMessage(ciphertext)
split, splitError := message.SeparateKeyAndData(estimatedSize, garbageCollector)
if attachmentProc.err != nil {
attachmentProc.err = splitError
}
attachmentProc.split = split
}()
var ew io.WriteCloser
var encryptErr error
ew, encryptErr = openpgp.Encrypt(writer, keyRing.entities, nil, hints, config)
if encryptErr != nil {
return nil, encryptErr
}
attachmentProc.w = &ew
attachmentProc.pipe = writer
return attachmentProc, nil
}
// EncryptAttachment encrypts a file given a PlainMessage and a fileName.
// Returns a PGPSplitMessage containing a session key packet and symmetrically encrypted data.
// Specifically designed for attachments rather than text messages.
func (keyRing *KeyRing) EncryptAttachment(message *PlainMessage, fileName string) (*PGPSplitMessage, error) {
ap, err := keyRing.newAttachmentProcessor(len(message.GetBinary()), fileName, -1)
if err != nil {
return nil, err
}
ap.Process(message.GetBinary())
split, err := ap.Finish()
if err != nil {
return nil, err
}
return split, nil
}
// NewLowMemoryAttachmentProcessor creates an AttachmentProcessor which can be used
// to encrypt a file. It takes an estimatedSize and fileName as hints about the
// file. It is optimized for low-memory environments and collects garbage every
// megabyte.
func (keyRing *KeyRing) NewLowMemoryAttachmentProcessor(
estimatedSize int, fileName string,
) (*AttachmentProcessor, error) {
return keyRing.newAttachmentProcessor(estimatedSize, fileName, 1<<20)
}
// DecryptAttachment takes a PGPSplitMessage, containing a session key packet and symmetrically encrypted data
// and returns a decrypted PlainMessage
// Specifically designed for attachments rather than text messages.
func (keyRing *KeyRing) DecryptAttachment(message *PGPSplitMessage) (*PlainMessage, error) {
privKeyEntries := keyRing.entities
keyReader := bytes.NewReader(message.GetBinaryKeyPacket())
dataReader := bytes.NewReader(message.GetBinaryDataPacket())
encryptedReader := io.MultiReader(keyReader, dataReader)
config := &packet.Config{Time: getTimeGenerator()}
md, err := openpgp.ReadMessage(encryptedReader, privKeyEntries, nil, config)
if err != nil {
return nil, err
}
decrypted := md.UnverifiedBody
b, err := ioutil.ReadAll(decrypted)
if err != nil {
return nil, err
}
return NewPlainMessage(b), nil
}