/
attachment_manual.go
188 lines (169 loc) · 5.38 KB
/
attachment_manual.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
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
package crypto
import (
"io"
"io/ioutil"
"runtime"
"runtime/debug"
"sync"
"time"
"github.com/ProtonMail/go-crypto/openpgp"
"github.com/ProtonMail/go-crypto/openpgp/packet"
"github.com/pkg/errors"
)
// ManualAttachmentProcessor keeps track of the progress of encrypting an attachment
// (optimized for encrypting large files).
// With this processor, the caller has to first allocate
// a buffer large enough to hold the whole data packet.
type ManualAttachmentProcessor struct {
keyPacket []byte
dataLength int
plaintextWriter io.WriteCloser
ciphertextWriter *io.PipeWriter
err error
done sync.WaitGroup
}
// GetKeyPacket returns the key packet for the attachment.
// This should be called only after Finish() has been called.
func (ap *ManualAttachmentProcessor) GetKeyPacket() []byte {
return ap.keyPacket
}
// GetDataLength returns the number of bytes in the DataPacket.
// This should be called only after Finish() has been called.
func (ap *ManualAttachmentProcessor) GetDataLength() int {
return ap.dataLength
}
// Process writes attachment data to be encrypted.
func (ap *ManualAttachmentProcessor) Process(plainData []byte) error {
defer runtime.GC()
_, err := ap.plaintextWriter.Write(plainData)
return errors.Wrap(err, "gopenpgp: couldn't write attachment data")
}
// Finish tells the processor to finalize encryption.
func (ap *ManualAttachmentProcessor) Finish() error {
defer runtime.GC()
if ap.err != nil {
return ap.err
}
if err := ap.plaintextWriter.Close(); err != nil {
return errors.Wrap(err, "gopengpp: unable to close the plaintext writer")
}
if err := ap.ciphertextWriter.Close(); err != nil {
return errors.Wrap(err, "gopengpp: unable to close the dataPacket writer")
}
ap.done.Wait()
if ap.err != nil {
return ap.err
}
return nil
}
// NewManualAttachmentProcessor creates an AttachmentProcessor which can be used
// to encrypt a file. It takes an estimatedSize and filename as hints about the
// file and a buffer to hold the DataPacket.
// It is optimized for low-memory environments and collects garbage every megabyte.
// The buffer for the data packet must be manually allocated by the caller.
// Make sure that the dataBuffer is large enough to hold the whole data packet
// otherwise Finish() will return an error.
func (keyRing *KeyRing) NewManualAttachmentProcessor(
estimatedSize int, filename string, dataBuffer []byte,
) (*ManualAttachmentProcessor, error) {
if len(dataBuffer) == 0 {
return nil, errors.New("gopenpgp: can't give a nil or empty buffer to process the attachement")
}
// forces the gc to be called often
debug.SetGCPercent(10)
attachmentProc := &ManualAttachmentProcessor{}
// hints for the encrypted file
isBinary := true
modTime := GetUnixTime()
hints := &openpgp.FileHints{
FileName: filename,
IsBinary: isBinary,
ModTime: time.Unix(modTime, 0),
}
// encryption config
config := &packet.Config{
DefaultCipher: packet.CipherAES256,
Time: getTimeGenerator(),
}
// goroutine that reads the key packet
// to be later returned to the caller via GetKeyPacket()
keyReader, keyWriter := io.Pipe()
attachmentProc.done.Add(1)
go func() {
defer attachmentProc.done.Done()
keyPacket, err := ioutil.ReadAll(keyReader)
if err != nil {
attachmentProc.err = err
} else {
attachmentProc.keyPacket = clone(keyPacket)
}
}()
// goroutine that reads the data packet into the provided buffer
dataReader, dataWriter := io.Pipe()
attachmentProc.done.Add(1)
go func() {
defer attachmentProc.done.Done()
totalRead, err := readAll(dataBuffer, dataReader)
if err != nil {
attachmentProc.err = err
} else {
attachmentProc.dataLength = totalRead
}
}()
// We generate the encrypting writer
var ew io.WriteCloser
var encryptErr error
ew, encryptErr = openpgp.EncryptSplit(keyWriter, dataWriter, keyRing.entities, nil, hints, config)
if encryptErr != nil {
return nil, errors.Wrap(encryptErr, "gopengpp: unable to encrypt attachment")
}
attachmentProc.plaintextWriter = ew
attachmentProc.ciphertextWriter = dataWriter
// The key packet should have been already written, so we can close
if err := keyWriter.Close(); err != nil {
return nil, errors.Wrap(err, "gopenpgp: couldn't close the keyPacket writer")
}
// Check if the goroutines encountered errors
if attachmentProc.err != nil {
return nil, attachmentProc.err
}
return attachmentProc, nil
}
// readAll works a bit like io.ReadAll
// but we can choose the buffer to write to
// and we don't grow the slice in case of overflow.
func readAll(buffer []byte, reader io.Reader) (int, error) {
bufferLen := len(buffer)
totalRead := 0
offset := 0
overflow := false
reset := false
for {
// We read into the buffer
n, err := reader.Read(buffer[offset:])
totalRead += n
offset += n
if !overflow && reset && n != 0 {
// In case we've started overwriting the beginning of the buffer
// We will return an error at Finish()
overflow = true
}
if err != nil {
if errors.Is(err, io.EOF) {
break
}
return 0, errors.Wrap(err, "gopenpgp: couldn't read data from the encrypted reader")
}
if offset == bufferLen {
// Here we've reached the end of the buffer
// But we need to keep reading to not block the Process()
// So we reset the buffer
reset = true
offset = 0
}
}
if overflow {
return 0, errors.New("gopenpgp: read more bytes that was allocated in the buffer")
}
return totalRead, nil
}