-
Notifications
You must be signed in to change notification settings - Fork 3.8k
/
forwarder.go
158 lines (131 loc) · 4.84 KB
/
forwarder.go
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
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
// Copyright 2022 The Cockroach Authors.
//
// Licensed as a CockroachDB Enterprise file under the Cockroach Community
// License (the "License"); you may not use this file except in compliance with
// the License. You may obtain a copy of the License at
//
// https://github.com/cockroachdb/cockroach/blob/master/licenses/CCL.txt
package sqlproxyccl
import (
"context"
"net"
"github.com/cockroachdb/cockroach/pkg/util/syncutil"
"github.com/cockroachdb/errors"
)
// ErrForwarderClosed indicates that the forwarder has been closed.
var ErrForwarderClosed = errors.New("forwarder has been closed")
// ErrForwarderStarted indicates that the forwarder has already started.
var ErrForwarderStarted = errors.New("forwarder has already started")
// forwarder is used to forward pgwire messages from the client to the server,
// and vice-versa. At the moment, this does a direct proxying, and there is
// no intercepting. Once https://github.com/cockroachdb/cockroach/issues/76000
// has been addressed, we will start intercepting pgwire messages at their
// boundaries here.
type forwarder struct {
// ctx is a single context used to control all goroutines spawned by the
// forwarder.
ctx context.Context
ctxCancel context.CancelFunc
// crdbConn is only set after the authentication phase for the initial
// connection. In the context of a connection migration, crdbConn is only
// replaced once the session has successfully been deserialized, and the
// old connection will be closed.
conn net.Conn // client <-> proxy
crdbConn net.Conn // proxy <-> client
// errChan is a buffered channel that contains the first forwarder error.
// This channel may receive nil errors.
errChan chan error
mu struct {
syncutil.Mutex
// closed is set to true whenever the forwarder is closed explicitly
// through Close, or when any of its main goroutines has terminated,
// whichever that happens first.
//
// A new forwarder instance will have to be recreated if one wants to
// reuse the same pair of connections.
closed bool
// started is set to true once Run has been invoked on the forwarder.
// This is a safeguard to prevent callers from starting the forwarding
// process twice. This will never be set back to false.
started bool
}
}
// newForwarder returns a new instance of forwarder. When this is called, it
// is expected that the caller passes ownerships of conn and crdbConn to the
// forwarder, which implies that these connections may be closed by the
// forwarder if necessary, but this is not guaranteed. The caller should still
// attempt to cleanup the connections.
func newForwarder(ctx context.Context, conn net.Conn, crdbConn net.Conn) *forwarder {
ctx, cancelFn := context.WithCancel(ctx)
return &forwarder{
ctx: ctx,
ctxCancel: cancelFn,
conn: conn,
crdbConn: crdbConn,
errChan: make(chan error, 1),
}
}
// Run starts the forwarding process for the associated connections, and can
// only be called once throughout the lifetime of the forwarder instance.
//
// All goroutines spun up must check on f.ctx to prevent leaks, if possible. If
// there was an error within the goroutines, the forwarder will be closed, and
// the first error can be found in f.errChan.
//
// This returns ErrForwarderClosed if the forwarder has been closed, and
// ErrForwarderStarted if the forwarder has already been started.
func (f *forwarder) Run() error {
f.mu.Lock()
defer f.mu.Unlock()
if f.mu.closed {
return ErrForwarderClosed
}
if f.mu.started {
return ErrForwarderStarted
}
go func() {
defer f.Close()
// Block until context is done.
<-f.ctx.Done()
// Closing the connection will terminate the goroutine below.
// This will be here temporarily because the only way to unblock
// io.Copy is to close one of the ends. Once we replace io.Copy
// with the interceptors, we could use ctx directly, and no longer
// need this goroutine.
f.crdbConn.Close()
}()
// Copy all pgwire messages from frontend to backend connection until we
// encounter an error or shutdown signal.
go func() {
defer f.Close()
err := ConnectionCopy(f.crdbConn, f.conn)
select {
case f.errChan <- err: /* error reported */
default: /* the channel already contains an error */
}
}()
f.mu.started = true
return nil
}
// Close closes the forwarder, and stops the forwarding process. This is
// idempotent.
func (f *forwarder) Close() {
f.mu.Lock()
defer f.mu.Unlock()
f.mu.closed = true
f.ctxCancel()
}
// IsClosed returns a boolean indicating whether the forwarder is closed.
func (f *forwarder) IsClosed() bool {
f.mu.Lock()
defer f.mu.Unlock()
return f.mu.closed
}
// IsStarted returns a boolean indicating whether the forwarder has been
// started during its lifetime. This does not indicate that the forwarder has
// been closed.
func (f *forwarder) IsStarted() bool {
f.mu.Lock()
defer f.mu.Unlock()
return f.mu.started
}