pkg/proxy: make/simplify interface more extensible

Extend proxy for more advanced corrupt and packet drop testing.

Signed-off-by: Gyuho Lee <gyuhox@gmail.com>

pkg/proxy/server_test.go

@@ -269,6 +269,77 @@ func TestServer_PauseTx(t *testing.T) {
 	}
 }
 
+func TestServer_ModifyTx_corrupt(t *testing.T) {
+	scheme := "unix"
+	srcAddr, dstAddr := newUnixAddr(), newUnixAddr()
+	defer func() {
+		os.RemoveAll(srcAddr)
+		os.RemoveAll(dstAddr)
+	}()
+	ln := listen(t, scheme, dstAddr, transport.TLSInfo{})
+	defer ln.Close()
+
+	p := NewServer(ServerConfig{
+		Logger: testLogger,
+		From:   url.URL{Scheme: scheme, Host: srcAddr},
+		To:     url.URL{Scheme: scheme, Host: dstAddr},
+	})
+	<-p.Ready()
+	defer p.Close()
+
+	p.ModifyTx(func(d []byte) []byte {
+		d[len(d)/2]++
+		return d
+	})
+	data := []byte("Hello World!")
+	send(t, data, scheme, srcAddr, transport.TLSInfo{})
+	if d := receive(t, ln); bytes.Equal(d, data) {
+		t.Fatalf("expected corrupted data, got %q", string(d))
+	}
+
+	p.UnmodifyTx()
+	send(t, data, scheme, srcAddr, transport.TLSInfo{})
+	if d := receive(t, ln); !bytes.Equal(d, data) {
+		t.Fatalf("expected uncorrupted data, got %q", string(d))
+	}
+}
+
+func TestServer_ModifyTx_packet_loss(t *testing.T) {
+	scheme := "unix"
+	srcAddr, dstAddr := newUnixAddr(), newUnixAddr()
+	defer func() {
+		os.RemoveAll(srcAddr)
+		os.RemoveAll(dstAddr)
+	}()
+	ln := listen(t, scheme, dstAddr, transport.TLSInfo{})
+	defer ln.Close()
+
+	p := NewServer(ServerConfig{
+		Logger: testLogger,
+		From:   url.URL{Scheme: scheme, Host: srcAddr},
+		To:     url.URL{Scheme: scheme, Host: dstAddr},
+	})
+	<-p.Ready()
+	defer p.Close()
+
+	// 50% packet loss
+	p.ModifyTx(func(d []byte) []byte {
+		half := len(d) / 2
+		return d[:half:half]
+	})
+	data := []byte("Hello World!")
+	send(t, data, scheme, srcAddr, transport.TLSInfo{})
+	if d := receive(t, ln); bytes.Equal(d, data) {
+		t.Fatalf("expected corrupted data, got %q", string(d))
+	}
+
+	p.UnmodifyTx()
+	send(t, data, scheme, srcAddr, transport.TLSInfo{})
+	if d := receive(t, ln); !bytes.Equal(d, data) {
+		t.Fatalf("expected uncorrupted data, got %q", string(d))
+	}
+}
+
 func TestServer_BlackholeTx(t *testing.T) {
 	scheme := "unix"
 	srcAddr, dstAddr := newUnixAddr(), newUnixAddr()
@@ -319,41 +390,6 @@ func TestServer_BlackholeTx(t *testing.T) {
 	}
 }
 
-func TestServer_CorruptTx(t *testing.T) {
-	scheme := "unix"
-	srcAddr, dstAddr := newUnixAddr(), newUnixAddr()
-	defer func() {
-		os.RemoveAll(srcAddr)
-		os.RemoveAll(dstAddr)
-	}()
-	ln := listen(t, scheme, dstAddr, transport.TLSInfo{})
-	defer ln.Close()
-
-	p := NewServer(ServerConfig{
-		Logger: testLogger,
-		From:   url.URL{Scheme: scheme, Host: srcAddr},
-		To:     url.URL{Scheme: scheme, Host: dstAddr},
-	})
-	<-p.Ready()
-	defer p.Close()
-
-	p.CorruptTx(func(d []byte) []byte {
-		d[len(d)/2]++
-		return d
-	})
-	data := []byte("Hello World!")
-	send(t, data, scheme, srcAddr, transport.TLSInfo{})
-	if d := receive(t, ln); bytes.Equal(d, data) {
-		t.Fatalf("expected corrupted data, got %q", string(d))
-	}
-
-	p.UncorruptTx()
-	send(t, data, scheme, srcAddr, transport.TLSInfo{})
-	if d := receive(t, ln); !bytes.Equal(d, data) {
-		t.Fatalf("expected uncorrupted data, got %q", string(d))
-	}
-}
-
 func TestServer_Shutdown(t *testing.T) {
 	scheme := "unix"
 	srcAddr, dstAddr := newUnixAddr(), newUnixAddr()
@@ -372,8 +408,8 @@ func TestServer_Shutdown(t *testing.T) {
 	<-p.Ready()
 	defer p.Close()
 
-	px, _ := p.(*proxyServer)
+	s, _ := p.(*server)
-	px.listener.Close()
+	s.listener.Close()
 	time.Sleep(200 * time.Millisecond)
 
 	data := []byte("Hello World!")