handshake: rewrite sequences for clarity and correctness

Author: ttaylorr

handshake/ack_sequence.go

@@ -0,0 +1,70 @@
+package handshake
+
+import (
+	"crypto/rand"
+	"io"
+)
+
+// ClientAckSequence is a handshake sequence that verifies the challenge
+// sequence sent by the client during a RTMP handshake. It is responsible for
+// reading and responding to the C1 packet (with the C2 packet), and sending the
+// server challenge in S1.
+type ClientAckSequence struct {
+	C1 *AckPacket
+	S1 *AckPacket
+}
+
+var _ Sequence = new(ClientAckSequence)
+
+// NewClientAckSequence initializes and returns a new *ClientAckSequence
+// initialized with an empty C1 packet and a new S1 packet, initialized with the
+// result of a "rand.Read" into its Payload header.
+func NewClientAckSequence() *ClientAckSequence {
+	c := &ClientAckSequence{
+		C1: new(AckPacket),
+		S1: new(AckPacket),
+	}
+
+	rand.Read(c.S1.Payload[:])
+
+	return c
+}
+
+// Read implements the Sequence.Read function. It reads the C1 packet and
+// returns any read error, if there was one. Otherwise, a value of "nil" is
+// returned instead.
+func (c *ClientAckSequence) Read(r io.Reader) error {
+	if err := c.C1.Read(r); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// Write implements the Sequence.Write function. It writes the S1 packet first
+// (returning any errors if there is one), and then writes the S2 packet with
+// the same data as was sent in the C1 packet (returning any error that was
+// encountered).
+//
+// A successful call to Write constitutes a value of `nil` being returned.
+func (c *ClientAckSequence) Write(w io.Writer) error {
+	if err := c.S1.Write(w); err != nil {
+		return err
+	}
+
+	s2 := &AckPacket{
+		Time1:   c.C1.Time1,
+		Payload: c.C1.Payload,
+	}
+
+	if err := s2.Write(w); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// Nex implements the Sequence.Next function.
+func (c *ClientAckSequence) Next() Sequence {
+	return NewServerAckSequence(c.S1)
+}

handshake/ack_sequence_test.go

@@ -0,0 +1,67 @@
+package handshake_test
+
+import (
+	"bytes"
+	"crypto/rand"
+	"io"
+	"testing"
+
+	"github.com/WatchBeam/rtmp/handshake"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestItReadsC1Packet(t *testing.T) {
+	payload := payload()
+
+	buf := bytes.NewBuffer([]byte{})
+	buf.Write([]byte{0x0, 0x0, 0x0, 0x1}) // Time 1
+	buf.Write([]byte{0x0, 0x0, 0x0, 0x0}) // Padding
+	buf.Write(payload[:])
+
+	c := handshake.NewClientAckSequence()
+	err := c.Read(buf)
+
+	assert.Nil(t, err)
+	assert.Equal(t, &handshake.AckPacket{
+		Time1:   1,
+		Time2:   0,
+		Payload: payload,
+	}, c.C1)
+}
+
+func TestItErorrsOnBadC1Packets(t *testing.T) {
+	buf := bytes.NewBuffer([]byte{
+	// Empty C1 ~> io.EOF
+	})
+
+	c := handshake.NewClientAckSequence()
+	err := c.Read(buf)
+
+	assert.Equal(t, io.EOF, err)
+}
+
+func TestItWritesS1AndMatchingS2(t *testing.T) {
+	buf := bytes.NewBuffer([]byte{})
+
+	c := handshake.NewClientAckSequence()
+	c.C1 = &handshake.AckPacket{
+		Payload: payload(),
+	}
+
+	start := 4 + 4 + 1528 + 4 + 4
+	end := start + 1528
+
+	err := c.Write(buf)
+
+	assert.Nil(t, err)
+	assert.Len(t, buf.Bytes(), 2*(4+4+1528))
+	assert.Equal(t, c.C1.Payload[:], buf.Bytes()[start:end])
+}
+
+func payload() [handshake.PayloadLen]byte {
+	var b [handshake.PayloadLen]byte
+
+	rand.Read(b[:])
+
+	return b
+}

handshake/client_ack_sequence.go

@@ -0,0 +1,56 @@
+package handshake
+
+import (
+	"bytes"
+	"errors"
+	"io"
+)
+
+var (
+	// MismatchedChallengeErr is an error which is returned in a situtation
+	// when the challenge sequence read from the client in C2 differs from
+	// the challenge sequence in S1.
+	MismatchedChallengeErr = errors.New("rtmp/handshake: mismatched challenge")
+)
+
+// ServerAckSequence is a type implementing the handshake.Sequence interface and
+// is repsonseible for reading and verifying the C2 packet written by the client
+// against the S1 packet from the server.
+type ServerAckSequence struct {
+	// S1 is the packet which C2 should acknowledge.
+	S1 *AckPacket
+}
+
+var _ Sequence = new(ServerAckSequence)
+
+// NewServerAckSequence returns a new *ServerAckSequence initialized with the
+// given S1 packet.
+func NewServerAckSequence(S1 *AckPacket) *ServerAckSequence {
+	return &ServerAckSequence{S1}
+}
+
+// Read implements the Handshake.Read method by reading the C2 packet and
+// comparing it to the stored S1 packet. If a read error occured while reading
+// C2, then it will be returned. If the payloads were not equal, then
+// MismatchedChallengeErr will be returned. Otherwise, in the successful case,
+// a value of nil will be returned.
+func (s *ServerAckSequence) Read(r io.Reader) error {
+	c2 := new(AckPacket)
+	if err := c2.Read(r); err != nil {
+		return err
+	}
+
+	if !bytes.Equal(s.S1.Payload[:], c2.Payload[:]) {
+		return MismatchedChallengeErr
+	}
+
+	return nil
+}
+
+// Write implements the Sequence.Write function. Since there is nothing to
+// write, a value of nil is always returned here.
+func (s *ServerAckSequence) Write(w io.Writer) error { return nil }
+
+// Next implements the Sequence.Next function. Since there is no next function,
+// this function always returns nil.
+func (s *ServerAckSequence) Next() Sequence { return nil }