forked from sahib/brig
-
Notifications
You must be signed in to change notification settings - Fork 0
/
writer.go
203 lines (161 loc) · 4.84 KB
/
writer.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
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
package encrypt
import (
"bytes"
"encoding/binary"
"errors"
"io"
)
var (
// ErrBadBlockSize is returned when the data is damaged and has an invalid block size
ErrBadBlockSize = errors.New("underlying reader failed to read full w.maxBlockSize")
// ErrMixedMethods is returned when calling Write() with ReadFrom() together.
ErrMixedMethods = errors.New("mixing Write() and ReadFrom() is not allowed")
)
// Writer encrypts the data stream before writing to Writer.
type Writer struct {
// Internal Writer we would write to.
io.Writer
// Common fields with Reader
aeadCommon
// A buffer that is max. w.maxBlockSize big.
// Used for caching leftover data between writes.
rbuf *bytes.Buffer
// Index of the currently written block.
blockCount uint64
// True after the first write.
headerWritten bool
// w.maxBlockSize is the maximum number of bytes a single payload may have
maxBlockSize int64
// Used encryption algorithm
cipher uint16
}
// GoodDecBufferSize returns a buffer size that is suitable for decryption.
func (w *Writer) GoodDecBufferSize() int64 {
return w.maxBlockSize
}
// GoodEncBufferSize returns a buffer size that is suitable for encryption.
func (w *Writer) GoodEncBufferSize() int64 {
return w.maxBlockSize + 40
}
func (w *Writer) emitHeaderIfNeeded() error {
if !w.headerWritten {
w.headerWritten = true
header := GenerateHeader(w.key, w.maxBlockSize, w.cipher)
if _, err := w.Writer.Write(header); err != nil {
return err
}
}
return nil
}
func (w *Writer) Write(p []byte) (int, error) {
if err := w.emitHeaderIfNeeded(); err != nil {
return 0, err
}
for int64(w.rbuf.Len()) >= w.maxBlockSize {
if _, err := w.flushPack(w.rbuf.Next(int(w.maxBlockSize))); err != nil {
return 0, err
}
}
// Remember left-overs for next write:
if _, err := w.rbuf.Write(p); err != nil {
return 0, nil
}
// Fake the amount of data we've written:
return len(p), nil
}
func (w *Writer) flushPack(pack []byte) (int, error) {
// Create a new Nonce for this block:
binary.LittleEndian.PutUint64(w.nonce, w.blockCount)
// Encrypt the text:
w.encBuf = w.aead.Seal(w.encBuf[:0], w.nonce, pack, nil)
// Pass it to the underlying writer:
nNonce, err := w.Writer.Write(w.nonce)
if err != nil {
return nNonce, err
}
w.blockCount++
nBuf, err := w.Writer.Write(w.encBuf)
return nNonce + nBuf, err
}
// Close the Writer and write any left-over blocks
// This does not close the underlying data stream.
func (w *Writer) Close() error {
if err := w.emitHeaderIfNeeded(); err != nil {
return err
}
// Flush last block of data if any:
for w.rbuf.Len() > 0 {
n := int64(w.rbuf.Len())
if n > w.maxBlockSize {
n = w.maxBlockSize
}
if _, err := w.flushPack(w.rbuf.Next(int(n))); err != nil {
return err
}
}
return nil
}
// ReadFrom writes all readable from `r` into `w`.
//
// It is intentend as optimized way to copy the whole stream without
// unneeded copying in between. io.Copy() will use this function automatically.
//
// It returns the number of read bytes and any encountered error (no io.EOF)
func (w *Writer) ReadFrom(r io.Reader) (int64, error) {
if err := w.emitHeaderIfNeeded(); err != nil {
return 0, err
}
n, nprev := int64(0), -1
buf := make([]byte, defaultDecBufferSize)
// Check if a previous Write() wrote to rbuf.
if w.rbuf.Len() > 0 {
return 0, ErrMixedMethods
}
for {
nread, rerr := r.Read(buf)
if rerr != nil && rerr != io.EOF {
return n, rerr
}
n += int64(nread)
// Sanity check: check if previous block was properly aligned:
if nprev >= 0 && int64(nprev) != w.maxBlockSize && rerr != io.EOF {
return n, ErrBadBlockSize
}
if nread > 0 {
_, werr := w.flushPack(buf[:nread])
w.rbuf.Reset()
if werr != nil {
return n, werr
}
}
nprev = nread
if rerr == io.EOF {
break
}
}
return n, nil
}
// NewWriter calls NewWriterWithTypeAndBlockSize with a sane default cipher type
// and a sane default max block size.
func NewWriter(w io.Writer, key []byte) (*Writer, error) {
return NewWriterWithType(w, key, defaultCipherType)
}
// NewWriterWithType calls NewWriterWithTypeAndBlockSize with a a sane default maxblocksize.
func NewWriterWithType(w io.Writer, key []byte, cipherType uint16) (*Writer, error) {
return NewWriterWithTypeAndBlockSize(w, key, cipherType, defaultMaxBlockSize)
}
// NewWriterWithTypeAndBlockSize returns a new Writer which encrypts data with a
// certain key. If `compressionFlag` is true, the compression
// flag in the file header will also be true. Otherwise no compression is done.
func NewWriterWithTypeAndBlockSize(w io.Writer, key []byte, cipherType uint16, maxBlockSize int64) (*Writer, error) {
ew := &Writer{
Writer: w,
rbuf: &bytes.Buffer{},
maxBlockSize: maxBlockSize,
cipher: cipherType,
}
if err := ew.initAeadCommon(key, cipherType, ew.maxBlockSize); err != nil {
return nil, err
}
return ew, nil
}