1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
|
package proxy_test
import (
"context"
"io"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gitlab.com/gitlab-org/gitaly/internal/metadata"
"gitlab.com/gitlab-org/gitaly/internal/praefect/grpc-proxy/proxy"
testservice "gitlab.com/gitlab-org/gitaly/internal/praefect/grpc-proxy/testdata"
"gitlab.com/gitlab-org/gitaly/internal/testhelper"
"google.golang.org/protobuf/proto"
)
// TestStreamPeeking demonstrates that a director function is able to peek
// into a stream. Further more, it demonstrates that peeking into a stream
// will not disturb the stream sent from the proxy client to the backend.
func TestStreamPeeking(t *testing.T) {
ctx := testhelper.Context(t)
backendCC, backendSrvr, cleanupPinger := newBackendPinger(t, ctx)
defer cleanupPinger()
pingReqSent := &testservice.PingRequest{Value: "hi"}
// director will peek into stream before routing traffic
director := func(ctx context.Context, fullMethodName string, peeker proxy.StreamPeeker) (*proxy.StreamParameters, error) {
peekedMsg, err := peeker.Peek()
require.NoError(t, err)
peekedRequest := &testservice.PingRequest{}
err = proto.Unmarshal(peekedMsg, peekedRequest)
require.NoError(t, err)
require.True(t, proto.Equal(pingReqSent, peekedRequest), "expected to be the same")
return proxy.NewStreamParameters(proxy.Destination{Ctx: metadata.IncomingToOutgoing(ctx), Conn: backendCC, Msg: peekedMsg}, nil, nil, nil), nil
}
pingResp := &testservice.PingResponse{
Counter: 1,
}
// we expect the backend server to receive the peeked message
backendSrvr.pingStream = func(stream testservice.TestService_PingStreamServer) error {
pingReqReceived, err := stream.Recv()
assert.NoError(t, err)
assert.True(t, proto.Equal(pingReqSent, pingReqReceived), "expected to be the same")
return stream.Send(pingResp)
}
proxyCC, cleanupProxy := newProxy(t, ctx, director, "mwitkow.testproto.TestService", "PingStream")
defer cleanupProxy()
proxyClient := testservice.NewTestServiceClient(proxyCC)
proxyClientPingStream, err := proxyClient.PingStream(ctx)
require.NoError(t, err)
defer func() {
require.NoError(t, proxyClientPingStream.CloseSend())
}()
require.NoError(t,
proxyClientPingStream.Send(pingReqSent),
)
resp, err := proxyClientPingStream.Recv()
require.NoError(t, err)
require.True(t, proto.Equal(resp, pingResp), "expected to be the same")
_, err = proxyClientPingStream.Recv()
require.Equal(t, io.EOF, err)
}
func TestStreamInjecting(t *testing.T) {
ctx := testhelper.Context(t)
backendCC, backendSrvr, cleanupPinger := newBackendPinger(t, ctx)
defer cleanupPinger()
pingReqSent := &testservice.PingRequest{Value: "hi"}
newValue := "bye"
// director will peek into stream and change some frames
director := func(ctx context.Context, fullMethodName string, peeker proxy.StreamPeeker) (*proxy.StreamParameters, error) {
peekedMsg, err := peeker.Peek()
require.NoError(t, err)
peekedRequest := &testservice.PingRequest{}
require.NoError(t, proto.Unmarshal(peekedMsg, peekedRequest))
require.Equal(t, "hi", peekedRequest.GetValue())
peekedRequest.Value = newValue
newPayload, err := proto.Marshal(peekedRequest)
require.NoError(t, err)
return proxy.NewStreamParameters(proxy.Destination{Ctx: metadata.IncomingToOutgoing(ctx), Conn: backendCC, Msg: newPayload}, nil, nil, nil), nil
}
pingResp := &testservice.PingResponse{
Counter: 1,
}
// we expect the backend server to receive the modified message
backendSrvr.pingStream = func(stream testservice.TestService_PingStreamServer) error {
pingReqReceived, err := stream.Recv()
assert.NoError(t, err)
assert.Equal(t, newValue, pingReqReceived.GetValue())
return stream.Send(pingResp)
}
proxyCC, cleanupProxy := newProxy(t, ctx, director, "mwitkow.testproto.TestService", "PingStream")
defer cleanupProxy()
proxyClient := testservice.NewTestServiceClient(proxyCC)
proxyClientPingStream, err := proxyClient.PingStream(ctx)
require.NoError(t, err)
defer func() {
require.NoError(t, proxyClientPingStream.CloseSend())
}()
require.NoError(t,
proxyClientPingStream.Send(pingReqSent),
)
resp, err := proxyClientPingStream.Recv()
require.NoError(t, err)
require.True(t, proto.Equal(resp, pingResp), "expected to be the same")
_, err = proxyClientPingStream.Recv()
require.Equal(t, io.EOF, err)
}
|
