diff --git a/proxy/packet_test.go b/proxy/packet_test.go index 615ff19..0bedc35 100644 --- a/proxy/packet_test.go +++ b/proxy/packet_test.go @@ -17,11 +17,11 @@ func TestAppendMac(t *testing.T) { appended := AppendMac(testMarshalled, testMac) t.Run("Length", func(t *testing.T) { - assert.Len(t, appended, len(testMarshalled)+8+4) + assert.Len(t, appended, len(testMarshalled)+4) }) t.Run("Mac", func(t *testing.T) { - assert.Equal(t, []byte{'a', 'b', 'c', 'd'}, appended[len(testMarshalled)+8:]) + assert.Equal(t, []byte{'a', 'b', 'c', 'd'}, appended[len(testMarshalled):]) }) t.Run("Original", func(t *testing.T) { diff --git a/tcp/flow_test.go b/tcp/flow_test.go index eba5a0a..559fa81 100644 --- a/tcp/flow_test.go +++ b/tcp/flow_test.go @@ -13,41 +13,75 @@ import ( func TestFlow_Consume(t *testing.T) { testContent := []byte("A test string is the content of this packet.") testPacket := proxy.SimplePacket(testContent) - testMac := mocks.AlmostUselessMac{} + testMac := mocks.AlmostUselessMac("abcd") + testMac2 := mocks.AlmostUselessMac("efgh") t.Run("Length", func(t *testing.T) { testConn := mocks.NewMockPerfectBiStreamConn(100) - flowA := NewFlowConn(context.Background(), testConn.SideA()) + flowA := NewFlowConn(context.Background(), testConn.SideA(), []proxy.MacVerifier{testMac}, []proxy.MacGenerator{testMac}) - err := flowA.Consume(context.Background(), testPacket, testMac) + err := flowA.Consume(context.Background(), testPacket) require.Nil(t, err) buf := make([]byte, 100) n, err := testConn.SideB().Read(buf) require.Nil(t, err) - assert.Equal(t, len(testContent)+8+4+4, n) - assert.Equal(t, uint32(len(testContent)+8+4), binary.LittleEndian.Uint32(buf[:len(buf)-4])) + assert.Equal(t, len(testContent)+4+4, n) + assert.Equal(t, uint32(len(testContent)+4), binary.LittleEndian.Uint32(buf[:len(buf)-4])) + }) + + t.Run("MultipleGeneratorsLength", func(t *testing.T) { + testConn := mocks.NewMockPerfectBiStreamConn(100) + + flowA := NewFlowConn(context.Background(), testConn.SideA(), []proxy.MacVerifier{testMac, testMac2}, []proxy.MacGenerator{testMac, testMac2}) + + err := flowA.Consume(context.Background(), testPacket) + require.Nil(t, err) + + buf := make([]byte, 100) + n, err := testConn.SideB().Read(buf) + require.Nil(t, err) + + assert.Equal(t, len(testContent)+4+4+4, n) + assert.Equal(t, uint32(len(testContent)+4+4), binary.LittleEndian.Uint32(buf[:len(buf)-4])) + }) + + t.Run("MultipleGeneratorsOrder", func(t *testing.T) { + testConn := mocks.NewMockPerfectBiStreamConn(100) + + flowA := NewFlowConn(context.Background(), testConn.SideA(), []proxy.MacVerifier{testMac, testMac2}, []proxy.MacGenerator{testMac, testMac2}) + + err := flowA.Consume(context.Background(), testPacket) + require.Nil(t, err) + + buf := make([]byte, 100) + n, err := testConn.SideB().Read(buf) + require.Nil(t, err) + + assert.Equal(t, len(testContent)+4+4+4, n) + assert.Equal(t, "abcdefgh", string(buf[n-8:n])) }) } func TestFlow_Produce(t *testing.T) { testContent := "A test string is the content of this packet." - testMarshalled := []byte("0000" + testContent + "00000000abcd") + testMarshalled := []byte("0000" + testContent + "abcd") binary.LittleEndian.PutUint32(testMarshalled, uint32(len(testMarshalled)-4)) - testMac := mocks.AlmostUselessMac{} + testMac := mocks.AlmostUselessMac("abcd") + testMac2 := mocks.AlmostUselessMac("efgh") t.Run("Length", func(t *testing.T) { testConn := mocks.NewMockPerfectBiStreamConn(100) - flowA := NewFlowConn(context.Background(), testConn.SideA()) + flowA := NewFlowConn(context.Background(), testConn.SideA(), []proxy.MacVerifier{testMac}, []proxy.MacGenerator{testMac}) _, err := testConn.SideB().Write(testMarshalled) require.Nil(t, err) - p, err := flowA.Produce(context.Background(), testMac) + p, err := flowA.Produce(context.Background()) require.Nil(t, err) assert.Equal(t, len(testContent), len(p.Contents())) }) @@ -55,12 +89,29 @@ func TestFlow_Produce(t *testing.T) { t.Run("Value", func(t *testing.T) { testConn := mocks.NewMockPerfectBiStreamConn(100) - flowA := NewFlowConn(context.Background(), testConn.SideA()) + flowA := NewFlowConn(context.Background(), testConn.SideA(), []proxy.MacVerifier{testMac}, []proxy.MacGenerator{testMac}) _, err := testConn.SideB().Write(testMarshalled) require.Nil(t, err) - p, err := flowA.Produce(context.Background(), testMac) + p, err := flowA.Produce(context.Background()) + require.Nil(t, err) + assert.Equal(t, testContent, string(p.Contents())) + }) + + t.Run("MultipleVerifiersStrip", func(t *testing.T) { + testContent := "A test string is the content of this packet." + testMarshalled := []byte("0000" + testContent + "abcdefgh") + binary.LittleEndian.PutUint32(testMarshalled, uint32(len(testMarshalled)-4)) + + testConn := mocks.NewMockPerfectBiStreamConn(100) + + flowA := NewFlowConn(context.Background(), testConn.SideA(), []proxy.MacVerifier{testMac, testMac2}, []proxy.MacGenerator{testMac, testMac2}) + + _, err := testConn.SideB().Write(testMarshalled) + require.Nil(t, err) + + p, err := flowA.Produce(context.Background()) require.Nil(t, err) assert.Equal(t, testContent, string(p.Contents())) })