forked from keybase/client
-
Notifications
You must be signed in to change notification settings - Fork 0
/
armor.go
275 lines (249 loc) · 7.27 KB
/
armor.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
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.
package kbcmf
import (
"bytes"
"fmt"
basex "github.com/keybase/client/go/encoding/basex"
"io"
"io/ioutil"
"strings"
)
// ArmorParams specify armor formatting, encoding and punctuation.
type ArmorParams struct {
// BytesPerWord is the number of characters in each "word" of output.
// We'll put spaces between words.
BytesPerWord int
// WordsPerLine is the number of words for each line of output. We'll
// put newlines between two subsequent lines of output.
WordsPerLine int
// Punctuation is the byte inserted after the three "sentences" of
// our encoding -- the header, the body and the footer.
Punctuation byte
// Encoding is the basex encoding to use, including strictness parameters
Encoding *basex.Encoding
}
type armorEncoderStream struct {
buf *bytes.Buffer
footer string
encoded io.Writer
encoder io.WriteCloser
nWords int
params ArmorParams
}
func (s *armorEncoderStream) Write(b []byte) (n int, err error) {
n, err = s.encoder.Write(b)
if err != nil {
return n, err
}
s.spaceAndOutputBuffer()
return n, err
}
func (s *armorEncoderStream) spaceAndOutputBuffer() error {
for s.buf.Len() > s.params.BytesPerWord {
buf := s.buf.Next(s.params.BytesPerWord)
s.nWords++
sep := byte(' ')
if s.nWords%s.params.WordsPerLine == 0 {
sep = byte('\n')
}
if _, err := s.encoded.Write(buf); err != nil {
return err
}
if _, err := s.encoded.Write([]byte{sep}); err != nil {
return err
}
}
return nil
}
func (s *armorEncoderStream) Close() (err error) {
if err = s.encoder.Close(); err != nil {
return err
}
if err := s.spaceAndOutputBuffer(); err != nil {
return err
}
lst := s.buf.Bytes()
if _, err := s.encoded.Write([]byte(lst)); err != nil {
return err
}
s.nWords++
pad := ""
if len(lst) == s.params.BytesPerWord {
if s.nWords%s.params.WordsPerLine == 0 {
pad = "\n"
} else {
pad = " "
}
}
if _, err := fmt.Fprintf(s.encoded, "%s%c\n\n%s%c\n", pad, s.params.Punctuation, s.footer, s.params.Punctuation); err != nil {
return err
}
return nil
}
// NewArmorEncoderStream makes a new Armor encoding stream, using the given encoding
// Pass it an `encoded` stream writer to write the
// encoded stream to. Also pass a header, and a footer string. It will
// return an io.WriteCloser on success, that you can write raw (unencoded) data to.
// An error will be returned if there is trouble writing the header to encoded.
//
// To make the output look pretty, a space is inserted every 15 characters of output,
// and a newline is inserted every 200 words.
func NewArmorEncoderStream(encoded io.Writer, header string, footer string, params ArmorParams) (io.WriteCloser, error) {
ret := &armorEncoderStream{
buf: new(bytes.Buffer),
encoded: encoded,
footer: footer,
params: params,
}
ret.encoder = basex.NewEncoder(params.Encoding, ret.buf)
if _, err := fmt.Fprintf(encoded, "%s%c\n\n", header, params.Punctuation); err != nil {
return nil, err
}
return ret, nil
}
// ArmorSeal takes an input plaintext and returns and output armor encoding
// as a string, or an error if a problem was encountered. Also provide a header
// and a footer to frame the message.
func ArmorSeal(plaintext []byte, header string, footer string, params ArmorParams) (string, error) {
var buf bytes.Buffer
enc, err := NewArmorEncoderStream(&buf, header, footer, params)
if err != nil {
return "", err
}
if _, err := enc.Write(plaintext); err != nil {
return "", err
}
if err := enc.Close(); err != nil {
return "", err
}
return buf.String(), nil
}
// Frame is a way to read the frame out of a Decoder stream.
type Frame interface {
// GetHeader() returns the frame associated with this stream, or an error
GetHeader() (string, error)
// GetFooter() returns the frame associated with this stream, or an error
GetFooter() (string, error)
}
type fdsState int
const (
fdsHeader fdsState = iota
fdsBody
fdsFooter
fdsEndOfStream
)
type framedDecoderStream struct {
header []byte
footer []byte
state fdsState
params ArmorParams
r *PunctuatedReader
}
// Read from a framedDeecoderStream. The frame is the "BEGIN FOO." block
// at the footer, and the "END FOO." block at the end.
func (s *framedDecoderStream) Read(p []byte) (n int, err error) {
// The largest frame we'll accept before we show an overflow.
frameLim := 8192
if s.state == fdsHeader {
s.header, err = s.r.ReadUntilPunctuation(frameLim)
if err != nil {
return 0, err
}
s.state = fdsBody
}
if s.state == fdsBody {
n, err = s.r.Read(p)
if err == ErrPunctuated {
err = nil
s.state = fdsFooter
}
if err == io.EOF {
err = io.ErrUnexpectedEOF
}
if err != nil {
return 0, err
}
}
if s.state == fdsFooter {
s.footer, err = s.r.ReadUntilPunctuation(frameLim)
if err != nil {
return 0, err
}
s.state = fdsEndOfStream
}
if s.state == fdsEndOfStream {
err = s.consumeUntilEOF()
if err == io.EOF && n > 0 {
err = nil
}
}
return n, err
}
// consume the stream until we hit an EOF. For all data we consume, make
// sure that it's a valid byte as far as our underlying decoder is concerned.
// We might considering clamping down here on the number of characters we're willing
// to accept after the message is over. But for now, we're quite liberal.
func (s *framedDecoderStream) consumeUntilEOF() error {
var buf [4096]byte
for {
n, err := s.r.Read(buf[:])
if err != nil {
return err
}
if n == 0 {
return io.EOF
}
if !s.isValidByteSequence(buf[0:n]) {
return ErrTrailingGarbage
}
}
}
// isValidByteSequence checks if the byte sequence is valid as far as our
// underlying encoder is concerned. This is usually quite liberal, and excludes
// non-ASCII and ASCII control characters, but leaves everything else in.
func (s *framedDecoderStream) isValidByteSequence(p []byte) bool {
for _, b := range p {
if !s.params.Encoding.IsValidByte(b) {
return false
}
}
return true
}
func (s *framedDecoderStream) toASCII(buf []byte) (string, error) {
if !s.isValidByteSequence(buf) {
return "", ErrBadArmorFrame
}
return strings.TrimSpace(string(buf)), nil
}
func (s *framedDecoderStream) GetFooter() (string, error) { return s.toASCII(s.footer) }
func (s *framedDecoderStream) GetHeader() (string, error) { return s.toASCII(s.header) }
// NewArmorDecoderStream is used to decode armored encoding. It returns a stream you
// can read from, and also a Frame you can query to see what the open/close
// frame markers were.
func NewArmorDecoderStream(r io.Reader, params ArmorParams) (io.Reader, Frame, error) {
fds := &framedDecoderStream{r: NewPunctuatedReader(r, params.Punctuation), params: params}
ret := basex.NewDecoder(params.Encoding, fds)
return ret, fds, nil
}
// ArmorOpen runs armor stream decoding, but on a string, and it outputs a string.
func ArmorOpen(msg string, params ArmorParams) (body []byte, header string, footer string, err error) {
var dec io.Reader
var frame Frame
buf := bytes.NewBufferString(msg)
dec, frame, err = NewArmorDecoderStream(buf, params)
if err != nil {
return
}
body, err = ioutil.ReadAll(dec)
if err != nil {
return
}
if header, err = frame.GetHeader(); err != nil {
return
}
if footer, err = frame.GetFooter(); err != nil {
return
}
return body, header, footer, nil
}