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,7 +269,7 @@ func TestServer_PauseTx(t *testing.T) {
 	}
 }
 
-func TestServer_BlackholeTx(t *testing.T) {
+func TestServer_ModifyTx_corrupt(t *testing.T) {
 	scheme := "unix"
 	srcAddr, dstAddr := newUnixAddr(), newUnixAddr()
 	defer func() {
@@ -287,39 +287,24 @@ func TestServer_BlackholeTx(t *testing.T) {
 	<-p.Ready()
 	defer p.Close()
 
-	p.BlackholeTx()
-
+	p.ModifyTx(func(d []byte) []byte {
+		d[len(d)/2]++
+		return d
+	})
 	data := []byte("Hello World!")
 	send(t, data, scheme, srcAddr, transport.TLSInfo{})
-
-	recvc := make(chan []byte)
-	go func() {
-		recvc <- receive(t, ln)
-	}()
-
-	select {
-	case d := <-recvc:
-		t.Fatalf("unexpected data receive %q during blackhole", string(d))
-	case <-time.After(200 * time.Millisecond):
+	if d := receive(t, ln); bytes.Equal(d, data) {
+		t.Fatalf("expected corrupted data, got %q", string(d))
 	}
 
-	p.UnblackholeTx()
-
-	// expect different data, old data dropped
-	data[0]++
+	p.UnmodifyTx()
 	send(t, data, scheme, srcAddr, transport.TLSInfo{})
-
-	select {
-	case d := <-recvc:
-		if !bytes.Equal(data, d) {
-			t.Fatalf("expected %q, got %q", string(data), string(d))
-		}
-	case <-time.After(2 * time.Second):
-		t.Fatal("took too long to receive after unblackhole")
+	if d := receive(t, ln); !bytes.Equal(d, data) {
+		t.Fatalf("expected uncorrupted data, got %q", string(d))
 	}
 }
 
-func TestServer_CorruptTx(t *testing.T) {
+func TestServer_ModifyTx_packet_loss(t *testing.T) {
 	scheme := "unix"
 	srcAddr, dstAddr := newUnixAddr(), newUnixAddr()
 	defer func() {
@@ -337,9 +322,10 @@ func TestServer_CorruptTx(t *testing.T) {
 	<-p.Ready()
 	defer p.Close()
 
-	p.CorruptTx(func(d []byte) []byte {
-		d[len(d)/2]++
-		return d
+	// 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{})
@@ -347,13 +333,63 @@ func TestServer_CorruptTx(t *testing.T) {
 		t.Fatalf("expected corrupted data, got %q", string(d))
 	}
 
-	p.UncorruptTx()
+	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()
+	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.BlackholeTx()
+
+	data := []byte("Hello World!")
+	send(t, data, scheme, srcAddr, transport.TLSInfo{})
+
+	recvc := make(chan []byte)
+	go func() {
+		recvc <- receive(t, ln)
+	}()
+
+	select {
+	case d := <-recvc:
+		t.Fatalf("unexpected data receive %q during blackhole", string(d))
+	case <-time.After(200 * time.Millisecond):
+	}
+
+	p.UnblackholeTx()
+
+	// expect different data, old data dropped
+	data[0]++
+	send(t, data, scheme, srcAddr, transport.TLSInfo{})
+
+	select {
+	case d := <-recvc:
+		if !bytes.Equal(data, d) {
+			t.Fatalf("expected %q, got %q", string(data), string(d))
+		}
+	case <-time.After(2 * time.Second):
+		t.Fatal("took too long to receive after unblackhole")
+	}
+}
+
 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)
-	px.listener.Close()
+	s, _ := p.(*server)
+	s.listener.Close()
 	time.Sleep(200 * time.Millisecond)
 
 	data := []byte("Hello World!")