1package main
2
3import (
4	"crypto/cipher"
5	"crypto/subtle"
6	"encoding/binary"
7	"errors"
8)
9
10// See draft-agl-tls-chacha20poly1305-04 and
11// draft-irtf-cfrg-chacha20-poly1305-10. Where the two differ, the
12// draft-agl-tls-chacha20poly1305-04 variant is implemented.
13
14func leftRotate(a uint32, n uint) uint32 {
15	return (a << n) | (a >> (32 - n))
16}
17
18func chaChaQuarterRound(state *[16]uint32, a, b, c, d int) {
19	state[a] += state[b]
20	state[d] = leftRotate(state[d]^state[a], 16)
21
22	state[c] += state[d]
23	state[b] = leftRotate(state[b]^state[c], 12)
24
25	state[a] += state[b]
26	state[d] = leftRotate(state[d]^state[a], 8)
27
28	state[c] += state[d]
29	state[b] = leftRotate(state[b]^state[c], 7)
30}
31
32func chaCha20Block(state *[16]uint32, out []byte) {
33	var workingState [16]uint32
34	copy(workingState[:], state[:])
35	for i := 0; i < 10; i++ {
36		chaChaQuarterRound(&workingState, 0, 4, 8, 12)
37		chaChaQuarterRound(&workingState, 1, 5, 9, 13)
38		chaChaQuarterRound(&workingState, 2, 6, 10, 14)
39		chaChaQuarterRound(&workingState, 3, 7, 11, 15)
40		chaChaQuarterRound(&workingState, 0, 5, 10, 15)
41		chaChaQuarterRound(&workingState, 1, 6, 11, 12)
42		chaChaQuarterRound(&workingState, 2, 7, 8, 13)
43		chaChaQuarterRound(&workingState, 3, 4, 9, 14)
44	}
45	for i := 0; i < 16; i++ {
46		binary.LittleEndian.PutUint32(out[i*4:i*4+4], workingState[i]+state[i])
47	}
48}
49
50// sliceForAppend takes a slice and a requested number of bytes. It returns a
51// slice with the contents of the given slice followed by that many bytes and a
52// second slice that aliases into it and contains only the extra bytes. If the
53// original slice has sufficient capacity then no allocation is performed.
54func sliceForAppend(in []byte, n int) (head, tail []byte) {
55	if total := len(in) + n; cap(in) >= total {
56		head = in[:total]
57	} else {
58		head = make([]byte, total)
59		copy(head, in)
60	}
61	tail = head[len(in):]
62	return
63}
64
65type chaCha20Poly1305 struct {
66	key [32]byte
67}
68
69func newChaCha20Poly1305(key []byte) (cipher.AEAD, error) {
70	if len(key) != 32 {
71		return nil, errors.New("bad key length")
72	}
73	aead := new(chaCha20Poly1305)
74	copy(aead.key[:], key)
75	return aead, nil
76}
77
78func (c *chaCha20Poly1305) NonceSize() int { return 8 }
79func (c *chaCha20Poly1305) Overhead() int  { return 16 }
80
81func (c *chaCha20Poly1305) chaCha20(out, in, nonce []byte, counter uint64) {
82	var state [16]uint32
83	state[0] = 0x61707865
84	state[1] = 0x3320646e
85	state[2] = 0x79622d32
86	state[3] = 0x6b206574
87	for i := 0; i < 8; i++ {
88		state[4+i] = binary.LittleEndian.Uint32(c.key[i*4 : i*4+4])
89	}
90	state[14] = binary.LittleEndian.Uint32(nonce[0:4])
91	state[15] = binary.LittleEndian.Uint32(nonce[4:8])
92
93	for i := 0; i < len(in); i += 64 {
94		state[12] = uint32(counter & 0xffffffff)
95		state[13] = uint32(counter >> 32)
96
97		var tmp [64]byte
98		chaCha20Block(&state, tmp[:])
99		count := 64
100		if len(in)-i < count {
101			count = len(in) - i
102		}
103		for j := 0; j < count; j++ {
104			out[i+j] = in[i+j] ^ tmp[j]
105		}
106
107		counter++
108	}
109}
110
111func (c *chaCha20Poly1305) poly1305(tag *[16]byte, nonce, ciphertext, additionalData []byte) {
112	input := make([]byte, 0, len(additionalData)+8+len(ciphertext)+8)
113	input = append(input, additionalData...)
114	input, out := sliceForAppend(input, 8)
115	binary.LittleEndian.PutUint64(out, uint64(len(additionalData)))
116	input = append(input, ciphertext...)
117	input, out = sliceForAppend(input, 8)
118	binary.LittleEndian.PutUint64(out, uint64(len(ciphertext)))
119
120	var poly1305Key [32]byte
121	c.chaCha20(poly1305Key[:], poly1305Key[:], nonce, 0)
122
123	poly1305Sum(tag, input, &poly1305Key)
124}
125
126func (c *chaCha20Poly1305) Seal(dst, nonce, plaintext, additionalData []byte) []byte {
127	if len(nonce) != 8 {
128		panic("Bad nonce length")
129	}
130
131	ret, out := sliceForAppend(dst, len(plaintext)+16)
132	c.chaCha20(out[:len(plaintext)], plaintext, nonce, 1)
133
134	var tag [16]byte
135	c.poly1305(&tag, nonce, out[:len(plaintext)], additionalData)
136	copy(out[len(plaintext):], tag[:])
137
138	return ret
139}
140
141func (c *chaCha20Poly1305) Open(dst, nonce, ciphertext, additionalData []byte) ([]byte, error) {
142	if len(nonce) != 8 {
143		panic("Bad nonce length")
144	}
145	if len(ciphertext) < 16 {
146		return nil, errors.New("chacha20: message authentication failed")
147	}
148	plaintextLen := len(ciphertext) - 16
149
150	var tag [16]byte
151	c.poly1305(&tag, nonce, ciphertext[:plaintextLen], additionalData)
152	if subtle.ConstantTimeCompare(tag[:], ciphertext[plaintextLen:]) != 1 {
153		return nil, errors.New("chacha20: message authentication failed")
154	}
155
156	ret, out := sliceForAppend(dst, plaintextLen)
157	c.chaCha20(out, ciphertext[:plaintextLen], nonce, 1)
158	return ret, nil
159}
160