implement writing of crypto messages

crypto_stream.go

@@ -1,8 +1,10 @@
 package quic
 
 import (
+	"bytes"
 	"encoding/binary"
 	"errors"
+	"sort"
 )
 
 // A Tag in the QUIC crypto
@@ -87,3 +89,37 @@ func ParseCryptoMessage(data []byte) (Tag, map[Tag][]byte, error) {
 
 	return messageTag, resultMap, nil
 }
+
+// WriteCryptoMessage writes a crypto message
+func WriteCryptoMessage(b *bytes.Buffer, messageTag Tag, data map[Tag][]byte) {
+	writeUint32(b, uint32(messageTag))
+	writeUint16(b, uint16(len(data)))
+	writeUint16(b, 0)
+
+	// Save current position in the buffer, so that we can update the index in-place later
+	indexStart := b.Len()
+
+	indexData := make([]byte, 8*len(data))
+	b.Write(indexData) // Will be updated later
+
+	// Sort the tags
+	tags := make([]uint32, len(data))
+	i := 0
+	for t := range data {
+		tags[i] = uint32(t)
+		i++
+	}
+	sort.Sort(Uint32Slice(tags))
+
+	offset := uint32(0)
+	for i, t := range tags {
+		v := data[Tag(t)]
+		b.Write(v)
+		offset += uint32(len(v))
+		binary.LittleEndian.PutUint32(indexData[i*8:], t)
+		binary.LittleEndian.PutUint32(indexData[i*8+4:], offset)
+	}
+
+	// Now we write the index data for real
+	copy(b.Bytes()[indexStart:], indexData)
+}

crypto_stream_test.go

@@ -1,7 +1,7 @@
 package quic
 
 import (
-	"strings"
+	"bytes"
 
 	. "github.com/onsi/ginkgo"
 	. "github.com/onsi/gomega"
@@ -13,24 +13,15 @@ var _ = Describe("CryptoStream", func() {
 			tag, msg, err := ParseCryptoMessage(sampleCHLO)
 			Expect(err).ToNot(HaveOccurred())
 			Expect(tag).To(Equal(TagCHLO))
-			Expect(msg).To(Equal(map[Tag][]byte{
-				TagPAD:  []byte(strings.Repeat("-", 1016)),
-				TagSNI:  []byte("www.example.org"),
-				TagVER:  []byte("Q030"),
-				TagCCS:  []byte("{&\xe9\xe7\xe4\\q\xff\x01\xe8\x81`\x92\x92\x1a\xe8"),
-				TagMSPC: []byte("d\x00\x00\x00"),
-				TagUAID: []byte("dev Chrome/51.0.2700.0 Intel Mac OS X 10_11_4"),
-				TagTCID: []byte("\x00\x00\x00\x00"),
-				TagSRBF: []byte("\x00\x00\x10\x00"),
-				TagICSL: []byte("\x1e\x00\x00\x00"),
-				TagNONP: []byte("\xe1\x84T\x1b\xe3\xd6|\x1fi\xb2N\x9eF\xf4Fݫ\xe5\xdef\x94\xf6\xb2\xee\x01ĥw\xfe\xc9\v\xa3"),
-				TagSCLS: []byte("\x01\x00\x00\x00"),
-				TagCSCT: []byte{},
-				TagCOPT: []byte("FIXD"),
-				TagSFCW: []byte("\x00\x00`\x00"),
-				TagCFCW: []byte("\x00\x00\xf0\x00"),
-				TagPDMD: []byte("X509"),
-			}))
+			Expect(msg).To(Equal(sampleCHLOMap))
+		})
+	})
+
+	Context("when writing", func() {
+		It("writes sample message", func() {
+			b := &bytes.Buffer{}
+			WriteCryptoMessage(b, TagCHLO, sampleCHLOMap)
+			Expect(b.Bytes()).To(Equal(sampleCHLO))
 		})
 	})
 })

utils.go

@@ -1,6 +1,9 @@
 package quic
 
-import "io"
+import (
+	"bytes"
+	"io"
+)
 
 func readUintN(b io.ByteReader, length uint8) (uint64, error) {
 	var res uint64
@@ -43,3 +46,22 @@ func readUint16(b io.ByteReader) (uint16, error) {
 	}
 	return uint16(b1) + uint16(b2)<<8, nil
 }
+
+func writeUint32(b *bytes.Buffer, i uint32) {
+	b.WriteByte(uint8(i & 0xff))
+	b.WriteByte(uint8((i >> 8) & 0xff))
+	b.WriteByte(uint8((i >> 16) & 0xff))
+	b.WriteByte(uint8((i >> 24) & 0xff))
+}
+
+func writeUint16(b *bytes.Buffer, i uint16) {
+	b.WriteByte(uint8(i & 0xff))
+	b.WriteByte(uint8((i >> 8) & 0xff))
+}
+
+// Uint32Slice attaches the methods of sort.Interface to []uint32, sorting in increasing order.
+type Uint32Slice []uint32
+
+func (s Uint32Slice) Len() int           { return len(s) }
+func (s Uint32Slice) Less(i, j int) bool { return s[i] < s[j] }
+func (s Uint32Slice) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }