Skip to content

Commit 72cc8cf

Browse files
jkh52jkh52
committed
proxy-server: Wrap Backend more completely.
The goal is to fix nil deref at #513, which was caused by ProxyServer.addBackend inconsistent return value.
1 parent 9c1359e commit 72cc8cf

File tree

6 files changed

+124
-126
lines changed

6 files changed

+124
-126
lines changed

‎pkg/server/backend_manager.go

Lines changed: 45 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ import (
2525
"sync"
2626
"time"
2727

28+
"google.golang.org/grpc/metadata"
2829
"k8s.io/klog/v2"
2930

3031
commonmetrics "sigs.k8s.io/apiserver-network-proxy/konnectivity-client/pkg/common/metrics"
@@ -70,18 +71,21 @@ func GenProxyStrategiesFromStr(proxyStrategies string) ([]ProxyStrategy, error)
7071
return ps, nil
7172
}
7273

74+
// Backend abstracts a connected Konnectivity agent.
75+
//
76+
// In the only currently supported case (gRPC), it wraps an
77+
// agent.AgentService_ConnectServer, provides synchronization and
78+
// emits common stream metrics.
7379
type Backend interface {
7480
Send(p *client.Packet) error
7581
Recv() (*client.Packet, error)
7682
Context() context.Context
83+
GetAgentIdentifiers() (header.Identifiers, error)
7784
}
7885

7986
var _ Backend = &backend{}
80-
var _ Backend = agent.AgentService_ConnectServer(nil)
8187

8288
type backend struct {
83-
// TODO: this is a multi-writer single-reader pattern, it's tricky to
84-
// write it using channel. Let's worry about performance later.
8589
sendLock sync.Mutex
8690
recvLock sync.Mutex
8791
conn agent.AgentService_ConnectServer
@@ -121,17 +125,34 @@ func (b *backend) Context() context.Context {
121125
return b.conn.Context()
122126
}
123127

124-
func newBackend(conn agent.AgentService_ConnectServer) *backend {
128+
func (b *backend) GetAgentIdentifiers() (header.Identifiers, error) {
129+
var agentIdentifiers header.Identifiers
130+
md, ok := metadata.FromIncomingContext(b.Context())
131+
if !ok {
132+
return agentIdentifiers, fmt.Errorf("failed to get metadata from context")
133+
}
134+
agentIDs := md.Get(header.AgentIdentifiers)
135+
if len(agentIDs) > 1 {
136+
return agentIdentifiers, fmt.Errorf("expected at most one set of agent IDs in the context, got %v", agentIDs)
137+
}
138+
if len(agentIDs) == 0 {
139+
return agentIdentifiers, nil
140+
}
141+
142+
return header.GenAgentIdentifiers(agentIDs[0])
143+
}
144+
145+
func NewBackend(conn agent.AgentService_ConnectServer) Backend {
125146
return &backend{conn: conn}
126147
}
127148

128149
// BackendStorage is an interface to manage the storage of the backend
129150
// connections, i.e., get, add and remove
130151
type BackendStorage interface {
131152
// AddBackend adds a backend.
132-
AddBackend(identifier string, idType header.IdentifierType, conn agent.AgentService_ConnectServer) Backend
153+
AddBackend(identifier string, idType header.IdentifierType, backend Backend)
133154
// RemoveBackend removes a backend.
134-
RemoveBackend(identifier string, idType header.IdentifierType, conn agent.AgentService_ConnectServer)
155+
RemoveBackend(identifier string, idType header.IdentifierType, backend Backend)
135156
// NumBackends returns the number of backends.
136157
NumBackends() int
137158
}
@@ -168,7 +189,7 @@ type DefaultBackendStorage struct {
168189
// For a given agent, ProxyServer prefers backends[agentID][0] to send
169190
// traffic, because backends[agentID][1:] are more likely to be closed
170191
// by the agent to deduplicate connections to the same server.
171-
backends map[string][]*backend
192+
backends map[string][]Backend
172193
// agentID is tracked in this slice to enable randomly picking an
173194
// agentID in the Backend() method. There is no reliable way to
174195
// randomly pick a key from a map (in this case, the backends) in
@@ -198,7 +219,7 @@ func NewDefaultBackendStorage(idTypes []header.IdentifierType) *DefaultBackendSt
198219
// no agent ever successfully connects.
199220
metrics.Metrics.SetBackendCount(0)
200221
return &DefaultBackendStorage{
201-
backends: make(map[string][]*backend),
222+
backends: make(map[string][]Backend),
202223
random: rand.New(rand.NewSource(time.Now().UnixNano())),
203224
idTypes: idTypes,
204225
} /* #nosec G404 */
@@ -214,42 +235,40 @@ func containIDType(idTypes []header.IdentifierType, idType header.IdentifierType
214235
}
215236

216237
// AddBackend adds a backend.
217-
func (s *DefaultBackendStorage) AddBackend(identifier string, idType header.IdentifierType, conn agent.AgentService_ConnectServer) Backend {
238+
func (s *DefaultBackendStorage) AddBackend(identifier string, idType header.IdentifierType, backend Backend) {
218239
if !containIDType(s.idTypes, idType) {
219240
klog.V(4).InfoS("fail to add backend", "backend", identifier, "error", &ErrWrongIDType{idType, s.idTypes})
220-
return nil
241+
return
221242
}
222-
klog.V(5).InfoS("Register backend for agent", "connection", conn, "agentID", identifier)
243+
klog.V(5).InfoS("Register backend for agent", "agentID", identifier)
223244
s.mu.Lock()
224245
defer s.mu.Unlock()
225246
_, ok := s.backends[identifier]
226-
addedBackend := newBackend(conn)
227247
if ok {
228-
for _, v := range s.backends[identifier] {
229-
if v.conn == conn {
230-
klog.V(1).InfoS("This should not happen. Adding existing backend for agent", "connection", conn, "agentID", identifier)
231-
return v
248+
for _, b := range s.backends[identifier] {
249+
if b == backend {
250+
klog.V(1).InfoS("This should not happen. Adding existing backend for agent", "agentID", identifier)
251+
return
232252
}
233253
}
234-
s.backends[identifier] = append(s.backends[identifier], addedBackend)
235-
return addedBackend
254+
s.backends[identifier] = append(s.backends[identifier], backend)
255+
return
236256
}
237-
s.backends[identifier] = []*backend{addedBackend}
257+
s.backends[identifier] = []Backend{backend}
238258
metrics.Metrics.SetBackendCount(len(s.backends))
239259
s.agentIDs = append(s.agentIDs, identifier)
240260
if idType == header.DefaultRoute {
241261
s.defaultRouteAgentIDs = append(s.defaultRouteAgentIDs, identifier)
242262
}
243-
return addedBackend
244263
}
245264

246265
// RemoveBackend removes a backend.
247-
func (s *DefaultBackendStorage) RemoveBackend(identifier string, idType header.IdentifierType, conn agent.AgentService_ConnectServer) {
266+
func (s *DefaultBackendStorage) RemoveBackend(identifier string, idType header.IdentifierType, backend Backend) {
248267
if !containIDType(s.idTypes, idType) {
249268
klog.ErrorS(&ErrWrongIDType{idType, s.idTypes}, "fail to remove backend")
250269
return
251270
}
252-
klog.V(5).InfoS("Remove connection for agent", "connection", conn, "identifier", identifier)
271+
klog.V(5).InfoS("Remove connection for agent", "agentID", identifier)
253272
s.mu.Lock()
254273
defer s.mu.Unlock()
255274
backends, ok := s.backends[identifier]
@@ -258,11 +277,11 @@ func (s *DefaultBackendStorage) RemoveBackend(identifier string, idType header.I
258277
return
259278
}
260279
var found bool
261-
for i, c := range backends {
262-
if c.conn == conn {
280+
for i, b := range backends {
281+
if b == backend {
263282
s.backends[identifier] = append(s.backends[identifier][:i], s.backends[identifier][i+1:]...)
264283
if i == 0 && len(s.backends[identifier]) != 0 {
265-
klog.V(1).InfoS("This should not happen. Removed connection that is not the first connection", "connection", conn, "remainingConnections", s.backends[identifier])
284+
klog.V(1).InfoS("This should not happen. Removed connection that is not the first connection", "agentID", identifier)
266285
}
267286
found = true
268287
}
@@ -286,7 +305,7 @@ func (s *DefaultBackendStorage) RemoveBackend(identifier string, idType header.I
286305
}
287306
}
288307
if !found {
289-
klog.V(1).InfoS("Could not find connection matching identifier to remove", "connection", conn, "identifier", identifier)
308+
klog.V(1).InfoS("Could not find connection matching identifier to remove", "agentID", identifier, "idType", idType)
290309
}
291310
metrics.Metrics.SetBackendCount(len(s.backends))
292311
}

‎pkg/server/backend_manager_test.go

Lines changed: 43 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -29,17 +29,17 @@ type fakeAgentServiceConnectServer struct {
2929
}
3030

3131
func TestAddRemoveBackends(t *testing.T) {
32-
conn1 := new(fakeAgentServiceConnectServer)
33-
conn12 := new(fakeAgentServiceConnectServer)
34-
conn2 := new(fakeAgentServiceConnectServer)
35-
conn22 := new(fakeAgentServiceConnectServer)
36-
conn3 := new(fakeAgentServiceConnectServer)
32+
backend1 := NewBackend(new(fakeAgentServiceConnectServer))
33+
backend12 := NewBackend(new(fakeAgentServiceConnectServer))
34+
backend2 := NewBackend(new(fakeAgentServiceConnectServer))
35+
backend22 := NewBackend(new(fakeAgentServiceConnectServer))
36+
backend3 := NewBackend(new(fakeAgentServiceConnectServer))
3737

3838
p := NewDefaultBackendManager()
3939

40-
p.AddBackend("agent1", header.UID, conn1)
41-
p.RemoveBackend("agent1", header.UID, conn1)
42-
expectedBackends := make(map[string][]*backend)
40+
p.AddBackend("agent1", header.UID, backend1)
41+
p.RemoveBackend("agent1", header.UID, backend1)
42+
expectedBackends := make(map[string][]Backend)
4343
expectedAgentIDs := []string{}
4444
if e, a := expectedBackends, p.backends; !reflect.DeepEqual(e, a) {
4545
t.Errorf("expected %v, got %v", e, a)
@@ -49,21 +49,21 @@ func TestAddRemoveBackends(t *testing.T) {
4949
}
5050

5151
p = NewDefaultBackendManager()
52-
p.AddBackend("agent1", header.UID, conn1)
53-
p.AddBackend("agent1", header.UID, conn12)
52+
p.AddBackend("agent1", header.UID, backend1)
53+
p.AddBackend("agent1", header.UID, backend12)
5454
// Adding the same connection again should be a no-op.
55-
p.AddBackend("agent1", header.UID, conn12)
56-
p.AddBackend("agent2", header.UID, conn2)
57-
p.AddBackend("agent2", header.UID, conn22)
58-
p.AddBackend("agent3", header.UID, conn3)
59-
p.RemoveBackend("agent2", header.UID, conn22)
60-
p.RemoveBackend("agent2", header.UID, conn2)
61-
p.RemoveBackend("agent1", header.UID, conn1)
62-
// This is invalid. agent1 doesn't have conn3. This should be a no-op.
63-
p.RemoveBackend("agent1", header.UID, conn3)
64-
expectedBackends = map[string][]*backend{
65-
"agent1": {newBackend(conn12)},
66-
"agent3": {newBackend(conn3)},
55+
p.AddBackend("agent1", header.UID, backend12)
56+
p.AddBackend("agent2", header.UID, backend2)
57+
p.AddBackend("agent2", header.UID, backend22)
58+
p.AddBackend("agent3", header.UID, backend3)
59+
p.RemoveBackend("agent2", header.UID, backend22)
60+
p.RemoveBackend("agent2", header.UID, backend2)
61+
p.RemoveBackend("agent1", header.UID, backend1)
62+
// This is invalid. agent1 doesn't have backend3. This should be a no-op.
63+
p.RemoveBackend("agent1", header.UID, backend3)
64+
expectedBackends = map[string][]Backend{
65+
"agent1": {backend12},
66+
"agent3": {backend3},
6767
}
6868
expectedAgentIDs = []string{"agent1", "agent3"}
6969
if e, a := expectedBackends, p.backends; !reflect.DeepEqual(e, a) {
@@ -75,17 +75,17 @@ func TestAddRemoveBackends(t *testing.T) {
7575
}
7676

7777
func TestAddRemoveBackendsWithDefaultRoute(t *testing.T) {
78-
conn1 := new(fakeAgentServiceConnectServer)
79-
conn12 := new(fakeAgentServiceConnectServer)
80-
conn2 := new(fakeAgentServiceConnectServer)
81-
conn22 := new(fakeAgentServiceConnectServer)
82-
conn3 := new(fakeAgentServiceConnectServer)
78+
backend1 := NewBackend(new(fakeAgentServiceConnectServer))
79+
backend12 := NewBackend(new(fakeAgentServiceConnectServer))
80+
backend2 := NewBackend(new(fakeAgentServiceConnectServer))
81+
backend22 := NewBackend(new(fakeAgentServiceConnectServer))
82+
backend3 := NewBackend(new(fakeAgentServiceConnectServer))
8383

8484
p := NewDefaultRouteBackendManager()
8585

86-
p.AddBackend("agent1", header.DefaultRoute, conn1)
87-
p.RemoveBackend("agent1", header.DefaultRoute, conn1)
88-
expectedBackends := make(map[string][]*backend)
86+
p.AddBackend("agent1", header.DefaultRoute, backend1)
87+
p.RemoveBackend("agent1", header.DefaultRoute, backend1)
88+
expectedBackends := make(map[string][]Backend)
8989
expectedAgentIDs := []string{}
9090
if e, a := expectedBackends, p.backends; !reflect.DeepEqual(e, a) {
9191
t.Errorf("expected %v, got %v", e, a)
@@ -98,22 +98,22 @@ func TestAddRemoveBackendsWithDefaultRoute(t *testing.T) {
9898
}
9999

100100
p = NewDefaultRouteBackendManager()
101-
p.AddBackend("agent1", header.DefaultRoute, conn1)
102-
p.AddBackend("agent1", header.DefaultRoute, conn12)
101+
p.AddBackend("agent1", header.DefaultRoute, backend1)
102+
p.AddBackend("agent1", header.DefaultRoute, backend12)
103103
// Adding the same connection again should be a no-op.
104-
p.AddBackend("agent1", header.DefaultRoute, conn12)
105-
p.AddBackend("agent2", header.DefaultRoute, conn2)
106-
p.AddBackend("agent2", header.DefaultRoute, conn22)
107-
p.AddBackend("agent3", header.DefaultRoute, conn3)
108-
p.RemoveBackend("agent2", header.DefaultRoute, conn22)
109-
p.RemoveBackend("agent2", header.DefaultRoute, conn2)
110-
p.RemoveBackend("agent1", header.DefaultRoute, conn1)
104+
p.AddBackend("agent1", header.DefaultRoute, backend12)
105+
p.AddBackend("agent2", header.DefaultRoute, backend2)
106+
p.AddBackend("agent2", header.DefaultRoute, backend22)
107+
p.AddBackend("agent3", header.DefaultRoute, backend3)
108+
p.RemoveBackend("agent2", header.DefaultRoute, backend22)
109+
p.RemoveBackend("agent2", header.DefaultRoute, backend2)
110+
p.RemoveBackend("agent1", header.DefaultRoute, backend1)
111111
// This is invalid. agent1 doesn't have conn3. This should be a no-op.
112-
p.RemoveBackend("agent1", header.DefaultRoute, conn3)
112+
p.RemoveBackend("agent1", header.DefaultRoute, backend3)
113113

114-
expectedBackends = map[string][]*backend{
115-
"agent1": {newBackend(conn12)},
116-
"agent3": {newBackend(conn3)},
114+
expectedBackends = map[string][]Backend{
115+
"agent1": {backend12},
116+
"agent3": {backend3},
117117
}
118118
expectedDefaultRouteAgentIDs := []string{"agent1", "agent3"}
119119

0 commit comments

Comments
 (0)