feature: Ping/Pong stream flags

Signed-off-by: Valery Piashchynski <[email protected]>

Author: rustatian

pkg/frame/frame.go

@@ -14,7 +14,6 @@ const WORD = 4
 type Frame struct {
 	// Payload, max length 4.2GB.
 	payload []byte
-
 	// Header
 	header []byte
 }
@@ -128,10 +127,31 @@ func (*Frame) SetStreamFlag(header []byte) {
 
 func (*Frame) IsStream(header []byte) bool {
 	_ = header[11]
-	return header[10]&STREAM == 1
+	return header[10]&STREAM != 0
+}
+
+func (*Frame) IsPing(header []byte) bool {
+	_ = header[11]
+	return header[10]&PING != 0
+}
+
+func (*Frame) SetPingBit(header []byte) {
+	_ = header[11]
+	header[10] |= PING
+}
+
+func (*Frame) IsPong(header []byte) bool {
+	_ = header[11]
+	return header[10]&PONG != 0
+}
+
+func (*Frame) SetPongBit(header []byte) {
+	_ = header[11]
+	header[10] |= PONG
 }
 
 func (*Frame) SetStopBit(header []byte) {
+	_ = header[11]
 	header[10] |= STOP
 }
 

pkg/frame/frame_flags.go

@@ -15,11 +15,15 @@ const (
 	Version1 byte = 0x01
 
 	/*
-		10th byte
+		10th byte, stream
 	*/
 
 	// STREAM bit
 	STREAM byte = 0x01
 	// STOP command
 	STOP byte = 0x02
+	// PING command
+	PING byte = 0x04
+	// PONG command
+	PONG byte = 0x08
 )

pkg/frame/frame_test.go

@@ -296,6 +296,81 @@ func TestFrame_Bytes(t *testing.T) {
 	assert.Equal(t, []uint32{323423432}, rf.ReadOptions(rf.Header()))
 }
 
+func TestFrame_NotPingPong(t *testing.T) {
+	nf := NewFrame()
+	nf.WriteVersion(nf.Header(), 1)
+	nf.WriteFlags(nf.Header(), CONTROL, CodecGob)
+	nf.WritePayloadLen(nf.Header(), uint32(len([]byte(TestPayload))))
+
+	nf.WriteOptions(nf.HeaderPtr(), 323423432)
+	assert.Equal(t, []uint32{323423432}, nf.ReadOptions(nf.Header()))
+	nf.WritePayload([]byte(TestPayload))
+
+	nf.WriteCRC(nf.Header())
+	assert.Equal(t, true, nf.VerifyCRC(nf.Header()))
+	data := nf.Bytes()
+
+	rf := ReadFrame(data)
+
+	assert.False(t, rf.IsPing(rf.Header()))
+	assert.False(t, rf.IsPong(rf.Header()))
+	assert.Equal(t, rf.ReadVersion(rf.Header()), nf.ReadVersion(nf.Header()))
+	assert.Equal(t, rf.ReadFlags(), nf.ReadFlags())
+	assert.Equal(t, rf.ReadPayloadLen(rf.Header()), nf.ReadPayloadLen(nf.Header()))
+	assert.Equal(t, true, rf.VerifyCRC(rf.Header()))
+	assert.Equal(t, []uint32{323423432}, rf.ReadOptions(rf.Header()))
+}
+
+func TestFrame_Ping(t *testing.T) {
+	nf := NewFrame()
+	nf.WriteVersion(nf.Header(), 1)
+	nf.WriteFlags(nf.Header(), CONTROL, CodecGob)
+	nf.WritePayloadLen(nf.Header(), uint32(len([]byte(TestPayload))))
+
+	nf.WriteOptions(nf.HeaderPtr(), 323423432)
+	assert.Equal(t, []uint32{323423432}, nf.ReadOptions(nf.Header()))
+	nf.WritePayload([]byte(TestPayload))
+	nf.SetPingBit(nf.Header())
+
+	nf.WriteCRC(nf.Header())
+	assert.Equal(t, true, nf.VerifyCRC(nf.Header()))
+	data := nf.Bytes()
+
+	rf := ReadFrame(data)
+
+	assert.True(t, rf.IsPing(rf.Header()))
+	assert.Equal(t, rf.ReadVersion(rf.Header()), nf.ReadVersion(nf.Header()))
+	assert.Equal(t, rf.ReadFlags(), nf.ReadFlags())
+	assert.Equal(t, rf.ReadPayloadLen(rf.Header()), nf.ReadPayloadLen(nf.Header()))
+	assert.Equal(t, true, rf.VerifyCRC(rf.Header()))
+	assert.Equal(t, []uint32{323423432}, rf.ReadOptions(rf.Header()))
+}
+
+func TestFrame_Pong(t *testing.T) {
+	nf := NewFrame()
+	nf.WriteVersion(nf.Header(), 1)
+	nf.WriteFlags(nf.Header(), CONTROL, CodecGob)
+	nf.WritePayloadLen(nf.Header(), uint32(len([]byte(TestPayload))))
+
+	nf.WriteOptions(nf.HeaderPtr(), 323423432)
+	assert.Equal(t, []uint32{323423432}, nf.ReadOptions(nf.Header()))
+	nf.WritePayload([]byte(TestPayload))
+	nf.SetPongBit(nf.Header())
+
+	nf.WriteCRC(nf.Header())
+	assert.Equal(t, true, nf.VerifyCRC(nf.Header()))
+	data := nf.Bytes()
+
+	rf := ReadFrame(data)
+
+	assert.True(t, rf.IsPong(rf.Header()))
+	assert.Equal(t, rf.ReadVersion(rf.Header()), nf.ReadVersion(nf.Header()))
+	assert.Equal(t, rf.ReadFlags(), nf.ReadFlags())
+	assert.Equal(t, rf.ReadPayloadLen(rf.Header()), nf.ReadPayloadLen(nf.Header()))
+	assert.Equal(t, true, rf.VerifyCRC(rf.Header()))
+	assert.Equal(t, []uint32{323423432}, rf.ReadOptions(rf.Header()))
+}
+
 func BenchmarkCRC32(b *testing.B) {
 	b.ReportAllocs()
 	b.ResetTimer()