ssh: update to use new ChaCha20 API

Use the new streaming API added in CL 104856.

Change-Id: I6453a54a9739f20fc4d05f476c6e26f720ccd354
Reviewed-on: https://go-review.googlesource.com/108656
Run-TryBot: Michael Munday <mike.munday@ibm.com>
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>

ssh/cipher.go

@@ -16,6 +16,7 @@ import (
 	"hash"
 	"io"
 	"io/ioutil"
+	"math/bits"
 
 	"golang.org/x/crypto/internal/chacha20"
 	"golang.org/x/crypto/poly1305"
@@ -641,8 +642,8 @@ const chacha20Poly1305ID = "chacha20-poly1305@openssh.com"
 // the methods here also implement padding, which RFC4253 Section 6
 // also requires of stream ciphers.
 type chacha20Poly1305Cipher struct {
-	lengthKey  [32]byte
-	contentKey [32]byte
+	lengthKey  [8]uint32
+	contentKey [8]uint32
 	buf        []byte
 }
 
@@ -655,28 +656,29 @@ func newChaCha20Cipher(key, unusedIV, unusedMACKey []byte, unusedAlgs directionA
 		buf: make([]byte, 256),
 	}
 
-	copy(c.contentKey[:], key[:32])
-	copy(c.lengthKey[:], key[32:])
+	for i := range c.contentKey {
+		c.contentKey[i] = binary.LittleEndian.Uint32(key[i*4 : (i+1)*4])
+	}
+	for i := range c.contentKey {
+		c.lengthKey[i] = binary.LittleEndian.Uint32(key[(i+8)*4 : (i+9)*4])
+	}
 	return c, nil
 }
 
-// The Poly1305 key is obtained by encrypting 32 0-bytes.
-var chacha20PolyKeyInput [32]byte
-
 func (c *chacha20Poly1305Cipher) readPacket(seqNum uint32, r io.Reader) ([]byte, error) {
-	var counter [16]byte
-	binary.BigEndian.PutUint64(counter[8:], uint64(seqNum))
-
+	nonce := [3]uint32{0, 0, bits.ReverseBytes32(seqNum)}
+	s := chacha20.New(c.contentKey, nonce)
 	var polyKey [32]byte
-	chacha20.XORKeyStream(polyKey[:], chacha20PolyKeyInput[:], &counter, &c.contentKey)
+	s.XORKeyStream(polyKey[:], polyKey[:])
+	s.Advance() // skip next 32 bytes
 
 	encryptedLength := c.buf[:4]
 	if _, err := io.ReadFull(r, encryptedLength); err != nil {
 		return nil, err
 	}
 
 	var lenBytes [4]byte
-	chacha20.XORKeyStream(lenBytes[:], encryptedLength, &counter, &c.lengthKey)
+	chacha20.New(c.lengthKey, nonce).XORKeyStream(lenBytes[:], encryptedLength)
 
 	length := binary.BigEndian.Uint32(lenBytes[:])
 	if length > maxPacket {
@@ -702,10 +704,8 @@ func (c *chacha20Poly1305Cipher) readPacket(seqNum uint32, r io.Reader) ([]byte,
 		return nil, errors.New("ssh: MAC failure")
 	}
 
-	counter[0] = 1
-
 	plain := c.buf[4:contentEnd]
-	chacha20.XORKeyStream(plain, plain, &counter, &c.contentKey)
+	s.XORKeyStream(plain, plain)
 
 	padding := plain[0]
 	if padding < 4 {
@@ -724,11 +724,11 @@ func (c *chacha20Poly1305Cipher) readPacket(seqNum uint32, r io.Reader) ([]byte,
 }
 
 func (c *chacha20Poly1305Cipher) writePacket(seqNum uint32, w io.Writer, rand io.Reader, payload []byte) error {
-	var counter [16]byte
-	binary.BigEndian.PutUint64(counter[8:], uint64(seqNum))
-
+	nonce := [3]uint32{0, 0, bits.ReverseBytes32(seqNum)}
+	s := chacha20.New(c.contentKey, nonce)
 	var polyKey [32]byte
-	chacha20.XORKeyStream(polyKey[:], chacha20PolyKeyInput[:], &counter, &c.contentKey)
+	s.XORKeyStream(polyKey[:], polyKey[:])
+	s.Advance() // skip next 32 bytes
 
 	// There is no blocksize, so fall back to multiple of 8 byte
 	// padding, as described in RFC 4253, Sec 6.
@@ -748,16 +748,15 @@ func (c *chacha20Poly1305Cipher) writePacket(seqNum uint32, w io.Writer, rand io
 	}
 
 	binary.BigEndian.PutUint32(c.buf, uint32(1+len(payload)+padding))
-	chacha20.XORKeyStream(c.buf, c.buf[:4], &counter, &c.lengthKey)
+	chacha20.New(c.lengthKey, nonce).XORKeyStream(c.buf, c.buf[:4])
 	c.buf[4] = byte(padding)
 	copy(c.buf[5:], payload)
 	packetEnd := 5 + len(payload) + padding
 	if _, err := io.ReadFull(rand, c.buf[5+len(payload):packetEnd]); err != nil {
 		return err
 	}
 
-	counter[0] = 1
-	chacha20.XORKeyStream(c.buf[4:], c.buf[4:packetEnd], &counter, &c.contentKey)
+	s.XORKeyStream(c.buf[4:], c.buf[4:packetEnd])
 
 	var mac [poly1305.TagSize]byte
 	poly1305.Sum(&mac, c.buf[:packetEnd], &polyKey)