diff --git a/mocks/mac.go b/mocks/mac.go index 5d7ab7a..caf99b3 100644 --- a/mocks/mac.go +++ b/mocks/mac.go @@ -4,19 +4,22 @@ import ( "mpbl3p/shared" ) -type AlmostUselessMac struct{} +type AlmostUselessMac string -func (AlmostUselessMac) CodeLength() int { - return 4 +func (a AlmostUselessMac) CodeLength() int { + return len(a) } -func (AlmostUselessMac) Generate([]byte) []byte { - return []byte{'a', 'b', 'c', 'd'} +func (a AlmostUselessMac) Generate([]byte) []byte { + return []byte(a) } -func (u AlmostUselessMac) Verify(_, sum []byte) error { - if !(sum[0] == 'a' && sum[1] == 'b' && sum[2] == 'c' && sum[3] == 'd') { - return shared.ErrBadChecksum +func (a AlmostUselessMac) Verify(_, sum []byte) error { + for i, c := range sum { + if a[i] != c { + return shared.ErrBadChecksum + } } + return nil } diff --git a/tcp/flow.go b/tcp/flow.go index a1fe39e..c0cdcba 100644 --- a/tcp/flow.go +++ b/tcp/flow.go @@ -194,9 +194,10 @@ func (f *Flow) Produce(ctx context.Context) (proxy.Packet, error) { return nil, err } - for _, v := range f.verifiers { - var err error + for i := range f.verifiers { + v := f.verifiers[len(f.verifiers)-i-1] + var err error data, err = proxy.StripMac(data, v) if err != nil { return nil, err diff --git a/udp/flow.go b/udp/flow.go index 5d54745..6ef2666 100644 --- a/udp/flow.go +++ b/udp/flow.go @@ -226,7 +226,9 @@ func (f *Flow) produceInternal(ctx context.Context, mustReturn bool) (proxy.Pack return nil, ctx.Err() } - for _, v := range f.verifiers { + for i := range f.verifiers { + v := f.verifiers[len(f.verifiers)-i-1] + var err error received, err = proxy.StripMac(received, v) if err != nil { diff --git a/udp/flow_test.go b/udp/flow_test.go index d044477..ee242f0 100644 --- a/udp/flow_test.go +++ b/udp/flow_test.go @@ -15,17 +15,17 @@ 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") - t.Run("Length", func(t *testing.T) { + t.Run("SingleGeneratorLength", func(t *testing.T) { testConn := mocks.NewMockPerfectBiPacketConn(10) - flowA := newFlow(congestion.NewNone(), testMac) + flowA := newFlow(congestion.NewNone(), []proxy.MacVerifier{testMac}, []proxy.MacGenerator{testMac}) flowA.writer = testConn.SideB() flowA.isAlive = true - err := flowA.Consume(context.Background(), testPacket, testMac) + err := flowA.Consume(context.Background(), testPacket) require.Nil(t, err) buf := make([]byte, 100) @@ -33,7 +33,50 @@ func TestFlow_Consume(t *testing.T) { require.Nil(t, err) // 12 header, 8 timestamp, 4 MAC - assert.Equal(t, len(testContent)+12+8+4, n) + assert.Equal(t, len(testContent)+12+4, n) + }) + + t.Run("MultipleGeneratorsLength", func(t *testing.T) { + testMac2 := mocks.AlmostUselessMac("efgh") + testConn := mocks.NewMockPerfectBiPacketConn(10) + + flowA := newFlow(congestion.NewNone(), []proxy.MacVerifier{testMac, testMac2}, []proxy.MacGenerator{testMac, testMac2}) + + flowA.writer = testConn.SideB() + flowA.isAlive = true + + err := flowA.Consume(context.Background(), testPacket) + require.Nil(t, err) + + buf := make([]byte, 100) + n, _, err := testConn.SideA().ReadFromUDP(buf) + require.Nil(t, err) + + // 12 header, 8 timestamp, 4 MAC + assert.Equal(t, len(testContent)+12+4+4, n) + }) + + t.Run("MultipleGeneratorsOrder", func(t *testing.T) { + testMac2 := mocks.AlmostUselessMac("efgh") + testConn := mocks.NewMockPerfectBiPacketConn(10) + + flowA := newFlow(congestion.NewNone(), []proxy.MacVerifier{testMac, testMac2}, []proxy.MacGenerator{testMac, testMac2}) + + flowA.writer = testConn.SideB() + flowA.isAlive = true + + err := flowA.Consume(context.Background(), testPacket) + require.Nil(t, err) + + buf := make([]byte, 100) + n, _, err := testConn.SideA().ReadFromUDP(buf) + require.Nil(t, err) + + // 12 header, 8 timestamp, 4 MAC + require.Equal(t, len(testContent)+12+4+4, n) + + macs := string(buf[n-8:n]) + assert.Equal(t, "abcdefgh", macs) }) } @@ -45,7 +88,7 @@ func TestFlow_Produce(t *testing.T) { seq: 128, data: proxy.SimplePacket(testContent), } - testMac := mocks.AlmostUselessMac{} + testMac := mocks.AlmostUselessMac("abcd") testMarshalled := proxy.AppendMac(testPacket.Marshal(), testMac) @@ -58,7 +101,7 @@ func TestFlow_Produce(t *testing.T) { _, err := testConn.SideA().Write(testMarshalled) require.Nil(t, err) - flowA := newFlow(congestion.NewNone(), testMac) + flowA := newFlow(congestion.NewNone(), []proxy.MacVerifier{testMac}, []proxy.MacGenerator{testMac}) flowA.writer = testConn.SideB() flowA.isAlive = true @@ -67,7 +110,43 @@ func TestFlow_Produce(t *testing.T) { err := flowA.readQueuePacket(context.Background(), testConn.SideB()) assert.Nil(t, err) }() - p, err := flowA.Produce(context.Background(), testMac) + p, err := flowA.Produce(context.Background()) + + require.Nil(t, err) + assert.Len(t, p.Contents(), len(testContent)) + + done <- struct{}{} + }() + + timer := time.NewTimer(500 * time.Millisecond) + select { + case <-done: + case <-timer.C: + fmt.Println("timed out") + t.FailNow() + } + }) + + t.Run("MultipleVerifiersStrip", func(t *testing.T) { + done := make(chan struct{}) + + go func() { + testMac2 := mocks.AlmostUselessMac("efgh") + testConn := mocks.NewMockPerfectBiPacketConn(10) + + _, err := testConn.SideA().Write(proxy.AppendMac(testMarshalled, testMac2)) + require.Nil(t, err) + + flowA := newFlow(congestion.NewNone(), []proxy.MacVerifier{testMac, testMac2}, []proxy.MacGenerator{testMac, testMac2}) + + flowA.writer = testConn.SideB() + flowA.isAlive = true + + go func() { + err := flowA.readQueuePacket(context.Background(), testConn.SideB()) + assert.Nil(t, err) + }() + p, err := flowA.Produce(context.Background()) require.Nil(t, err) assert.Len(t, p.Contents(), len(testContent))