TUN-8861: Add session limiter to UDP session manager
## Summary In order to make cloudflared behavior more predictable and prevent an exhaustion of resources, we have decided to add session limits that can be configured by the user. This first commit introduces the session limiter and adds it to the UDP handling path. For now the limiter is set to run only in unlimited mode.
This commit is contained in:
parent
8918b6729e
commit
bf4954e96a
4
Makefile
4
Makefile
|
@ -265,3 +265,7 @@ fmt-check:
|
|||
.PHONY: lint
|
||||
lint:
|
||||
@golangci-lint run
|
||||
|
||||
.PHONY: mocks
|
||||
mocks:
|
||||
go generate mocks/mockgen.go
|
||||
|
|
|
@ -67,6 +67,7 @@ For example, as of January 2023 Cloudflare will support cloudflared version 2023
|
|||
- [capnpc-go](https://pkg.go.dev/zombiezen.com/go/capnproto2/capnpc-go)
|
||||
- [goimports](https://pkg.go.dev/golang.org/x/tools/cmd/goimports)
|
||||
- [golangci-lint](https://github.com/golangci/golangci-lint)
|
||||
- [gomocks](https://pkg.go.dev/go.uber.org/mock)
|
||||
|
||||
### Build
|
||||
To build cloudflared locally run `make cloudflared`
|
||||
|
@ -76,3 +77,6 @@ To locally run the tests run `make test`
|
|||
|
||||
### Linting
|
||||
To format the code and keep a good code quality use `make fmt` and `make lint`
|
||||
|
||||
### Mocks
|
||||
After changes on interfaces you might need to regenerate the mocks, so run `make mock`
|
||||
|
|
|
@ -28,6 +28,8 @@ import (
|
|||
"github.com/stretchr/testify/require"
|
||||
"golang.org/x/net/nettest"
|
||||
|
||||
cfdsession "github.com/cloudflare/cloudflared/session"
|
||||
|
||||
"github.com/cloudflare/cloudflared/datagramsession"
|
||||
"github.com/cloudflare/cloudflared/ingress"
|
||||
"github.com/cloudflare/cloudflared/packet"
|
||||
|
@ -53,7 +55,8 @@ var _ ReadWriteAcker = (*streamReadWriteAcker)(nil)
|
|||
func TestQUICServer(t *testing.T) {
|
||||
// This is simply a sample websocket frame message.
|
||||
wsBuf := &bytes.Buffer{}
|
||||
wsutil.WriteClientBinary(wsBuf, []byte("Hello"))
|
||||
err := wsutil.WriteClientBinary(wsBuf, []byte("Hello"))
|
||||
require.NoError(t, err)
|
||||
|
||||
var tests = []struct {
|
||||
desc string
|
||||
|
@ -158,17 +161,19 @@ func TestQUICServer(t *testing.T) {
|
|||
|
||||
serverDone := make(chan struct{})
|
||||
go func() {
|
||||
// nolint: testifylint
|
||||
quicServer(
|
||||
ctx, t, quicListener, test.dest, test.connectionType, test.metadata, test.message, test.expectedResponse,
|
||||
)
|
||||
close(serverDone)
|
||||
}()
|
||||
|
||||
// nolint: gosec
|
||||
tunnelConn, _ := testTunnelConnection(t, netip.MustParseAddrPort(udpListener.LocalAddr().String()), uint8(i))
|
||||
|
||||
connDone := make(chan struct{})
|
||||
go func() {
|
||||
tunnelConn.Serve(ctx)
|
||||
_ = tunnelConn.Serve(ctx)
|
||||
close(connDone)
|
||||
}()
|
||||
|
||||
|
@ -254,14 +259,14 @@ func (moc *mockOriginProxyWithRequest) ProxyHTTP(w ResponseWriter, tr *tracing.T
|
|||
case "/ok":
|
||||
originRespEndpoint(w, http.StatusOK, []byte(http.StatusText(http.StatusOK)))
|
||||
case "/slow_echo_body":
|
||||
time.Sleep(5)
|
||||
time.Sleep(5 * time.Nanosecond)
|
||||
fallthrough
|
||||
case "/echo_body":
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
}
|
||||
_ = w.WriteRespHeaders(resp.StatusCode, resp.Header)
|
||||
io.Copy(w, r.Body)
|
||||
_, _ = io.Copy(w, r.Body)
|
||||
case "/error":
|
||||
return fmt.Errorf("Failed to proxy to origin")
|
||||
default:
|
||||
|
@ -493,16 +498,16 @@ func TestBuildHTTPRequest(t *testing.T) {
|
|||
test := test // capture range variable
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
req, err := buildHTTPRequest(context.Background(), test.connectRequest, test.body, 0, &log)
|
||||
assert.NoError(t, err)
|
||||
require.NoError(t, err)
|
||||
test.req = test.req.WithContext(req.Context())
|
||||
assert.Equal(t, test.req, req.Request)
|
||||
require.Equal(t, test.req, req.Request)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func (moc *mockOriginProxyWithRequest) ProxyTCP(ctx context.Context, rwa ReadWriteAcker, tcpRequest *TCPRequest) error {
|
||||
rwa.AckConnection("")
|
||||
io.Copy(rwa, rwa)
|
||||
_ = rwa.AckConnection("")
|
||||
_, _ = io.Copy(rwa, rwa)
|
||||
return nil
|
||||
}
|
||||
|
||||
|
@ -520,16 +525,19 @@ func TestServeUDPSession(t *testing.T) {
|
|||
edgeQUICSessionChan := make(chan quic.Connection)
|
||||
go func() {
|
||||
earlyListener, err := quic.Listen(udpListener, testTLSServerConfig, testQUICConfig)
|
||||
require.NoError(t, err)
|
||||
assert.NoError(t, err)
|
||||
|
||||
edgeQUICSession, err := earlyListener.Accept(ctx)
|
||||
require.NoError(t, err)
|
||||
assert.NoError(t, err)
|
||||
|
||||
edgeQUICSessionChan <- edgeQUICSession
|
||||
}()
|
||||
|
||||
// Random index to avoid reusing port
|
||||
tunnelConn, datagramConn := testTunnelConnection(t, netip.MustParseAddrPort(udpListener.LocalAddr().String()), 28)
|
||||
go tunnelConn.Serve(ctx)
|
||||
go func() {
|
||||
_ = tunnelConn.Serve(ctx)
|
||||
}()
|
||||
|
||||
edgeQUICSession := <-edgeQUICSessionChan
|
||||
|
||||
|
@ -545,14 +553,14 @@ func TestNopCloserReadWriterCloseBeforeEOF(t *testing.T) {
|
|||
|
||||
n, err := readerWriter.Read(buffer)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, n, 5)
|
||||
require.Equal(t, 5, n)
|
||||
|
||||
// close
|
||||
require.NoError(t, readerWriter.Close())
|
||||
|
||||
// read should get error
|
||||
n, err = readerWriter.Read(buffer)
|
||||
require.Equal(t, n, 0)
|
||||
require.Equal(t, 0, n)
|
||||
require.Equal(t, err, fmt.Errorf("closed by handler"))
|
||||
}
|
||||
|
||||
|
@ -562,7 +570,7 @@ func TestNopCloserReadWriterCloseAfterEOF(t *testing.T) {
|
|||
|
||||
n, err := readerWriter.Read(buffer)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, n, 9)
|
||||
require.Equal(t, 9, n)
|
||||
|
||||
// force another read to read eof
|
||||
_, err = readerWriter.Read(buffer)
|
||||
|
@ -573,7 +581,7 @@ func TestNopCloserReadWriterCloseAfterEOF(t *testing.T) {
|
|||
|
||||
// read should get EOF still
|
||||
n, err = readerWriter.Read(buffer)
|
||||
require.Equal(t, n, 0)
|
||||
require.Equal(t, 0, n)
|
||||
require.Equal(t, err, io.EOF)
|
||||
}
|
||||
|
||||
|
@ -669,6 +677,7 @@ func serveSession(ctx context.Context, datagramConn *datagramV2Connection, edgeQ
|
|||
unregisterReason: expectedReason,
|
||||
calledUnregisterChan: unregisterFromEdgeChan,
|
||||
}
|
||||
// nolint: testifylint
|
||||
go runRPCServer(ctx, edgeQUICSession, sessionRPCServer, nil, t)
|
||||
|
||||
<-unregisterFromEdgeChan
|
||||
|
@ -729,6 +738,7 @@ func (s mockSessionRPCServer) UnregisterUdpSession(ctx context.Context, sessionI
|
|||
|
||||
func testTunnelConnection(t *testing.T, serverAddr netip.AddrPort, index uint8) (TunnelConnection, *datagramV2Connection) {
|
||||
tlsClientConfig := &tls.Config{
|
||||
// nolint: gosec
|
||||
InsecureSkipVerify: true,
|
||||
NextProtos: []string{"argotunnel"},
|
||||
}
|
||||
|
@ -747,6 +757,7 @@ func testTunnelConnection(t *testing.T, serverAddr netip.AddrPort, index uint8)
|
|||
index,
|
||||
&log,
|
||||
)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Start a session manager for the connection
|
||||
sessionDemuxChan := make(chan *packet.Session, 4)
|
||||
|
@ -757,7 +768,9 @@ func testTunnelConnection(t *testing.T, serverAddr netip.AddrPort, index uint8)
|
|||
|
||||
datagramConn := &datagramV2Connection{
|
||||
conn,
|
||||
index,
|
||||
sessionManager,
|
||||
cfdsession.NewLimiter(0),
|
||||
datagramMuxer,
|
||||
packetRouter,
|
||||
15 * time.Second,
|
||||
|
@ -796,6 +809,7 @@ func (m *mockReaderNoopWriter) Close() error {
|
|||
|
||||
// GenerateTLSConfig sets up a bare-bones TLS config for a QUIC server
|
||||
func GenerateTLSConfig() *tls.Config {
|
||||
// nolint: gosec
|
||||
key, err := rsa.GenerateKey(rand.Reader, 1024)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
|
@ -812,6 +826,7 @@ func GenerateTLSConfig() *tls.Config {
|
|||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
// nolint: gosec
|
||||
return &tls.Config{
|
||||
Certificates: []tls.Certificate{tlsCert},
|
||||
NextProtos: []string{"argotunnel"},
|
||||
|
|
|
@ -7,12 +7,15 @@ import (
|
|||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
pkgerrors "github.com/pkg/errors"
|
||||
"github.com/quic-go/quic-go"
|
||||
"github.com/rs/zerolog"
|
||||
"go.opentelemetry.io/otel/attribute"
|
||||
"go.opentelemetry.io/otel/trace"
|
||||
"golang.org/x/sync/errgroup"
|
||||
|
||||
cfdsession "github.com/cloudflare/cloudflared/session"
|
||||
|
||||
"github.com/cloudflare/cloudflared/datagramsession"
|
||||
"github.com/cloudflare/cloudflared/ingress"
|
||||
"github.com/cloudflare/cloudflared/management"
|
||||
|
@ -38,10 +41,14 @@ type DatagramSessionHandler interface {
|
|||
}
|
||||
|
||||
type datagramV2Connection struct {
|
||||
conn quic.Connection
|
||||
conn quic.Connection
|
||||
index uint8
|
||||
|
||||
// sessionManager tracks active sessions. It receives datagrams from quic connection via datagramMuxer
|
||||
sessionManager datagramsession.Manager
|
||||
// sessionLimiter tracks active sessions across the tunnel and limits new sessions if they are above the limit.
|
||||
sessionLimiter cfdsession.Limiter
|
||||
|
||||
// datagramMuxer mux/demux datagrams from quic connection
|
||||
datagramMuxer *cfdquic.DatagramMuxerV2
|
||||
packetRouter *ingress.PacketRouter
|
||||
|
@ -58,6 +65,7 @@ func NewDatagramV2Connection(ctx context.Context,
|
|||
index uint8,
|
||||
rpcTimeout time.Duration,
|
||||
streamWriteTimeout time.Duration,
|
||||
sessionLimiter cfdsession.Limiter,
|
||||
logger *zerolog.Logger,
|
||||
) DatagramSessionHandler {
|
||||
sessionDemuxChan := make(chan *packet.Session, demuxChanCapacity)
|
||||
|
@ -66,13 +74,15 @@ func NewDatagramV2Connection(ctx context.Context,
|
|||
packetRouter := ingress.NewPacketRouter(icmpRouter, datagramMuxer, index, logger)
|
||||
|
||||
return &datagramV2Connection{
|
||||
conn,
|
||||
sessionManager,
|
||||
datagramMuxer,
|
||||
packetRouter,
|
||||
rpcTimeout,
|
||||
streamWriteTimeout,
|
||||
logger,
|
||||
conn: conn,
|
||||
index: index,
|
||||
sessionManager: sessionManager,
|
||||
sessionLimiter: sessionLimiter,
|
||||
datagramMuxer: datagramMuxer,
|
||||
packetRouter: packetRouter,
|
||||
rpcTimeout: rpcTimeout,
|
||||
streamWriteTimeout: streamWriteTimeout,
|
||||
logger: logger,
|
||||
}
|
||||
}
|
||||
|
||||
|
@ -109,12 +119,23 @@ func (q *datagramV2Connection) RegisterUdpSession(ctx context.Context, sessionID
|
|||
attribute.String("dst", fmt.Sprintf("%s:%d", dstIP, dstPort)),
|
||||
))
|
||||
log := q.logger.With().Int(management.EventTypeKey, int(management.UDP)).Logger()
|
||||
|
||||
// Try to start a new session
|
||||
if err := q.sessionLimiter.Acquire(management.UDP.String()); err != nil {
|
||||
log.Warn().Msgf("Too many concurrent sessions being handled, rejecting udp proxy to %s:%d", dstIP, dstPort)
|
||||
|
||||
err := pkgerrors.Wrap(err, "failed to start udp session due to rate limiting")
|
||||
tracing.EndWithErrorStatus(registerSpan, err)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Each session is a series of datagram from an eyeball to a dstIP:dstPort.
|
||||
// (src port, dst IP, dst port) uniquely identifies a session, so it needs a dedicated connected socket.
|
||||
originProxy, err := ingress.DialUDP(dstIP, dstPort)
|
||||
if err != nil {
|
||||
log.Err(err).Msgf("Failed to create udp proxy to %s:%d", dstIP, dstPort)
|
||||
tracing.EndWithErrorStatus(registerSpan, err)
|
||||
q.sessionLimiter.Release()
|
||||
return nil, err
|
||||
}
|
||||
registerSpan.SetAttributes(
|
||||
|
@ -127,10 +148,14 @@ func (q *datagramV2Connection) RegisterUdpSession(ctx context.Context, sessionID
|
|||
originProxy.Close()
|
||||
log.Err(err).Str(datagramsession.LogFieldSessionID, datagramsession.FormatSessionID(sessionID)).Msgf("Failed to register udp session")
|
||||
tracing.EndWithErrorStatus(registerSpan, err)
|
||||
q.sessionLimiter.Release()
|
||||
return nil, err
|
||||
}
|
||||
|
||||
go q.serveUDPSession(session, closeAfterIdleHint)
|
||||
go func() {
|
||||
defer q.sessionLimiter.Release() // we do the release here, instead of inside the `serveUDPSession` just to keep all acquire/release calls in the same method.
|
||||
q.serveUDPSession(session, closeAfterIdleHint)
|
||||
}()
|
||||
|
||||
log.Debug().
|
||||
Str(datagramsession.LogFieldSessionID, datagramsession.FormatSessionID(sessionID)).
|
||||
|
@ -170,7 +195,7 @@ func (q *datagramV2Connection) serveUDPSession(session *datagramsession.Session,
|
|||
|
||||
// closeUDPSession first unregisters the session from session manager, then it tries to unregister from edge
|
||||
func (q *datagramV2Connection) closeUDPSession(ctx context.Context, sessionID uuid.UUID, message string) {
|
||||
q.sessionManager.UnregisterSession(ctx, sessionID, message, false)
|
||||
_ = q.sessionManager.UnregisterSession(ctx, sessionID, message, false)
|
||||
quicStream, err := q.conn.OpenStream()
|
||||
if err != nil {
|
||||
// Log this at debug because this is not an error if session was closed due to lost connection
|
||||
|
|
|
@ -0,0 +1,96 @@
|
|||
package connection
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/quic-go/quic-go"
|
||||
"github.com/rs/zerolog"
|
||||
"github.com/stretchr/testify/require"
|
||||
"go.uber.org/mock/gomock"
|
||||
|
||||
"github.com/cloudflare/cloudflared/mocks"
|
||||
cfdsession "github.com/cloudflare/cloudflared/session"
|
||||
)
|
||||
|
||||
type mockQuicConnection struct {
|
||||
}
|
||||
|
||||
func (m *mockQuicConnection) AcceptStream(_ context.Context) (quic.Stream, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (m *mockQuicConnection) AcceptUniStream(_ context.Context) (quic.ReceiveStream, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (m *mockQuicConnection) OpenStream() (quic.Stream, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (m *mockQuicConnection) OpenStreamSync(_ context.Context) (quic.Stream, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (m *mockQuicConnection) OpenUniStream() (quic.SendStream, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (m *mockQuicConnection) OpenUniStreamSync(_ context.Context) (quic.SendStream, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (m *mockQuicConnection) LocalAddr() net.Addr {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *mockQuicConnection) RemoteAddr() net.Addr {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *mockQuicConnection) CloseWithError(_ quic.ApplicationErrorCode, s string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *mockQuicConnection) Context() context.Context {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *mockQuicConnection) ConnectionState() quic.ConnectionState {
|
||||
panic("not meant to be called")
|
||||
}
|
||||
|
||||
func (m *mockQuicConnection) SendDatagram(_ []byte) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *mockQuicConnection) ReceiveDatagram(_ context.Context) ([]byte, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func TestRateLimitOnNewDatagramV2UDPSession(t *testing.T) {
|
||||
log := zerolog.Nop()
|
||||
conn := &mockQuicConnection{}
|
||||
ctrl := gomock.NewController(t)
|
||||
sessionLimiterMock := mocks.NewMockLimiter(ctrl)
|
||||
|
||||
datagramConn := NewDatagramV2Connection(
|
||||
context.Background(),
|
||||
conn,
|
||||
nil,
|
||||
0,
|
||||
0*time.Second,
|
||||
0*time.Second,
|
||||
sessionLimiterMock,
|
||||
&log,
|
||||
)
|
||||
|
||||
sessionLimiterMock.EXPECT().Acquire("udp").Return(cfdsession.ErrTooManyActiveSessions)
|
||||
sessionLimiterMock.EXPECT().Release().Times(0)
|
||||
|
||||
_, err := datagramConn.RegisterUdpSession(context.Background(), uuid.New(), net.IPv4(0, 0, 0, 0), 1000, 1*time.Second, "")
|
||||
require.ErrorIs(t, err, cfdsession.ErrTooManyActiveSessions)
|
||||
}
|
16
go.mod
16
go.mod
|
@ -35,11 +35,12 @@ require (
|
|||
go.opentelemetry.io/otel/trace v1.26.0
|
||||
go.opentelemetry.io/proto/otlp v1.2.0
|
||||
go.uber.org/automaxprocs v1.4.0
|
||||
golang.org/x/crypto v0.23.0
|
||||
golang.org/x/net v0.25.0
|
||||
go.uber.org/mock v0.5.0
|
||||
golang.org/x/crypto v0.24.0
|
||||
golang.org/x/net v0.26.0
|
||||
golang.org/x/sync v0.7.0
|
||||
golang.org/x/sys v0.20.0
|
||||
golang.org/x/term v0.20.0
|
||||
golang.org/x/sys v0.21.0
|
||||
golang.org/x/term v0.21.0
|
||||
google.golang.org/protobuf v1.34.1
|
||||
gopkg.in/natefinch/lumberjack.v2 v2.0.0
|
||||
gopkg.in/yaml.v3 v3.0.1
|
||||
|
@ -83,12 +84,11 @@ require (
|
|||
github.com/prometheus/procfs v0.12.0 // indirect
|
||||
github.com/russross/blackfriday/v2 v2.1.0 // indirect
|
||||
go.opentelemetry.io/otel/metric v1.26.0 // indirect
|
||||
go.uber.org/mock v0.4.0 // indirect
|
||||
golang.org/x/exp v0.0.0-20240506185415-9bf2ced13842 // indirect
|
||||
golang.org/x/mod v0.17.0 // indirect
|
||||
golang.org/x/mod v0.18.0 // indirect
|
||||
golang.org/x/oauth2 v0.18.0 // indirect
|
||||
golang.org/x/text v0.15.0 // indirect
|
||||
golang.org/x/tools v0.21.0 // indirect
|
||||
golang.org/x/text v0.16.0 // indirect
|
||||
golang.org/x/tools v0.22.0 // indirect
|
||||
google.golang.org/appengine v1.6.8 // indirect
|
||||
google.golang.org/genproto/googleapis/api v0.0.0-20240311132316-a219d84964c2 // indirect
|
||||
google.golang.org/genproto/googleapis/rpc v0.0.0-20240318140521-94a12d6c2237 // indirect
|
||||
|
|
32
go.sum
32
go.sum
|
@ -217,26 +217,26 @@ go.opentelemetry.io/proto/otlp v1.2.0 h1:pVeZGk7nXDC9O2hncA6nHldxEjm6LByfA2aN8IO
|
|||
go.opentelemetry.io/proto/otlp v1.2.0/go.mod h1:gGpR8txAl5M03pDhMC79G6SdqNV26naRm/KDsgaHD8A=
|
||||
go.uber.org/automaxprocs v1.4.0 h1:CpDZl6aOlLhReez+8S3eEotD7Jx0Os++lemPlMULQP0=
|
||||
go.uber.org/automaxprocs v1.4.0/go.mod h1:/mTEdr7LvHhs0v7mjdxDreTz1OG5zdZGqgOnhWiR/+Q=
|
||||
go.uber.org/mock v0.4.0 h1:VcM4ZOtdbR4f6VXfiOpwpVJDL6lCReaZ6mw31wqh7KU=
|
||||
go.uber.org/mock v0.4.0/go.mod h1:a6FSlNadKUHUa9IP5Vyt1zh4fC7uAwxMutEAscFbkZc=
|
||||
go.uber.org/mock v0.5.0 h1:KAMbZvZPyBPWgD14IrIQ38QCyjwpvVVV6K/bHl1IwQU=
|
||||
go.uber.org/mock v0.5.0/go.mod h1:ge71pBPLYDk7QIi1LupWxdAykm7KIEFchiOqd6z7qMM=
|
||||
golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
|
||||
golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
|
||||
golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=
|
||||
golang.org/x/crypto v0.23.0 h1:dIJU/v2J8Mdglj/8rJ6UUOM3Zc9zLZxVZwwxMooUSAI=
|
||||
golang.org/x/crypto v0.23.0/go.mod h1:CKFgDieR+mRhux2Lsu27y0fO304Db0wZe70UKqHu0v8=
|
||||
golang.org/x/crypto v0.24.0 h1:mnl8DM0o513X8fdIkmyFE/5hTYxbwYOjDS/+rK6qpRI=
|
||||
golang.org/x/crypto v0.24.0/go.mod h1:Z1PMYSOR5nyMcyAVAIQSKCDwalqy85Aqn1x3Ws4L5DM=
|
||||
golang.org/x/exp v0.0.0-20240506185415-9bf2ced13842 h1:vr/HnozRka3pE4EsMEg1lgkXJkTFJCVUX+S/ZT6wYzM=
|
||||
golang.org/x/exp v0.0.0-20240506185415-9bf2ced13842/go.mod h1:XtvwrStGgqGPLc4cjQfWqZHG1YFdYs6swckp8vpsjnc=
|
||||
golang.org/x/lint v0.0.0-20200302205851-738671d3881b/go.mod h1:3xt1FjdF8hUf6vQPIChWIBhFzV8gjjsPE/fR3IyQdNY=
|
||||
golang.org/x/mod v0.1.1-0.20191105210325-c90efee705ee/go.mod h1:QqPTAvyqsEbceGzBzNggFXnrqF1CaUcvgkdR5Ot7KZg=
|
||||
golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4=
|
||||
golang.org/x/mod v0.17.0 h1:zY54UmvipHiNd+pm+m0x9KhZ9hl1/7QNMyxXbc6ICqA=
|
||||
golang.org/x/mod v0.17.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c=
|
||||
golang.org/x/mod v0.18.0 h1:5+9lSbEzPSdWkH32vYPBwEpX8KwDbM52Ud9xBUvNlb0=
|
||||
golang.org/x/mod v0.18.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c=
|
||||
golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
|
||||
golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
|
||||
golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
|
||||
golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug+ECip1KBveYUHfp+8e9klMJ9c=
|
||||
golang.org/x/net v0.25.0 h1:d/OCCoBEUq33pjydKrGQhw7IlUPI2Oylr+8qLx49kac=
|
||||
golang.org/x/net v0.25.0/go.mod h1:JkAGAh7GEvH74S6FOH42FLoXpXbE/aqXSrIQjXgsiwM=
|
||||
golang.org/x/net v0.26.0 h1:soB7SVo0PWrY4vPW/+ay0jKDNScG2X9wFeYlXIvJsOQ=
|
||||
golang.org/x/net v0.26.0/go.mod h1:5YKkiSynbBIh3p6iOc/vibscux0x38BZDkn8sCUPxHE=
|
||||
golang.org/x/oauth2 v0.18.0 h1:09qnuIAgzdx1XplqJvW6CQqMCtGZykZWcXzPMPUusvI=
|
||||
golang.org/x/oauth2 v0.18.0/go.mod h1:Wf7knwG0MPoWIMMBgFlEaSUDaKskp0dCfrlJRJXbBi8=
|
||||
golang.org/x/sync v0.0.0-20181221193216-37e7f081c4d4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
|
||||
|
@ -254,19 +254,19 @@ golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBc
|
|||
golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.20.0 h1:Od9JTbYCk261bKm4M/mw7AklTlFYIa0bIp9BgSm1S8Y=
|
||||
golang.org/x/sys v0.20.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
|
||||
golang.org/x/sys v0.21.0 h1:rF+pYz3DAGSQAxAu1CbC7catZg4ebC4UIeIhKxBZvws=
|
||||
golang.org/x/sys v0.21.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
|
||||
golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
|
||||
golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
|
||||
golang.org/x/term v0.20.0 h1:VnkxpohqXaOBYJtBmEppKUG6mXpi+4O6purfc2+sMhw=
|
||||
golang.org/x/term v0.20.0/go.mod h1:8UkIAJTvZgivsXaD6/pH6U9ecQzZ45awqEOzuCvwpFY=
|
||||
golang.org/x/term v0.21.0 h1:WVXCp+/EBEHOj53Rvu+7KiT/iElMrO8ACK16SMZ3jaA=
|
||||
golang.org/x/term v0.21.0/go.mod h1:ooXLefLobQVslOqselCNF4SxFAaoS6KujMbsGzSDmX0=
|
||||
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
|
||||
golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk=
|
||||
golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
|
||||
golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
|
||||
golang.org/x/text v0.3.8/go.mod h1:E6s5w1FMmriuDzIBO73fBruAKo1PCIq6d2Q6DHfQ8WQ=
|
||||
golang.org/x/text v0.15.0 h1:h1V/4gjBv8v9cjcR6+AR5+/cIYK5N/WAgiv4xlsEtAk=
|
||||
golang.org/x/text v0.15.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
|
||||
golang.org/x/text v0.16.0 h1:a94ExnEXNtEwYLGJSIUxnWoxoRz/ZcCsV63ROupILh4=
|
||||
golang.org/x/text v0.16.0/go.mod h1:GhwF1Be+LQoKShO3cGOHzqOgRrGaYc9AvblQOmPVHnI=
|
||||
golang.org/x/time v0.0.0-20191024005414-555d28b269f0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
|
||||
golang.org/x/time v0.5.0 h1:o7cqy6amK/52YcAKIPlM3a+Fpj35zvRj2TP+e1xFSfk=
|
||||
golang.org/x/time v0.5.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM=
|
||||
|
@ -275,8 +275,8 @@ golang.org/x/tools v0.0.0-20190828213141-aed303cbaa74/go.mod h1:b+2E5dAYhXwXZwtn
|
|||
golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
|
||||
golang.org/x/tools v0.0.0-20200130002326-2f3ba24bd6e7/go.mod h1:TB2adYChydJhpapKDTa4BR/hXlZSLoq2Wpct/0txZ28=
|
||||
golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc=
|
||||
golang.org/x/tools v0.21.0 h1:qc0xYgIbsSDt9EyWz05J5wfa7LOVW0YTLOXrqdLAWIw=
|
||||
golang.org/x/tools v0.21.0/go.mod h1:aiJjzUbINMkxbQROHiO6hDPo2LHcIPhhQsa9DLh0yGk=
|
||||
golang.org/x/tools v0.22.0 h1:gqSGLZqv+AI9lIQzniJ0nZDRG5GBPsSi+DRNHWNz6yA=
|
||||
golang.org/x/tools v0.22.0/go.mod h1:aCwcsjqvq7Yqt6TNyX7QMU2enbQ/Gt0bo6krSeEri+c=
|
||||
golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
|
||||
golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
|
||||
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
|
||||
|
|
|
@ -0,0 +1,150 @@
|
|||
// Code generated by MockGen. DO NOT EDIT.
|
||||
// Source: ../session/limiter.go
|
||||
//
|
||||
// Generated by this command:
|
||||
//
|
||||
// mockgen -typed -build_flags=-tags=gomock -package mocks -destination mock_limiter.go -source=../session/limiter.go Limiter
|
||||
//
|
||||
|
||||
// Package mocks is a generated GoMock package.
|
||||
package mocks
|
||||
|
||||
import (
|
||||
reflect "reflect"
|
||||
|
||||
gomock "go.uber.org/mock/gomock"
|
||||
)
|
||||
|
||||
// MockLimiter is a mock of Limiter interface.
|
||||
type MockLimiter struct {
|
||||
ctrl *gomock.Controller
|
||||
recorder *MockLimiterMockRecorder
|
||||
isgomock struct{}
|
||||
}
|
||||
|
||||
// MockLimiterMockRecorder is the mock recorder for MockLimiter.
|
||||
type MockLimiterMockRecorder struct {
|
||||
mock *MockLimiter
|
||||
}
|
||||
|
||||
// NewMockLimiter creates a new mock instance.
|
||||
func NewMockLimiter(ctrl *gomock.Controller) *MockLimiter {
|
||||
mock := &MockLimiter{ctrl: ctrl}
|
||||
mock.recorder = &MockLimiterMockRecorder{mock}
|
||||
return mock
|
||||
}
|
||||
|
||||
// EXPECT returns an object that allows the caller to indicate expected use.
|
||||
func (m *MockLimiter) EXPECT() *MockLimiterMockRecorder {
|
||||
return m.recorder
|
||||
}
|
||||
|
||||
// Acquire mocks base method.
|
||||
func (m *MockLimiter) Acquire(sessionType string) error {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "Acquire", sessionType)
|
||||
ret0, _ := ret[0].(error)
|
||||
return ret0
|
||||
}
|
||||
|
||||
// Acquire indicates an expected call of Acquire.
|
||||
func (mr *MockLimiterMockRecorder) Acquire(sessionType any) *MockLimiterAcquireCall {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Acquire", reflect.TypeOf((*MockLimiter)(nil).Acquire), sessionType)
|
||||
return &MockLimiterAcquireCall{Call: call}
|
||||
}
|
||||
|
||||
// MockLimiterAcquireCall wrap *gomock.Call
|
||||
type MockLimiterAcquireCall struct {
|
||||
*gomock.Call
|
||||
}
|
||||
|
||||
// Return rewrite *gomock.Call.Return
|
||||
func (c *MockLimiterAcquireCall) Return(arg0 error) *MockLimiterAcquireCall {
|
||||
c.Call = c.Call.Return(arg0)
|
||||
return c
|
||||
}
|
||||
|
||||
// Do rewrite *gomock.Call.Do
|
||||
func (c *MockLimiterAcquireCall) Do(f func(string) error) *MockLimiterAcquireCall {
|
||||
c.Call = c.Call.Do(f)
|
||||
return c
|
||||
}
|
||||
|
||||
// DoAndReturn rewrite *gomock.Call.DoAndReturn
|
||||
func (c *MockLimiterAcquireCall) DoAndReturn(f func(string) error) *MockLimiterAcquireCall {
|
||||
c.Call = c.Call.DoAndReturn(f)
|
||||
return c
|
||||
}
|
||||
|
||||
// Release mocks base method.
|
||||
func (m *MockLimiter) Release() {
|
||||
m.ctrl.T.Helper()
|
||||
m.ctrl.Call(m, "Release")
|
||||
}
|
||||
|
||||
// Release indicates an expected call of Release.
|
||||
func (mr *MockLimiterMockRecorder) Release() *MockLimiterReleaseCall {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Release", reflect.TypeOf((*MockLimiter)(nil).Release))
|
||||
return &MockLimiterReleaseCall{Call: call}
|
||||
}
|
||||
|
||||
// MockLimiterReleaseCall wrap *gomock.Call
|
||||
type MockLimiterReleaseCall struct {
|
||||
*gomock.Call
|
||||
}
|
||||
|
||||
// Return rewrite *gomock.Call.Return
|
||||
func (c *MockLimiterReleaseCall) Return() *MockLimiterReleaseCall {
|
||||
c.Call = c.Call.Return()
|
||||
return c
|
||||
}
|
||||
|
||||
// Do rewrite *gomock.Call.Do
|
||||
func (c *MockLimiterReleaseCall) Do(f func()) *MockLimiterReleaseCall {
|
||||
c.Call = c.Call.Do(f)
|
||||
return c
|
||||
}
|
||||
|
||||
// DoAndReturn rewrite *gomock.Call.DoAndReturn
|
||||
func (c *MockLimiterReleaseCall) DoAndReturn(f func()) *MockLimiterReleaseCall {
|
||||
c.Call = c.Call.DoAndReturn(f)
|
||||
return c
|
||||
}
|
||||
|
||||
// SetLimit mocks base method.
|
||||
func (m *MockLimiter) SetLimit(arg0 uint64) {
|
||||
m.ctrl.T.Helper()
|
||||
m.ctrl.Call(m, "SetLimit", arg0)
|
||||
}
|
||||
|
||||
// SetLimit indicates an expected call of SetLimit.
|
||||
func (mr *MockLimiterMockRecorder) SetLimit(arg0 any) *MockLimiterSetLimitCall {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetLimit", reflect.TypeOf((*MockLimiter)(nil).SetLimit), arg0)
|
||||
return &MockLimiterSetLimitCall{Call: call}
|
||||
}
|
||||
|
||||
// MockLimiterSetLimitCall wrap *gomock.Call
|
||||
type MockLimiterSetLimitCall struct {
|
||||
*gomock.Call
|
||||
}
|
||||
|
||||
// Return rewrite *gomock.Call.Return
|
||||
func (c *MockLimiterSetLimitCall) Return() *MockLimiterSetLimitCall {
|
||||
c.Call = c.Call.Return()
|
||||
return c
|
||||
}
|
||||
|
||||
// Do rewrite *gomock.Call.Do
|
||||
func (c *MockLimiterSetLimitCall) Do(f func(uint64)) *MockLimiterSetLimitCall {
|
||||
c.Call = c.Call.Do(f)
|
||||
return c
|
||||
}
|
||||
|
||||
// DoAndReturn rewrite *gomock.Call.DoAndReturn
|
||||
func (c *MockLimiterSetLimitCall) DoAndReturn(f func(uint64)) *MockLimiterSetLimitCall {
|
||||
c.Call = c.Call.DoAndReturn(f)
|
||||
return c
|
||||
}
|
|
@ -0,0 +1,5 @@
|
|||
//go:build gomock || generate
|
||||
|
||||
package mocks
|
||||
|
||||
//go:generate sh -c "go run go.uber.org/mock/mockgen -typed -build_flags=\"-tags=gomock\" -package mocks -destination mock_limiter.go -source=../session/limiter.go Limiter"
|
|
@ -10,6 +10,8 @@ import (
|
|||
"github.com/pkg/errors"
|
||||
"github.com/rs/zerolog"
|
||||
|
||||
cfdsession "github.com/cloudflare/cloudflared/session"
|
||||
|
||||
"github.com/cloudflare/cloudflared/config"
|
||||
"github.com/cloudflare/cloudflared/connection"
|
||||
"github.com/cloudflare/cloudflared/ingress"
|
||||
|
@ -33,7 +35,9 @@ type Orchestrator struct {
|
|||
// cloudflared Configuration
|
||||
config *Config
|
||||
tags []pogs.Tag
|
||||
log *zerolog.Logger
|
||||
// sessionLimiter tracks active sessions across the tunnel and limits new sessions if they are above the limit.
|
||||
sessionLimiter cfdsession.Limiter
|
||||
log *zerolog.Logger
|
||||
|
||||
// orchestrator must not handle any more updates after shutdownC is closed
|
||||
shutdownC <-chan struct{}
|
||||
|
@ -54,6 +58,7 @@ func NewOrchestrator(ctx context.Context,
|
|||
internalRules: internalRules,
|
||||
config: config,
|
||||
tags: tags,
|
||||
sessionLimiter: cfdsession.NewLimiter(0),
|
||||
log: log,
|
||||
shutdownC: ctx.Done(),
|
||||
}
|
||||
|
@ -208,6 +213,12 @@ func (o *Orchestrator) GetOriginProxy() (connection.OriginProxy, error) {
|
|||
return proxy, nil
|
||||
}
|
||||
|
||||
// GetSessionLimiter returns the session limiter used across cloudflared, that can be hot reload when
|
||||
// the configuration changes.
|
||||
func (o *Orchestrator) GetSessionLimiter() cfdsession.Limiter {
|
||||
return o.sessionLimiter
|
||||
}
|
||||
|
||||
func (o *Orchestrator) waitToCloseLastProxy() {
|
||||
<-o.shutdownC
|
||||
o.lock.Lock()
|
||||
|
|
|
@ -116,7 +116,7 @@ func (s *UDPSessionRegistrationDatagram) MarshalBinary() (data []byte, err error
|
|||
data = make([]byte, sessionRegistrationIPv4DatagramHeaderLen+len(s.Payload))
|
||||
}
|
||||
data[0] = byte(UDPSessionRegistrationType)
|
||||
data[1] = byte(flags)
|
||||
data[1] = flags
|
||||
binary.BigEndian.PutUint16(data[2:4], s.Dest.Port())
|
||||
binary.BigEndian.PutUint16(data[4:6], uint16(s.IdleDurationHint.Seconds()))
|
||||
err = s.RequestID.MarshalBinaryTo(data[6:22])
|
||||
|
@ -284,6 +284,8 @@ const (
|
|||
ResponseDestinationUnreachable SessionRegistrationResp = 0x01
|
||||
// Session registration was unable to bind to a local UDP socket.
|
||||
ResponseUnableToBindSocket SessionRegistrationResp = 0x02
|
||||
// Session registration failed due to the number of session being higher than the limit.
|
||||
ResponseTooManyActiveSessions SessionRegistrationResp = 0x03
|
||||
// Session registration failed with an unexpected error but provided a message.
|
||||
ResponseErrorWithMsg SessionRegistrationResp = 0xff
|
||||
)
|
||||
|
@ -311,6 +313,7 @@ func (s *UDPSessionRegistrationResponseDatagram) MarshalBinary() (data []byte, e
|
|||
if len(s.ErrorMsg) > maxResponseErrorMessageLen {
|
||||
return nil, wrapMarshalErr(ErrDatagramResponseMsgInvalidSize)
|
||||
}
|
||||
// nolint: gosec
|
||||
errMsgLen := uint16(len(s.ErrorMsg))
|
||||
|
||||
data = make([]byte, datagramSessionRegistrationResponseLen+errMsgLen)
|
||||
|
|
|
@ -7,6 +7,10 @@ import (
|
|||
"sync"
|
||||
|
||||
"github.com/rs/zerolog"
|
||||
|
||||
"github.com/cloudflare/cloudflared/management"
|
||||
|
||||
cfdsession "github.com/cloudflare/cloudflared/session"
|
||||
)
|
||||
|
||||
var (
|
||||
|
@ -16,6 +20,8 @@ var (
|
|||
ErrSessionBoundToOtherConn = errors.New("flow is in use by another connection")
|
||||
// ErrSessionAlreadyRegistered is returned when a registration already exists for this connection.
|
||||
ErrSessionAlreadyRegistered = errors.New("flow is already registered for this connection")
|
||||
// ErrSessionRegistrationRateLimited is returned when a registration fails due to rate limiting on the number of active sessions.
|
||||
ErrSessionRegistrationRateLimited = errors.New("flow registration rate limited")
|
||||
)
|
||||
|
||||
type SessionManager interface {
|
||||
|
@ -38,14 +44,16 @@ type sessionManager struct {
|
|||
sessions map[RequestID]Session
|
||||
mutex sync.RWMutex
|
||||
originDialer DialUDP
|
||||
limiter cfdsession.Limiter
|
||||
metrics Metrics
|
||||
log *zerolog.Logger
|
||||
}
|
||||
|
||||
func NewSessionManager(metrics Metrics, log *zerolog.Logger, originDialer DialUDP) SessionManager {
|
||||
func NewSessionManager(metrics Metrics, log *zerolog.Logger, originDialer DialUDP, limiter cfdsession.Limiter) SessionManager {
|
||||
return &sessionManager{
|
||||
sessions: make(map[RequestID]Session),
|
||||
originDialer: originDialer,
|
||||
limiter: limiter,
|
||||
metrics: metrics,
|
||||
log: log,
|
||||
}
|
||||
|
@ -61,6 +69,12 @@ func (s *sessionManager) RegisterSession(request *UDPSessionRegistrationDatagram
|
|||
}
|
||||
return nil, ErrSessionBoundToOtherConn
|
||||
}
|
||||
|
||||
// Try to start a new session
|
||||
if err := s.limiter.Acquire(management.UDP.String()); err != nil {
|
||||
return nil, ErrSessionRegistrationRateLimited
|
||||
}
|
||||
|
||||
// Attempt to bind the UDP socket for the new session
|
||||
origin, err := s.originDialer(request.Dest)
|
||||
if err != nil {
|
||||
|
@ -100,4 +114,5 @@ func (s *sessionManager) UnregisterSession(requestID RequestID) {
|
|||
_ = session.Close()
|
||||
}
|
||||
delete(s.sessions, requestID)
|
||||
s.limiter.Release()
|
||||
}
|
||||
|
|
|
@ -8,14 +8,19 @@ import (
|
|||
"time"
|
||||
|
||||
"github.com/rs/zerolog"
|
||||
"github.com/stretchr/testify/require"
|
||||
"go.uber.org/mock/gomock"
|
||||
|
||||
"github.com/cloudflare/cloudflared/mocks"
|
||||
|
||||
"github.com/cloudflare/cloudflared/ingress"
|
||||
v3 "github.com/cloudflare/cloudflared/quic/v3"
|
||||
cfdsession "github.com/cloudflare/cloudflared/session"
|
||||
)
|
||||
|
||||
func TestRegisterSession(t *testing.T) {
|
||||
log := zerolog.Nop()
|
||||
manager := v3.NewSessionManager(&noopMetrics{}, &log, ingress.DialUDPAddrPort)
|
||||
manager := v3.NewSessionManager(&noopMetrics{}, &log, ingress.DialUDPAddrPort, cfdsession.NewLimiter(0))
|
||||
|
||||
request := v3.UDPSessionRegistrationDatagram{
|
||||
RequestID: testRequestID,
|
||||
|
@ -71,10 +76,32 @@ func TestRegisterSession(t *testing.T) {
|
|||
|
||||
func TestGetSession_Empty(t *testing.T) {
|
||||
log := zerolog.Nop()
|
||||
manager := v3.NewSessionManager(&noopMetrics{}, &log, ingress.DialUDPAddrPort)
|
||||
manager := v3.NewSessionManager(&noopMetrics{}, &log, ingress.DialUDPAddrPort, cfdsession.NewLimiter(0))
|
||||
|
||||
_, err := manager.GetSession(testRequestID)
|
||||
if !errors.Is(err, v3.ErrSessionNotFound) {
|
||||
t.Fatalf("get session find no session: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegisterSessionRateLimit(t *testing.T) {
|
||||
log := zerolog.Nop()
|
||||
ctrl := gomock.NewController(t)
|
||||
|
||||
sessionLimiterMock := mocks.NewMockLimiter(ctrl)
|
||||
|
||||
sessionLimiterMock.EXPECT().Acquire("udp").Return(cfdsession.ErrTooManyActiveSessions)
|
||||
sessionLimiterMock.EXPECT().Release().Times(0)
|
||||
|
||||
manager := v3.NewSessionManager(&noopMetrics{}, &log, ingress.DialUDPAddrPort, sessionLimiterMock)
|
||||
|
||||
request := v3.UDPSessionRegistrationDatagram{
|
||||
RequestID: testRequestID,
|
||||
Dest: netip.MustParseAddrPort("127.0.0.1:5000"),
|
||||
Traced: false,
|
||||
IdleDurationHint: 5 * time.Second,
|
||||
Payload: nil,
|
||||
}
|
||||
_, err := manager.RegisterSession(&request, &noopEyeball{})
|
||||
require.ErrorIs(t, err, v3.ErrSessionRegistrationRateLimited)
|
||||
}
|
||||
|
|
|
@ -143,8 +143,6 @@ func (c *datagramConn) SendICMPTTLExceed(icmp *packet.ICMP, rawPacket packet.Raw
|
|||
return c.SendICMPPacket(c.icmpRouter.ConvertToTTLExceeded(icmp, rawPacket))
|
||||
}
|
||||
|
||||
var errReadTimeout error = errors.New("receive datagram timeout")
|
||||
|
||||
// pollDatagrams will read datagrams from the underlying connection until the provided context is done.
|
||||
func (c *datagramConn) pollDatagrams(ctx context.Context) {
|
||||
for ctx.Err() == nil {
|
||||
|
@ -256,8 +254,12 @@ func (c *datagramConn) handleSessionRegistrationDatagram(ctx context.Context, da
|
|||
// Session is already registered but to a different connection
|
||||
c.handleSessionMigration(datagram.RequestID, &log)
|
||||
return
|
||||
case ErrSessionRegistrationRateLimited:
|
||||
// There are too many concurrent sessions so we return an error to force a retry later
|
||||
c.handleSessionRegistrationRateLimited(datagram, &log)
|
||||
return
|
||||
default:
|
||||
log.Err(err).Msgf("flow registration failure")
|
||||
log.Err(err).Msg("flow registration failure")
|
||||
c.handleSessionRegistrationFailure(datagram.RequestID, &log)
|
||||
return
|
||||
}
|
||||
|
@ -278,7 +280,7 @@ func (c *datagramConn) handleSessionRegistrationDatagram(ctx context.Context, da
|
|||
// [Session.Serve] is blocking and will continue this go routine till the end of the session lifetime.
|
||||
start := time.Now()
|
||||
err = session.Serve(ctx)
|
||||
elapsedMS := time.Now().Sub(start).Milliseconds()
|
||||
elapsedMS := time.Since(start).Milliseconds()
|
||||
log = log.With().Int64(logDurationKey, elapsedMS).Logger()
|
||||
if err == nil {
|
||||
// We typically don't expect a session to close without some error response. [SessionIdleErr] is the typical
|
||||
|
@ -346,6 +348,16 @@ func (c *datagramConn) handleSessionRegistrationFailure(requestID RequestID, log
|
|||
}
|
||||
}
|
||||
|
||||
func (c *datagramConn) handleSessionRegistrationRateLimited(datagram *UDPSessionRegistrationDatagram, logger *zerolog.Logger) {
|
||||
c.logger.Warn().Msg("Too many concurrent sessions being handled, rejecting udp proxy")
|
||||
|
||||
rateLimitResponse := ResponseTooManyActiveSessions
|
||||
err := c.SendUDPSessionResponse(datagram.RequestID, rateLimitResponse)
|
||||
if err != nil {
|
||||
logger.Err(err).Msgf("unable to send flow registration error response (%d)", rateLimitResponse)
|
||||
}
|
||||
}
|
||||
|
||||
// Handles incoming datagrams that need to be sent to a registered session.
|
||||
func (c *datagramConn) handleSessionPayloadDatagram(datagram *UDPSessionPayloadDatagram, logger *zerolog.Logger) {
|
||||
s, err := c.sessionManager.GetSession(datagram.RequestID)
|
||||
|
|
|
@ -13,16 +13,17 @@ import (
|
|||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
|
||||
"github.com/google/gopacket/layers"
|
||||
"github.com/rs/zerolog"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"golang.org/x/net/icmp"
|
||||
"golang.org/x/net/ipv4"
|
||||
|
||||
"github.com/cloudflare/cloudflared/ingress"
|
||||
"github.com/cloudflare/cloudflared/packet"
|
||||
v3 "github.com/cloudflare/cloudflared/quic/v3"
|
||||
cfdsession "github.com/cloudflare/cloudflared/session"
|
||||
)
|
||||
|
||||
type noopEyeball struct {
|
||||
|
@ -87,7 +88,7 @@ func (m *mockEyeball) SendICMPTTLExceed(icmp *packet.ICMP, rawPacket packet.RawP
|
|||
|
||||
func TestDatagramConn_New(t *testing.T) {
|
||||
log := zerolog.Nop()
|
||||
conn := v3.NewDatagramConn(newMockQuicConn(), v3.NewSessionManager(&noopMetrics{}, &log, ingress.DialUDPAddrPort), &noopICMPRouter{}, 0, &noopMetrics{}, &log)
|
||||
conn := v3.NewDatagramConn(newMockQuicConn(), v3.NewSessionManager(&noopMetrics{}, &log, ingress.DialUDPAddrPort, cfdsession.NewLimiter(0)), &noopICMPRouter{}, 0, &noopMetrics{}, &log)
|
||||
if conn == nil {
|
||||
t.Fatal("expected valid connection")
|
||||
}
|
||||
|
@ -96,10 +97,12 @@ func TestDatagramConn_New(t *testing.T) {
|
|||
func TestDatagramConn_SendUDPSessionDatagram(t *testing.T) {
|
||||
log := zerolog.Nop()
|
||||
quic := newMockQuicConn()
|
||||
conn := v3.NewDatagramConn(quic, v3.NewSessionManager(&noopMetrics{}, &log, ingress.DialUDPAddrPort), &noopICMPRouter{}, 0, &noopMetrics{}, &log)
|
||||
conn := v3.NewDatagramConn(quic, v3.NewSessionManager(&noopMetrics{}, &log, ingress.DialUDPAddrPort, cfdsession.NewLimiter(0)), &noopICMPRouter{}, 0, &noopMetrics{}, &log)
|
||||
|
||||
payload := []byte{0xef, 0xef}
|
||||
conn.SendUDPSessionDatagram(payload)
|
||||
err := conn.SendUDPSessionDatagram(payload)
|
||||
require.NoError(t, err)
|
||||
|
||||
p := <-quic.recv
|
||||
if !slices.Equal(p, payload) {
|
||||
t.Fatal("datagram sent does not match datagram received on quic side")
|
||||
|
@ -109,15 +112,16 @@ func TestDatagramConn_SendUDPSessionDatagram(t *testing.T) {
|
|||
func TestDatagramConn_SendUDPSessionResponse(t *testing.T) {
|
||||
log := zerolog.Nop()
|
||||
quic := newMockQuicConn()
|
||||
conn := v3.NewDatagramConn(quic, v3.NewSessionManager(&noopMetrics{}, &log, ingress.DialUDPAddrPort), &noopICMPRouter{}, 0, &noopMetrics{}, &log)
|
||||
conn := v3.NewDatagramConn(quic, v3.NewSessionManager(&noopMetrics{}, &log, ingress.DialUDPAddrPort, cfdsession.NewLimiter(0)), &noopICMPRouter{}, 0, &noopMetrics{}, &log)
|
||||
|
||||
err := conn.SendUDPSessionResponse(testRequestID, v3.ResponseDestinationUnreachable)
|
||||
require.NoError(t, err)
|
||||
|
||||
conn.SendUDPSessionResponse(testRequestID, v3.ResponseDestinationUnreachable)
|
||||
resp := <-quic.recv
|
||||
var response v3.UDPSessionRegistrationResponseDatagram
|
||||
err := response.UnmarshalBinary(resp)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
err = response.UnmarshalBinary(resp)
|
||||
require.NoError(t, err)
|
||||
|
||||
expected := v3.UDPSessionRegistrationResponseDatagram{
|
||||
RequestID: testRequestID,
|
||||
ResponseType: v3.ResponseDestinationUnreachable,
|
||||
|
@ -130,7 +134,7 @@ func TestDatagramConn_SendUDPSessionResponse(t *testing.T) {
|
|||
func TestDatagramConnServe_ApplicationClosed(t *testing.T) {
|
||||
log := zerolog.Nop()
|
||||
quic := newMockQuicConn()
|
||||
conn := v3.NewDatagramConn(quic, v3.NewSessionManager(&noopMetrics{}, &log, ingress.DialUDPAddrPort), &noopICMPRouter{}, 0, &noopMetrics{}, &log)
|
||||
conn := v3.NewDatagramConn(quic, v3.NewSessionManager(&noopMetrics{}, &log, ingress.DialUDPAddrPort, cfdsession.NewLimiter(0)), &noopICMPRouter{}, 0, &noopMetrics{}, &log)
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
|
||||
defer cancel()
|
||||
|
@ -146,7 +150,7 @@ func TestDatagramConnServe_ConnectionClosed(t *testing.T) {
|
|||
ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
|
||||
defer cancel()
|
||||
quic.ctx = ctx
|
||||
conn := v3.NewDatagramConn(quic, v3.NewSessionManager(&noopMetrics{}, &log, ingress.DialUDPAddrPort), &noopICMPRouter{}, 0, &noopMetrics{}, &log)
|
||||
conn := v3.NewDatagramConn(quic, v3.NewSessionManager(&noopMetrics{}, &log, ingress.DialUDPAddrPort, cfdsession.NewLimiter(0)), &noopICMPRouter{}, 0, &noopMetrics{}, &log)
|
||||
|
||||
err := conn.Serve(context.Background())
|
||||
if !errors.Is(err, context.DeadlineExceeded) {
|
||||
|
@ -157,7 +161,7 @@ func TestDatagramConnServe_ConnectionClosed(t *testing.T) {
|
|||
func TestDatagramConnServe_ReceiveDatagramError(t *testing.T) {
|
||||
log := zerolog.Nop()
|
||||
quic := &mockQuicConnReadError{err: net.ErrClosed}
|
||||
conn := v3.NewDatagramConn(quic, v3.NewSessionManager(&noopMetrics{}, &log, ingress.DialUDPAddrPort), &noopICMPRouter{}, 0, &noopMetrics{}, &log)
|
||||
conn := v3.NewDatagramConn(quic, v3.NewSessionManager(&noopMetrics{}, &log, ingress.DialUDPAddrPort, cfdsession.NewLimiter(0)), &noopICMPRouter{}, 0, &noopMetrics{}, &log)
|
||||
|
||||
err := conn.Serve(context.Background())
|
||||
if !errors.Is(err, net.ErrClosed) {
|
||||
|
@ -165,6 +169,38 @@ func TestDatagramConnServe_ReceiveDatagramError(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestDatagramConnServe_SessionRegistrationRateLimit(t *testing.T) {
|
||||
log := zerolog.Nop()
|
||||
quic := newMockQuicConn()
|
||||
sessionManager := &mockSessionManager{
|
||||
expectedRegErr: v3.ErrSessionRegistrationRateLimited,
|
||||
}
|
||||
conn := v3.NewDatagramConn(quic, sessionManager, &noopICMPRouter{}, 0, &noopMetrics{}, &log)
|
||||
|
||||
// Setup the muxer
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
done := make(chan error, 1)
|
||||
go func() {
|
||||
done <- conn.Serve(ctx)
|
||||
}()
|
||||
|
||||
// Send new session registration
|
||||
datagram := newRegisterSessionDatagram(testRequestID)
|
||||
quic.send <- datagram
|
||||
|
||||
// Wait for session registration response with failure
|
||||
datagram = <-quic.recv
|
||||
var resp v3.UDPSessionRegistrationResponseDatagram
|
||||
err := resp.UnmarshalBinary(datagram)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
require.EqualValues(t, testRequestID, resp.RequestID)
|
||||
require.EqualValues(t, v3.ResponseTooManyActiveSessions, resp.ResponseType)
|
||||
}
|
||||
|
||||
func TestDatagramConnServe_ErrorDatagramTypes(t *testing.T) {
|
||||
for _, test := range []struct {
|
||||
name string
|
||||
|
@ -354,11 +390,9 @@ func TestDatagramConnServeDecodeMultipleICMPInParallel(t *testing.T) {
|
|||
var receivedPackets []*packet.ICMP
|
||||
go func() {
|
||||
for ctx.Err() == nil {
|
||||
select {
|
||||
case icmpPacket := <-router.recv:
|
||||
receivedPackets = append(receivedPackets, icmpPacket)
|
||||
wg.Done()
|
||||
}
|
||||
icmpPacket := <-router.recv
|
||||
receivedPackets = append(receivedPackets, icmpPacket)
|
||||
wg.Done()
|
||||
}
|
||||
}()
|
||||
|
||||
|
@ -677,7 +711,7 @@ func TestDatagramConnServe_ICMPDatagram_TTLExceeded(t *testing.T) {
|
|||
datagram := newICMPDatagram(expectedICMP)
|
||||
quic.send <- datagram
|
||||
|
||||
// Origin should not recieve a packet
|
||||
// Origin should not receive a packet
|
||||
select {
|
||||
case <-router.recv:
|
||||
t.Fatalf("TTL should be expired and no origin ICMP sent")
|
||||
|
@ -719,18 +753,6 @@ func newRegisterSessionDatagram(id v3.RequestID) []byte {
|
|||
return payload
|
||||
}
|
||||
|
||||
func newRegisterResponseSessionDatagram(id v3.RequestID, resp v3.SessionRegistrationResp) []byte {
|
||||
datagram := v3.UDPSessionRegistrationResponseDatagram{
|
||||
RequestID: id,
|
||||
ResponseType: resp,
|
||||
}
|
||||
payload, err := datagram.MarshalBinary()
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
return payload
|
||||
}
|
||||
|
||||
func newSessionPayloadDatagram(id v3.RequestID, payload []byte) []byte {
|
||||
datagram := make([]byte, len(payload)+17)
|
||||
err := v3.MarshalPayloadHeaderTo(id, datagram[:])
|
||||
|
|
|
@ -0,0 +1,77 @@
|
|||
package session
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"sync"
|
||||
)
|
||||
|
||||
const (
|
||||
unlimitedActiveSessions = 0
|
||||
)
|
||||
|
||||
var (
|
||||
ErrTooManyActiveSessions = errors.New("too many active sessions")
|
||||
)
|
||||
|
||||
type Limiter interface {
|
||||
// Acquire tries to acquire a free slot for a session, if the value of sessions is already above
|
||||
// the maximum it returns ErrTooManyActiveSessions.
|
||||
Acquire(sessionType string) error
|
||||
// Release releases a slot for a session.
|
||||
Release()
|
||||
// SetLimit allows to hot swap the limit value of the limiter.
|
||||
SetLimit(uint64)
|
||||
}
|
||||
|
||||
type sessionLimiter struct {
|
||||
limiterLock sync.Mutex
|
||||
activeSessionsCounter uint64
|
||||
maxActiveSessions uint64
|
||||
unlimited bool
|
||||
}
|
||||
|
||||
func NewLimiter(maxActiveSessions uint64) Limiter {
|
||||
sessionLimiter := &sessionLimiter{
|
||||
maxActiveSessions: maxActiveSessions,
|
||||
unlimited: isUnlimited(maxActiveSessions),
|
||||
}
|
||||
|
||||
return sessionLimiter
|
||||
}
|
||||
|
||||
func (s *sessionLimiter) Acquire(sessionType string) error {
|
||||
s.limiterLock.Lock()
|
||||
defer s.limiterLock.Unlock()
|
||||
|
||||
if !s.unlimited && s.activeSessionsCounter >= s.maxActiveSessions {
|
||||
sessionRegistrationsDropped.WithLabelValues(sessionType).Inc()
|
||||
return ErrTooManyActiveSessions
|
||||
}
|
||||
|
||||
s.activeSessionsCounter++
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *sessionLimiter) Release() {
|
||||
s.limiterLock.Lock()
|
||||
defer s.limiterLock.Unlock()
|
||||
|
||||
if s.activeSessionsCounter <= 0 {
|
||||
return
|
||||
}
|
||||
|
||||
s.activeSessionsCounter--
|
||||
}
|
||||
|
||||
func (s *sessionLimiter) SetLimit(newMaxActiveSessions uint64) {
|
||||
s.limiterLock.Lock()
|
||||
defer s.limiterLock.Unlock()
|
||||
|
||||
s.maxActiveSessions = newMaxActiveSessions
|
||||
s.unlimited = isUnlimited(newMaxActiveSessions)
|
||||
}
|
||||
|
||||
// isUnlimited checks if the value received matches the configuration for the unlimited session limiter.
|
||||
func isUnlimited(value uint64) bool {
|
||||
return value == unlimitedActiveSessions
|
||||
}
|
|
@ -0,0 +1,119 @@
|
|||
package session_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/cloudflare/cloudflared/session"
|
||||
)
|
||||
|
||||
func TestSessionLimiter_Unlimited(t *testing.T) {
|
||||
unlimitedLimiter := session.NewLimiter(0)
|
||||
|
||||
for i := 0; i < 1000; i++ {
|
||||
err := unlimitedLimiter.Acquire("test")
|
||||
require.NoError(t, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionLimiter_Limited(t *testing.T) {
|
||||
maxSessions := uint64(5)
|
||||
limiter := session.NewLimiter(maxSessions)
|
||||
|
||||
for i := uint64(0); i < maxSessions; i++ {
|
||||
err := limiter.Acquire("test")
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
err := limiter.Acquire("should fail")
|
||||
require.ErrorIs(t, err, session.ErrTooManyActiveSessions)
|
||||
}
|
||||
|
||||
func TestSessionLimiter_AcquireAndReleaseSession(t *testing.T) {
|
||||
maxSessions := uint64(5)
|
||||
limiter := session.NewLimiter(maxSessions)
|
||||
|
||||
// Acquire the maximum number of sessions
|
||||
for i := uint64(0); i < maxSessions; i++ {
|
||||
err := limiter.Acquire("test")
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// Validate acquire 1 more sessions fails
|
||||
err := limiter.Acquire("should fail")
|
||||
require.ErrorIs(t, err, session.ErrTooManyActiveSessions)
|
||||
|
||||
// Release the maximum number of sessions
|
||||
for i := uint64(0); i < maxSessions; i++ {
|
||||
limiter.Release()
|
||||
}
|
||||
|
||||
// Validate acquire 1 more sessions works
|
||||
err = limiter.Acquire("shouldn't fail")
|
||||
require.NoError(t, err)
|
||||
|
||||
// Release a 10x the number of max sessions
|
||||
for i := uint64(0); i < 10*maxSessions; i++ {
|
||||
limiter.Release()
|
||||
}
|
||||
|
||||
// Validate it still can only acquire a value = number max sessions.
|
||||
for i := uint64(0); i < maxSessions; i++ {
|
||||
err := limiter.Acquire("test")
|
||||
require.NoError(t, err)
|
||||
}
|
||||
err = limiter.Acquire("should fail")
|
||||
require.ErrorIs(t, err, session.ErrTooManyActiveSessions)
|
||||
}
|
||||
|
||||
func TestSessionLimiter_SetLimit(t *testing.T) {
|
||||
maxSessions := uint64(5)
|
||||
limiter := session.NewLimiter(maxSessions)
|
||||
|
||||
// Acquire the maximum number of sessions
|
||||
for i := uint64(0); i < maxSessions; i++ {
|
||||
err := limiter.Acquire("test")
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// Validate acquire 1 more sessions fails
|
||||
err := limiter.Acquire("should fail")
|
||||
require.ErrorIs(t, err, session.ErrTooManyActiveSessions)
|
||||
|
||||
// Set the session limiter to support one more request
|
||||
limiter.SetLimit(maxSessions + 1)
|
||||
|
||||
// Validate acquire 1 more sessions now works
|
||||
err = limiter.Acquire("shouldn't fail")
|
||||
require.NoError(t, err)
|
||||
|
||||
// Validate acquire 1 more sessions doesn't work because we already reached the limit
|
||||
err = limiter.Acquire("should fail")
|
||||
require.ErrorIs(t, err, session.ErrTooManyActiveSessions)
|
||||
|
||||
// Release all sessions
|
||||
for i := uint64(0); i < maxSessions+1; i++ {
|
||||
limiter.Release()
|
||||
}
|
||||
|
||||
// Validate 1 session works again
|
||||
err = limiter.Acquire("shouldn't fail")
|
||||
require.NoError(t, err)
|
||||
|
||||
// Set the session limit to 1
|
||||
limiter.SetLimit(1)
|
||||
|
||||
// Validate acquire 1 more sessions doesn't work
|
||||
err = limiter.Acquire("should fail")
|
||||
require.ErrorIs(t, err, session.ErrTooManyActiveSessions)
|
||||
|
||||
// Set the session limit to unlimited
|
||||
limiter.SetLimit(0)
|
||||
|
||||
// Validate it can acquire a lot of sessions because it is now unlimited.
|
||||
for i := uint64(0); i < 10*maxSessions; i++ {
|
||||
err := limiter.Acquire("shouldn't fail")
|
||||
require.NoError(t, err)
|
||||
}
|
||||
}
|
|
@ -0,0 +1,23 @@
|
|||
package session
|
||||
|
||||
import (
|
||||
"github.com/prometheus/client_golang/prometheus"
|
||||
"github.com/prometheus/client_golang/prometheus/promauto"
|
||||
)
|
||||
|
||||
const (
|
||||
namespace = "session"
|
||||
)
|
||||
|
||||
var (
|
||||
labels = []string{"session_type"}
|
||||
|
||||
sessionRegistrationsDropped = promauto.NewCounterVec(prometheus.CounterOpts{
|
||||
Namespace: namespace,
|
||||
Subsystem: "client",
|
||||
Name: "registrations_rate_limited_total",
|
||||
Help: "Count registrations dropped due to high number of concurrent sessions being handled",
|
||||
},
|
||||
labels,
|
||||
)
|
||||
)
|
|
@ -26,12 +26,6 @@ const (
|
|||
tunnelRetryDuration = time.Second * 10
|
||||
// Interval between registering new tunnels
|
||||
registrationInterval = time.Second
|
||||
|
||||
subsystemRefreshAuth = "refresh_auth"
|
||||
// Maximum exponent for 'Authenticate' exponential backoff
|
||||
refreshAuthMaxBackoff = 10
|
||||
// Waiting time before retrying a failed 'Authenticate' connection
|
||||
refreshAuthRetryDuration = time.Second * 10
|
||||
)
|
||||
|
||||
// Supervisor manages non-declarative tunnels. Establishes TCP connections with the edge, and
|
||||
|
@ -84,7 +78,7 @@ func NewSupervisor(config *TunnelConfig, orchestrator *orchestration.Orchestrato
|
|||
edgeBindAddr := config.EdgeBindAddr
|
||||
|
||||
datagramMetrics := v3.NewMetrics(prometheus.DefaultRegisterer)
|
||||
sessionManager := v3.NewSessionManager(datagramMetrics, config.Log, ingress.DialUDPAddrPort)
|
||||
sessionManager := v3.NewSessionManager(datagramMetrics, config.Log, ingress.DialUDPAddrPort, orchestrator.GetSessionLimiter())
|
||||
|
||||
edgeTunnelServer := EdgeTunnelServer{
|
||||
config: config,
|
||||
|
@ -313,6 +307,7 @@ func (s *Supervisor) startTunnel(
|
|||
s.tunnelErrors <- tunnelError{index: index, err: err}
|
||||
}()
|
||||
|
||||
// nolint: gosec
|
||||
err = s.edgeTunnelServer.Serve(ctx, uint8(index), s.tunnelsProtocolFallback[index], connectedSignal)
|
||||
}
|
||||
|
||||
|
@ -334,7 +329,3 @@ func (s *Supervisor) waitForNextTunnel(index int) bool {
|
|||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (s *Supervisor) unusedIPs() bool {
|
||||
return s.edgeIPs.AvailableAddrs() > s.config.HAConnections
|
||||
}
|
||||
|
|
|
@ -459,6 +459,7 @@ func (e *EdgeTunnelServer) serveConnection(
|
|||
|
||||
switch protocol {
|
||||
case connection.QUIC:
|
||||
// nolint: gosec
|
||||
connOptions := e.config.connectionOptions(addr.UDP.String(), uint8(backoff.Retries()))
|
||||
return e.serveQUIC(ctx,
|
||||
addr.UDP.AddrPort(),
|
||||
|
@ -474,6 +475,7 @@ func (e *EdgeTunnelServer) serveConnection(
|
|||
return err, true
|
||||
}
|
||||
|
||||
// nolint: gosec
|
||||
connOptions := e.config.connectionOptions(edgeConn.LocalAddr().String(), uint8(backoff.Retries()))
|
||||
if err := e.serveHTTP2(
|
||||
ctx,
|
||||
|
@ -615,6 +617,7 @@ func (e *EdgeTunnelServer) serveQUIC(
|
|||
connIndex,
|
||||
e.config.RPCTimeout,
|
||||
e.config.WriteStreamTimeout,
|
||||
e.orchestrator.GetSessionLimiter(),
|
||||
connLogger.Logger(),
|
||||
)
|
||||
}
|
||||
|
|
|
@ -0,0 +1,506 @@
|
|||
// Copyright 2010 Google Inc.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package gomock
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"reflect"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Call represents an expected call to a mock.
|
||||
type Call struct {
|
||||
t TestHelper // for triggering test failures on invalid call setup
|
||||
|
||||
receiver any // the receiver of the method call
|
||||
method string // the name of the method
|
||||
methodType reflect.Type // the type of the method
|
||||
args []Matcher // the args
|
||||
origin string // file and line number of call setup
|
||||
|
||||
preReqs []*Call // prerequisite calls
|
||||
|
||||
// Expectations
|
||||
minCalls, maxCalls int
|
||||
|
||||
numCalls int // actual number made
|
||||
|
||||
// actions are called when this Call is called. Each action gets the args and
|
||||
// can set the return values by returning a non-nil slice. Actions run in the
|
||||
// order they are created.
|
||||
actions []func([]any) []any
|
||||
}
|
||||
|
||||
// newCall creates a *Call. It requires the method type in order to support
|
||||
// unexported methods.
|
||||
func newCall(t TestHelper, receiver any, method string, methodType reflect.Type, args ...any) *Call {
|
||||
t.Helper()
|
||||
|
||||
// TODO: check arity, types.
|
||||
mArgs := make([]Matcher, len(args))
|
||||
for i, arg := range args {
|
||||
if m, ok := arg.(Matcher); ok {
|
||||
mArgs[i] = m
|
||||
} else if arg == nil {
|
||||
// Handle nil specially so that passing a nil interface value
|
||||
// will match the typed nils of concrete args.
|
||||
mArgs[i] = Nil()
|
||||
} else {
|
||||
mArgs[i] = Eq(arg)
|
||||
}
|
||||
}
|
||||
|
||||
// callerInfo's skip should be updated if the number of calls between the user's test
|
||||
// and this line changes, i.e. this code is wrapped in another anonymous function.
|
||||
// 0 is us, 1 is RecordCallWithMethodType(), 2 is the generated recorder, and 3 is the user's test.
|
||||
origin := callerInfo(3)
|
||||
actions := []func([]any) []any{func([]any) []any {
|
||||
// Synthesize the zero value for each of the return args' types.
|
||||
rets := make([]any, methodType.NumOut())
|
||||
for i := 0; i < methodType.NumOut(); i++ {
|
||||
rets[i] = reflect.Zero(methodType.Out(i)).Interface()
|
||||
}
|
||||
return rets
|
||||
}}
|
||||
return &Call{
|
||||
t: t, receiver: receiver, method: method, methodType: methodType,
|
||||
args: mArgs, origin: origin, minCalls: 1, maxCalls: 1, actions: actions,
|
||||
}
|
||||
}
|
||||
|
||||
// AnyTimes allows the expectation to be called 0 or more times
|
||||
func (c *Call) AnyTimes() *Call {
|
||||
c.minCalls, c.maxCalls = 0, 1e8 // close enough to infinity
|
||||
return c
|
||||
}
|
||||
|
||||
// MinTimes requires the call to occur at least n times. If AnyTimes or MaxTimes have not been called or if MaxTimes
|
||||
// was previously called with 1, MinTimes also sets the maximum number of calls to infinity.
|
||||
func (c *Call) MinTimes(n int) *Call {
|
||||
c.minCalls = n
|
||||
if c.maxCalls == 1 {
|
||||
c.maxCalls = 1e8
|
||||
}
|
||||
return c
|
||||
}
|
||||
|
||||
// MaxTimes limits the number of calls to n times. If AnyTimes or MinTimes have not been called or if MinTimes was
|
||||
// previously called with 1, MaxTimes also sets the minimum number of calls to 0.
|
||||
func (c *Call) MaxTimes(n int) *Call {
|
||||
c.maxCalls = n
|
||||
if c.minCalls == 1 {
|
||||
c.minCalls = 0
|
||||
}
|
||||
return c
|
||||
}
|
||||
|
||||
// DoAndReturn declares the action to run when the call is matched.
|
||||
// The return values from this function are returned by the mocked function.
|
||||
// It takes an any argument to support n-arity functions.
|
||||
// The anonymous function must match the function signature mocked method.
|
||||
func (c *Call) DoAndReturn(f any) *Call {
|
||||
// TODO: Check arity and types here, rather than dying badly elsewhere.
|
||||
v := reflect.ValueOf(f)
|
||||
|
||||
c.addAction(func(args []any) []any {
|
||||
c.t.Helper()
|
||||
ft := v.Type()
|
||||
if c.methodType.NumIn() != ft.NumIn() {
|
||||
if ft.IsVariadic() {
|
||||
c.t.Fatalf("wrong number of arguments in DoAndReturn func for %T.%v The function signature must match the mocked method, a variadic function cannot be used.",
|
||||
c.receiver, c.method)
|
||||
} else {
|
||||
c.t.Fatalf("wrong number of arguments in DoAndReturn func for %T.%v: got %d, want %d [%s]",
|
||||
c.receiver, c.method, ft.NumIn(), c.methodType.NumIn(), c.origin)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
vArgs := make([]reflect.Value, len(args))
|
||||
for i := 0; i < len(args); i++ {
|
||||
if args[i] != nil {
|
||||
vArgs[i] = reflect.ValueOf(args[i])
|
||||
} else {
|
||||
// Use the zero value for the arg.
|
||||
vArgs[i] = reflect.Zero(ft.In(i))
|
||||
}
|
||||
}
|
||||
vRets := v.Call(vArgs)
|
||||
rets := make([]any, len(vRets))
|
||||
for i, ret := range vRets {
|
||||
rets[i] = ret.Interface()
|
||||
}
|
||||
return rets
|
||||
})
|
||||
return c
|
||||
}
|
||||
|
||||
// Do declares the action to run when the call is matched. The function's
|
||||
// return values are ignored to retain backward compatibility. To use the
|
||||
// return values call DoAndReturn.
|
||||
// It takes an any argument to support n-arity functions.
|
||||
// The anonymous function must match the function signature mocked method.
|
||||
func (c *Call) Do(f any) *Call {
|
||||
// TODO: Check arity and types here, rather than dying badly elsewhere.
|
||||
v := reflect.ValueOf(f)
|
||||
|
||||
c.addAction(func(args []any) []any {
|
||||
c.t.Helper()
|
||||
ft := v.Type()
|
||||
if c.methodType.NumIn() != ft.NumIn() {
|
||||
if ft.IsVariadic() {
|
||||
c.t.Fatalf("wrong number of arguments in Do func for %T.%v The function signature must match the mocked method, a variadic function cannot be used.",
|
||||
c.receiver, c.method)
|
||||
} else {
|
||||
c.t.Fatalf("wrong number of arguments in Do func for %T.%v: got %d, want %d [%s]",
|
||||
c.receiver, c.method, ft.NumIn(), c.methodType.NumIn(), c.origin)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
vArgs := make([]reflect.Value, len(args))
|
||||
for i := 0; i < len(args); i++ {
|
||||
if args[i] != nil {
|
||||
vArgs[i] = reflect.ValueOf(args[i])
|
||||
} else {
|
||||
// Use the zero value for the arg.
|
||||
vArgs[i] = reflect.Zero(ft.In(i))
|
||||
}
|
||||
}
|
||||
v.Call(vArgs)
|
||||
return nil
|
||||
})
|
||||
return c
|
||||
}
|
||||
|
||||
// Return declares the values to be returned by the mocked function call.
|
||||
func (c *Call) Return(rets ...any) *Call {
|
||||
c.t.Helper()
|
||||
|
||||
mt := c.methodType
|
||||
if len(rets) != mt.NumOut() {
|
||||
c.t.Fatalf("wrong number of arguments to Return for %T.%v: got %d, want %d [%s]",
|
||||
c.receiver, c.method, len(rets), mt.NumOut(), c.origin)
|
||||
}
|
||||
for i, ret := range rets {
|
||||
if got, want := reflect.TypeOf(ret), mt.Out(i); got == want {
|
||||
// Identical types; nothing to do.
|
||||
} else if got == nil {
|
||||
// Nil needs special handling.
|
||||
switch want.Kind() {
|
||||
case reflect.Chan, reflect.Func, reflect.Interface, reflect.Map, reflect.Ptr, reflect.Slice:
|
||||
// ok
|
||||
default:
|
||||
c.t.Fatalf("argument %d to Return for %T.%v is nil, but %v is not nillable [%s]",
|
||||
i, c.receiver, c.method, want, c.origin)
|
||||
}
|
||||
} else if got.AssignableTo(want) {
|
||||
// Assignable type relation. Make the assignment now so that the generated code
|
||||
// can return the values with a type assertion.
|
||||
v := reflect.New(want).Elem()
|
||||
v.Set(reflect.ValueOf(ret))
|
||||
rets[i] = v.Interface()
|
||||
} else {
|
||||
c.t.Fatalf("wrong type of argument %d to Return for %T.%v: %v is not assignable to %v [%s]",
|
||||
i, c.receiver, c.method, got, want, c.origin)
|
||||
}
|
||||
}
|
||||
|
||||
c.addAction(func([]any) []any {
|
||||
return rets
|
||||
})
|
||||
|
||||
return c
|
||||
}
|
||||
|
||||
// Times declares the exact number of times a function call is expected to be executed.
|
||||
func (c *Call) Times(n int) *Call {
|
||||
c.minCalls, c.maxCalls = n, n
|
||||
return c
|
||||
}
|
||||
|
||||
// SetArg declares an action that will set the nth argument's value,
|
||||
// indirected through a pointer. Or, in the case of a slice and map, SetArg
|
||||
// will copy value's elements/key-value pairs into the nth argument.
|
||||
func (c *Call) SetArg(n int, value any) *Call {
|
||||
c.t.Helper()
|
||||
|
||||
mt := c.methodType
|
||||
// TODO: This will break on variadic methods.
|
||||
// We will need to check those at invocation time.
|
||||
if n < 0 || n >= mt.NumIn() {
|
||||
c.t.Fatalf("SetArg(%d, ...) called for a method with %d args [%s]",
|
||||
n, mt.NumIn(), c.origin)
|
||||
}
|
||||
// Permit setting argument through an interface.
|
||||
// In the interface case, we don't (nay, can't) check the type here.
|
||||
at := mt.In(n)
|
||||
switch at.Kind() {
|
||||
case reflect.Ptr:
|
||||
dt := at.Elem()
|
||||
if vt := reflect.TypeOf(value); !vt.AssignableTo(dt) {
|
||||
c.t.Fatalf("SetArg(%d, ...) argument is a %v, not assignable to %v [%s]",
|
||||
n, vt, dt, c.origin)
|
||||
}
|
||||
case reflect.Interface, reflect.Slice, reflect.Map:
|
||||
// nothing to do
|
||||
default:
|
||||
c.t.Fatalf("SetArg(%d, ...) referring to argument of non-pointer non-interface non-slice non-map type %v [%s]",
|
||||
n, at, c.origin)
|
||||
}
|
||||
|
||||
c.addAction(func(args []any) []any {
|
||||
v := reflect.ValueOf(value)
|
||||
switch reflect.TypeOf(args[n]).Kind() {
|
||||
case reflect.Slice:
|
||||
setSlice(args[n], v)
|
||||
case reflect.Map:
|
||||
setMap(args[n], v)
|
||||
default:
|
||||
reflect.ValueOf(args[n]).Elem().Set(v)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
return c
|
||||
}
|
||||
|
||||
// isPreReq returns true if other is a direct or indirect prerequisite to c.
|
||||
func (c *Call) isPreReq(other *Call) bool {
|
||||
for _, preReq := range c.preReqs {
|
||||
if other == preReq || preReq.isPreReq(other) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// After declares that the call may only match after preReq has been exhausted.
|
||||
func (c *Call) After(preReq *Call) *Call {
|
||||
c.t.Helper()
|
||||
|
||||
if c == preReq {
|
||||
c.t.Fatalf("A call isn't allowed to be its own prerequisite")
|
||||
}
|
||||
if preReq.isPreReq(c) {
|
||||
c.t.Fatalf("Loop in call order: %v is a prerequisite to %v (possibly indirectly).", c, preReq)
|
||||
}
|
||||
|
||||
c.preReqs = append(c.preReqs, preReq)
|
||||
return c
|
||||
}
|
||||
|
||||
// Returns true if the minimum number of calls have been made.
|
||||
func (c *Call) satisfied() bool {
|
||||
return c.numCalls >= c.minCalls
|
||||
}
|
||||
|
||||
// Returns true if the maximum number of calls have been made.
|
||||
func (c *Call) exhausted() bool {
|
||||
return c.numCalls >= c.maxCalls
|
||||
}
|
||||
|
||||
func (c *Call) String() string {
|
||||
args := make([]string, len(c.args))
|
||||
for i, arg := range c.args {
|
||||
args[i] = arg.String()
|
||||
}
|
||||
arguments := strings.Join(args, ", ")
|
||||
return fmt.Sprintf("%T.%v(%s) %s", c.receiver, c.method, arguments, c.origin)
|
||||
}
|
||||
|
||||
// Tests if the given call matches the expected call.
|
||||
// If yes, returns nil. If no, returns error with message explaining why it does not match.
|
||||
func (c *Call) matches(args []any) error {
|
||||
if !c.methodType.IsVariadic() {
|
||||
if len(args) != len(c.args) {
|
||||
return fmt.Errorf("expected call at %s has the wrong number of arguments. Got: %d, want: %d",
|
||||
c.origin, len(args), len(c.args))
|
||||
}
|
||||
|
||||
for i, m := range c.args {
|
||||
if !m.Matches(args[i]) {
|
||||
return fmt.Errorf(
|
||||
"expected call at %s doesn't match the argument at index %d.\nGot: %v\nWant: %v",
|
||||
c.origin, i, formatGottenArg(m, args[i]), m,
|
||||
)
|
||||
}
|
||||
}
|
||||
} else {
|
||||
if len(c.args) < c.methodType.NumIn()-1 {
|
||||
return fmt.Errorf("expected call at %s has the wrong number of matchers. Got: %d, want: %d",
|
||||
c.origin, len(c.args), c.methodType.NumIn()-1)
|
||||
}
|
||||
if len(c.args) != c.methodType.NumIn() && len(args) != len(c.args) {
|
||||
return fmt.Errorf("expected call at %s has the wrong number of arguments. Got: %d, want: %d",
|
||||
c.origin, len(args), len(c.args))
|
||||
}
|
||||
if len(args) < len(c.args)-1 {
|
||||
return fmt.Errorf("expected call at %s has the wrong number of arguments. Got: %d, want: greater than or equal to %d",
|
||||
c.origin, len(args), len(c.args)-1)
|
||||
}
|
||||
|
||||
for i, m := range c.args {
|
||||
if i < c.methodType.NumIn()-1 {
|
||||
// Non-variadic args
|
||||
if !m.Matches(args[i]) {
|
||||
return fmt.Errorf("expected call at %s doesn't match the argument at index %s.\nGot: %v\nWant: %v",
|
||||
c.origin, strconv.Itoa(i), formatGottenArg(m, args[i]), m)
|
||||
}
|
||||
continue
|
||||
}
|
||||
// The last arg has a possibility of a variadic argument, so let it branch
|
||||
|
||||
// sample: Foo(a int, b int, c ...int)
|
||||
if i < len(c.args) && i < len(args) {
|
||||
if m.Matches(args[i]) {
|
||||
// Got Foo(a, b, c) want Foo(matcherA, matcherB, gomock.Any())
|
||||
// Got Foo(a, b, c) want Foo(matcherA, matcherB, someSliceMatcher)
|
||||
// Got Foo(a, b, c) want Foo(matcherA, matcherB, matcherC)
|
||||
// Got Foo(a, b) want Foo(matcherA, matcherB)
|
||||
// Got Foo(a, b, c, d) want Foo(matcherA, matcherB, matcherC, matcherD)
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
// The number of actual args don't match the number of matchers,
|
||||
// or the last matcher is a slice and the last arg is not.
|
||||
// If this function still matches it is because the last matcher
|
||||
// matches all the remaining arguments or the lack of any.
|
||||
// Convert the remaining arguments, if any, into a slice of the
|
||||
// expected type.
|
||||
vArgsType := c.methodType.In(c.methodType.NumIn() - 1)
|
||||
vArgs := reflect.MakeSlice(vArgsType, 0, len(args)-i)
|
||||
for _, arg := range args[i:] {
|
||||
vArgs = reflect.Append(vArgs, reflect.ValueOf(arg))
|
||||
}
|
||||
if m.Matches(vArgs.Interface()) {
|
||||
// Got Foo(a, b, c, d, e) want Foo(matcherA, matcherB, gomock.Any())
|
||||
// Got Foo(a, b, c, d, e) want Foo(matcherA, matcherB, someSliceMatcher)
|
||||
// Got Foo(a, b) want Foo(matcherA, matcherB, gomock.Any())
|
||||
// Got Foo(a, b) want Foo(matcherA, matcherB, someEmptySliceMatcher)
|
||||
break
|
||||
}
|
||||
// Wrong number of matchers or not match. Fail.
|
||||
// Got Foo(a, b) want Foo(matcherA, matcherB, matcherC, matcherD)
|
||||
// Got Foo(a, b, c) want Foo(matcherA, matcherB, matcherC, matcherD)
|
||||
// Got Foo(a, b, c, d) want Foo(matcherA, matcherB, matcherC, matcherD, matcherE)
|
||||
// Got Foo(a, b, c, d, e) want Foo(matcherA, matcherB, matcherC, matcherD)
|
||||
// Got Foo(a, b, c) want Foo(matcherA, matcherB)
|
||||
|
||||
return fmt.Errorf("expected call at %s doesn't match the argument at index %s.\nGot: %v\nWant: %v",
|
||||
c.origin, strconv.Itoa(i), formatGottenArg(m, args[i:]), c.args[i])
|
||||
}
|
||||
}
|
||||
|
||||
// Check that all prerequisite calls have been satisfied.
|
||||
for _, preReqCall := range c.preReqs {
|
||||
if !preReqCall.satisfied() {
|
||||
return fmt.Errorf("expected call at %s doesn't have a prerequisite call satisfied:\n%v\nshould be called before:\n%v",
|
||||
c.origin, preReqCall, c)
|
||||
}
|
||||
}
|
||||
|
||||
// Check that the call is not exhausted.
|
||||
if c.exhausted() {
|
||||
return fmt.Errorf("expected call at %s has already been called the max number of times", c.origin)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// dropPrereqs tells the expected Call to not re-check prerequisite calls any
|
||||
// longer, and to return its current set.
|
||||
func (c *Call) dropPrereqs() (preReqs []*Call) {
|
||||
preReqs = c.preReqs
|
||||
c.preReqs = nil
|
||||
return
|
||||
}
|
||||
|
||||
func (c *Call) call() []func([]any) []any {
|
||||
c.numCalls++
|
||||
return c.actions
|
||||
}
|
||||
|
||||
// InOrder declares that the given calls should occur in order.
|
||||
// It panics if the type of any of the arguments isn't *Call or a generated
|
||||
// mock with an embedded *Call.
|
||||
func InOrder(args ...any) {
|
||||
calls := make([]*Call, 0, len(args))
|
||||
for i := 0; i < len(args); i++ {
|
||||
if call := getCall(args[i]); call != nil {
|
||||
calls = append(calls, call)
|
||||
continue
|
||||
}
|
||||
panic(fmt.Sprintf(
|
||||
"invalid argument at position %d of type %T, InOrder expects *gomock.Call or generated mock types with an embedded *gomock.Call",
|
||||
i,
|
||||
args[i],
|
||||
))
|
||||
}
|
||||
for i := 1; i < len(calls); i++ {
|
||||
calls[i].After(calls[i-1])
|
||||
}
|
||||
}
|
||||
|
||||
// getCall checks if the parameter is a *Call or a generated struct
|
||||
// that wraps a *Call and returns the *Call pointer - if neither, it returns nil.
|
||||
func getCall(arg any) *Call {
|
||||
if call, ok := arg.(*Call); ok {
|
||||
return call
|
||||
}
|
||||
t := reflect.ValueOf(arg)
|
||||
if t.Kind() != reflect.Ptr && t.Kind() != reflect.Interface {
|
||||
return nil
|
||||
}
|
||||
t = t.Elem()
|
||||
for i := 0; i < t.NumField(); i++ {
|
||||
f := t.Field(i)
|
||||
if !f.CanInterface() {
|
||||
continue
|
||||
}
|
||||
if call, ok := f.Interface().(*Call); ok {
|
||||
return call
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func setSlice(arg any, v reflect.Value) {
|
||||
va := reflect.ValueOf(arg)
|
||||
for i := 0; i < v.Len(); i++ {
|
||||
va.Index(i).Set(v.Index(i))
|
||||
}
|
||||
}
|
||||
|
||||
func setMap(arg any, v reflect.Value) {
|
||||
va := reflect.ValueOf(arg)
|
||||
for _, e := range va.MapKeys() {
|
||||
va.SetMapIndex(e, reflect.Value{})
|
||||
}
|
||||
for _, e := range v.MapKeys() {
|
||||
va.SetMapIndex(e, v.MapIndex(e))
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Call) addAction(action func([]any) []any) {
|
||||
c.actions = append(c.actions, action)
|
||||
}
|
||||
|
||||
func formatGottenArg(m Matcher, arg any) string {
|
||||
got := fmt.Sprintf("%v (%T)", arg, arg)
|
||||
if gs, ok := m.(GotFormatter); ok {
|
||||
got = gs.Got(arg)
|
||||
}
|
||||
return got
|
||||
}
|
|
@ -0,0 +1,164 @@
|
|||
// Copyright 2011 Google Inc.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package gomock
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// callSet represents a set of expected calls, indexed by receiver and method
|
||||
// name.
|
||||
type callSet struct {
|
||||
// Calls that are still expected.
|
||||
expected map[callSetKey][]*Call
|
||||
expectedMu *sync.Mutex
|
||||
// Calls that have been exhausted.
|
||||
exhausted map[callSetKey][]*Call
|
||||
// when set to true, existing call expectations are overridden when new call expectations are made
|
||||
allowOverride bool
|
||||
}
|
||||
|
||||
// callSetKey is the key in the maps in callSet
|
||||
type callSetKey struct {
|
||||
receiver any
|
||||
fname string
|
||||
}
|
||||
|
||||
func newCallSet() *callSet {
|
||||
return &callSet{
|
||||
expected: make(map[callSetKey][]*Call),
|
||||
expectedMu: &sync.Mutex{},
|
||||
exhausted: make(map[callSetKey][]*Call),
|
||||
}
|
||||
}
|
||||
|
||||
func newOverridableCallSet() *callSet {
|
||||
return &callSet{
|
||||
expected: make(map[callSetKey][]*Call),
|
||||
expectedMu: &sync.Mutex{},
|
||||
exhausted: make(map[callSetKey][]*Call),
|
||||
allowOverride: true,
|
||||
}
|
||||
}
|
||||
|
||||
// Add adds a new expected call.
|
||||
func (cs callSet) Add(call *Call) {
|
||||
key := callSetKey{call.receiver, call.method}
|
||||
|
||||
cs.expectedMu.Lock()
|
||||
defer cs.expectedMu.Unlock()
|
||||
|
||||
m := cs.expected
|
||||
if call.exhausted() {
|
||||
m = cs.exhausted
|
||||
}
|
||||
if cs.allowOverride {
|
||||
m[key] = make([]*Call, 0)
|
||||
}
|
||||
|
||||
m[key] = append(m[key], call)
|
||||
}
|
||||
|
||||
// Remove removes an expected call.
|
||||
func (cs callSet) Remove(call *Call) {
|
||||
key := callSetKey{call.receiver, call.method}
|
||||
|
||||
cs.expectedMu.Lock()
|
||||
defer cs.expectedMu.Unlock()
|
||||
|
||||
calls := cs.expected[key]
|
||||
for i, c := range calls {
|
||||
if c == call {
|
||||
// maintain order for remaining calls
|
||||
cs.expected[key] = append(calls[:i], calls[i+1:]...)
|
||||
cs.exhausted[key] = append(cs.exhausted[key], call)
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// FindMatch searches for a matching call. Returns error with explanation message if no call matched.
|
||||
func (cs callSet) FindMatch(receiver any, method string, args []any) (*Call, error) {
|
||||
key := callSetKey{receiver, method}
|
||||
|
||||
cs.expectedMu.Lock()
|
||||
defer cs.expectedMu.Unlock()
|
||||
|
||||
// Search through the expected calls.
|
||||
expected := cs.expected[key]
|
||||
var callsErrors bytes.Buffer
|
||||
for _, call := range expected {
|
||||
err := call.matches(args)
|
||||
if err != nil {
|
||||
_, _ = fmt.Fprintf(&callsErrors, "\n%v", err)
|
||||
} else {
|
||||
return call, nil
|
||||
}
|
||||
}
|
||||
|
||||
// If we haven't found a match then search through the exhausted calls so we
|
||||
// get useful error messages.
|
||||
exhausted := cs.exhausted[key]
|
||||
for _, call := range exhausted {
|
||||
if err := call.matches(args); err != nil {
|
||||
_, _ = fmt.Fprintf(&callsErrors, "\n%v", err)
|
||||
continue
|
||||
}
|
||||
_, _ = fmt.Fprintf(
|
||||
&callsErrors, "all expected calls for method %q have been exhausted", method,
|
||||
)
|
||||
}
|
||||
|
||||
if len(expected)+len(exhausted) == 0 {
|
||||
_, _ = fmt.Fprintf(&callsErrors, "there are no expected calls of the method %q for that receiver", method)
|
||||
}
|
||||
|
||||
return nil, errors.New(callsErrors.String())
|
||||
}
|
||||
|
||||
// Failures returns the calls that are not satisfied.
|
||||
func (cs callSet) Failures() []*Call {
|
||||
cs.expectedMu.Lock()
|
||||
defer cs.expectedMu.Unlock()
|
||||
|
||||
failures := make([]*Call, 0, len(cs.expected))
|
||||
for _, calls := range cs.expected {
|
||||
for _, call := range calls {
|
||||
if !call.satisfied() {
|
||||
failures = append(failures, call)
|
||||
}
|
||||
}
|
||||
}
|
||||
return failures
|
||||
}
|
||||
|
||||
// Satisfied returns true in case all expected calls in this callSet are satisfied.
|
||||
func (cs callSet) Satisfied() bool {
|
||||
cs.expectedMu.Lock()
|
||||
defer cs.expectedMu.Unlock()
|
||||
|
||||
for _, calls := range cs.expected {
|
||||
for _, call := range calls {
|
||||
if !call.satisfied() {
|
||||
return false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
|
@ -0,0 +1,326 @@
|
|||
// Copyright 2010 Google Inc.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package gomock
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"reflect"
|
||||
"runtime"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// A TestReporter is something that can be used to report test failures. It
|
||||
// is satisfied by the standard library's *testing.T.
|
||||
type TestReporter interface {
|
||||
Errorf(format string, args ...any)
|
||||
Fatalf(format string, args ...any)
|
||||
}
|
||||
|
||||
// TestHelper is a TestReporter that has the Helper method. It is satisfied
|
||||
// by the standard library's *testing.T.
|
||||
type TestHelper interface {
|
||||
TestReporter
|
||||
Helper()
|
||||
}
|
||||
|
||||
// cleanuper is used to check if TestHelper also has the `Cleanup` method. A
|
||||
// common pattern is to pass in a `*testing.T` to
|
||||
// `NewController(t TestReporter)`. In Go 1.14+, `*testing.T` has a cleanup
|
||||
// method. This can be utilized to call `Finish()` so the caller of this library
|
||||
// does not have to.
|
||||
type cleanuper interface {
|
||||
Cleanup(func())
|
||||
}
|
||||
|
||||
// A Controller represents the top-level control of a mock ecosystem. It
|
||||
// defines the scope and lifetime of mock objects, as well as their
|
||||
// expectations. It is safe to call Controller's methods from multiple
|
||||
// goroutines. Each test should create a new Controller.
|
||||
//
|
||||
// func TestFoo(t *testing.T) {
|
||||
// ctrl := gomock.NewController(t)
|
||||
// // ..
|
||||
// }
|
||||
//
|
||||
// func TestBar(t *testing.T) {
|
||||
// t.Run("Sub-Test-1", st) {
|
||||
// ctrl := gomock.NewController(st)
|
||||
// // ..
|
||||
// })
|
||||
// t.Run("Sub-Test-2", st) {
|
||||
// ctrl := gomock.NewController(st)
|
||||
// // ..
|
||||
// })
|
||||
// })
|
||||
type Controller struct {
|
||||
// T should only be called within a generated mock. It is not intended to
|
||||
// be used in user code and may be changed in future versions. T is the
|
||||
// TestReporter passed in when creating the Controller via NewController.
|
||||
// If the TestReporter does not implement a TestHelper it will be wrapped
|
||||
// with a nopTestHelper.
|
||||
T TestHelper
|
||||
mu sync.Mutex
|
||||
expectedCalls *callSet
|
||||
finished bool
|
||||
}
|
||||
|
||||
// NewController returns a new Controller. It is the preferred way to create a Controller.
|
||||
//
|
||||
// Passing [*testing.T] registers cleanup function to automatically call [Controller.Finish]
|
||||
// when the test and all its subtests complete.
|
||||
func NewController(t TestReporter, opts ...ControllerOption) *Controller {
|
||||
h, ok := t.(TestHelper)
|
||||
if !ok {
|
||||
h = &nopTestHelper{t}
|
||||
}
|
||||
ctrl := &Controller{
|
||||
T: h,
|
||||
expectedCalls: newCallSet(),
|
||||
}
|
||||
for _, opt := range opts {
|
||||
opt.apply(ctrl)
|
||||
}
|
||||
if c, ok := isCleanuper(ctrl.T); ok {
|
||||
c.Cleanup(func() {
|
||||
ctrl.T.Helper()
|
||||
ctrl.finish(true, nil)
|
||||
})
|
||||
}
|
||||
|
||||
return ctrl
|
||||
}
|
||||
|
||||
// ControllerOption configures how a Controller should behave.
|
||||
type ControllerOption interface {
|
||||
apply(*Controller)
|
||||
}
|
||||
|
||||
type overridableExpectationsOption struct{}
|
||||
|
||||
// WithOverridableExpectations allows for overridable call expectations
|
||||
// i.e., subsequent call expectations override existing call expectations
|
||||
func WithOverridableExpectations() overridableExpectationsOption {
|
||||
return overridableExpectationsOption{}
|
||||
}
|
||||
|
||||
func (o overridableExpectationsOption) apply(ctrl *Controller) {
|
||||
ctrl.expectedCalls = newOverridableCallSet()
|
||||
}
|
||||
|
||||
type cancelReporter struct {
|
||||
t TestHelper
|
||||
cancel func()
|
||||
}
|
||||
|
||||
func (r *cancelReporter) Errorf(format string, args ...any) {
|
||||
r.t.Errorf(format, args...)
|
||||
}
|
||||
|
||||
func (r *cancelReporter) Fatalf(format string, args ...any) {
|
||||
defer r.cancel()
|
||||
r.t.Fatalf(format, args...)
|
||||
}
|
||||
|
||||
func (r *cancelReporter) Helper() {
|
||||
r.t.Helper()
|
||||
}
|
||||
|
||||
// WithContext returns a new Controller and a Context, which is cancelled on any
|
||||
// fatal failure.
|
||||
func WithContext(ctx context.Context, t TestReporter) (*Controller, context.Context) {
|
||||
h, ok := t.(TestHelper)
|
||||
if !ok {
|
||||
h = &nopTestHelper{t: t}
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(ctx)
|
||||
return NewController(&cancelReporter{t: h, cancel: cancel}), ctx
|
||||
}
|
||||
|
||||
type nopTestHelper struct {
|
||||
t TestReporter
|
||||
}
|
||||
|
||||
func (h *nopTestHelper) Errorf(format string, args ...any) {
|
||||
h.t.Errorf(format, args...)
|
||||
}
|
||||
|
||||
func (h *nopTestHelper) Fatalf(format string, args ...any) {
|
||||
h.t.Fatalf(format, args...)
|
||||
}
|
||||
|
||||
func (h nopTestHelper) Helper() {}
|
||||
|
||||
// RecordCall is called by a mock. It should not be called by user code.
|
||||
func (ctrl *Controller) RecordCall(receiver any, method string, args ...any) *Call {
|
||||
ctrl.T.Helper()
|
||||
|
||||
recv := reflect.ValueOf(receiver)
|
||||
for i := 0; i < recv.Type().NumMethod(); i++ {
|
||||
if recv.Type().Method(i).Name == method {
|
||||
return ctrl.RecordCallWithMethodType(receiver, method, recv.Method(i).Type(), args...)
|
||||
}
|
||||
}
|
||||
ctrl.T.Fatalf("gomock: failed finding method %s on %T", method, receiver)
|
||||
panic("unreachable")
|
||||
}
|
||||
|
||||
// RecordCallWithMethodType is called by a mock. It should not be called by user code.
|
||||
func (ctrl *Controller) RecordCallWithMethodType(receiver any, method string, methodType reflect.Type, args ...any) *Call {
|
||||
ctrl.T.Helper()
|
||||
|
||||
call := newCall(ctrl.T, receiver, method, methodType, args...)
|
||||
|
||||
ctrl.mu.Lock()
|
||||
defer ctrl.mu.Unlock()
|
||||
ctrl.expectedCalls.Add(call)
|
||||
|
||||
return call
|
||||
}
|
||||
|
||||
// Call is called by a mock. It should not be called by user code.
|
||||
func (ctrl *Controller) Call(receiver any, method string, args ...any) []any {
|
||||
ctrl.T.Helper()
|
||||
|
||||
// Nest this code so we can use defer to make sure the lock is released.
|
||||
actions := func() []func([]any) []any {
|
||||
ctrl.T.Helper()
|
||||
ctrl.mu.Lock()
|
||||
defer ctrl.mu.Unlock()
|
||||
|
||||
expected, err := ctrl.expectedCalls.FindMatch(receiver, method, args)
|
||||
if err != nil {
|
||||
// callerInfo's skip should be updated if the number of calls between the user's test
|
||||
// and this line changes, i.e. this code is wrapped in another anonymous function.
|
||||
// 0 is us, 1 is controller.Call(), 2 is the generated mock, and 3 is the user's test.
|
||||
origin := callerInfo(3)
|
||||
stringArgs := make([]string, len(args))
|
||||
for i, arg := range args {
|
||||
stringArgs[i] = getString(arg)
|
||||
}
|
||||
ctrl.T.Fatalf("Unexpected call to %T.%v(%v) at %s because: %s", receiver, method, stringArgs, origin, err)
|
||||
}
|
||||
|
||||
// Two things happen here:
|
||||
// * the matching call no longer needs to check prerequisite calls,
|
||||
// * and the prerequisite calls are no longer expected, so remove them.
|
||||
preReqCalls := expected.dropPrereqs()
|
||||
for _, preReqCall := range preReqCalls {
|
||||
ctrl.expectedCalls.Remove(preReqCall)
|
||||
}
|
||||
|
||||
actions := expected.call()
|
||||
if expected.exhausted() {
|
||||
ctrl.expectedCalls.Remove(expected)
|
||||
}
|
||||
return actions
|
||||
}()
|
||||
|
||||
var rets []any
|
||||
for _, action := range actions {
|
||||
if r := action(args); r != nil {
|
||||
rets = r
|
||||
}
|
||||
}
|
||||
|
||||
return rets
|
||||
}
|
||||
|
||||
// Finish checks to see if all the methods that were expected to be called were called.
|
||||
// It is not idempotent and therefore can only be invoked once.
|
||||
//
|
||||
// Note: If you pass a *testing.T into [NewController], you no longer
|
||||
// need to call ctrl.Finish() in your test methods.
|
||||
func (ctrl *Controller) Finish() {
|
||||
// If we're currently panicking, probably because this is a deferred call.
|
||||
// This must be recovered in the deferred function.
|
||||
err := recover()
|
||||
ctrl.finish(false, err)
|
||||
}
|
||||
|
||||
// Satisfied returns whether all expected calls bound to this Controller have been satisfied.
|
||||
// Calling Finish is then guaranteed to not fail due to missing calls.
|
||||
func (ctrl *Controller) Satisfied() bool {
|
||||
ctrl.mu.Lock()
|
||||
defer ctrl.mu.Unlock()
|
||||
return ctrl.expectedCalls.Satisfied()
|
||||
}
|
||||
|
||||
func (ctrl *Controller) finish(cleanup bool, panicErr any) {
|
||||
ctrl.T.Helper()
|
||||
|
||||
ctrl.mu.Lock()
|
||||
defer ctrl.mu.Unlock()
|
||||
|
||||
if ctrl.finished {
|
||||
if _, ok := isCleanuper(ctrl.T); !ok {
|
||||
ctrl.T.Fatalf("Controller.Finish was called more than once. It has to be called exactly once.")
|
||||
}
|
||||
return
|
||||
}
|
||||
ctrl.finished = true
|
||||
|
||||
// Short-circuit, pass through the panic.
|
||||
if panicErr != nil {
|
||||
panic(panicErr)
|
||||
}
|
||||
|
||||
// Check that all remaining expected calls are satisfied.
|
||||
failures := ctrl.expectedCalls.Failures()
|
||||
for _, call := range failures {
|
||||
ctrl.T.Errorf("missing call(s) to %v", call)
|
||||
}
|
||||
if len(failures) != 0 {
|
||||
if !cleanup {
|
||||
ctrl.T.Fatalf("aborting test due to missing call(s)")
|
||||
return
|
||||
}
|
||||
ctrl.T.Errorf("aborting test due to missing call(s)")
|
||||
}
|
||||
}
|
||||
|
||||
// callerInfo returns the file:line of the call site. skip is the number
|
||||
// of stack frames to skip when reporting. 0 is callerInfo's call site.
|
||||
func callerInfo(skip int) string {
|
||||
if _, file, line, ok := runtime.Caller(skip + 1); ok {
|
||||
return fmt.Sprintf("%s:%d", file, line)
|
||||
}
|
||||
return "unknown file"
|
||||
}
|
||||
|
||||
// isCleanuper checks it if t's base TestReporter has a Cleanup method.
|
||||
func isCleanuper(t TestReporter) (cleanuper, bool) {
|
||||
tr := unwrapTestReporter(t)
|
||||
c, ok := tr.(cleanuper)
|
||||
return c, ok
|
||||
}
|
||||
|
||||
// unwrapTestReporter unwraps TestReporter to the base implementation.
|
||||
func unwrapTestReporter(t TestReporter) TestReporter {
|
||||
tr := t
|
||||
switch nt := t.(type) {
|
||||
case *cancelReporter:
|
||||
tr = nt.t
|
||||
if h, check := tr.(*nopTestHelper); check {
|
||||
tr = h.t
|
||||
}
|
||||
case *nopTestHelper:
|
||||
tr = nt.t
|
||||
default:
|
||||
// not wrapped
|
||||
}
|
||||
return tr
|
||||
}
|
|
@ -0,0 +1,60 @@
|
|||
// Copyright 2022 Google LLC
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
// Package gomock is a mock framework for Go.
|
||||
//
|
||||
// Standard usage:
|
||||
//
|
||||
// (1) Define an interface that you wish to mock.
|
||||
// type MyInterface interface {
|
||||
// SomeMethod(x int64, y string)
|
||||
// }
|
||||
// (2) Use mockgen to generate a mock from the interface.
|
||||
// (3) Use the mock in a test:
|
||||
// func TestMyThing(t *testing.T) {
|
||||
// mockCtrl := gomock.NewController(t)
|
||||
// mockObj := something.NewMockMyInterface(mockCtrl)
|
||||
// mockObj.EXPECT().SomeMethod(4, "blah")
|
||||
// // pass mockObj to a real object and play with it.
|
||||
// }
|
||||
//
|
||||
// By default, expected calls are not enforced to run in any particular order.
|
||||
// Call order dependency can be enforced by use of InOrder and/or Call.After.
|
||||
// Call.After can create more varied call order dependencies, but InOrder is
|
||||
// often more convenient.
|
||||
//
|
||||
// The following examples create equivalent call order dependencies.
|
||||
//
|
||||
// Example of using Call.After to chain expected call order:
|
||||
//
|
||||
// firstCall := mockObj.EXPECT().SomeMethod(1, "first")
|
||||
// secondCall := mockObj.EXPECT().SomeMethod(2, "second").After(firstCall)
|
||||
// mockObj.EXPECT().SomeMethod(3, "third").After(secondCall)
|
||||
//
|
||||
// Example of using InOrder to declare expected call order:
|
||||
//
|
||||
// gomock.InOrder(
|
||||
// mockObj.EXPECT().SomeMethod(1, "first"),
|
||||
// mockObj.EXPECT().SomeMethod(2, "second"),
|
||||
// mockObj.EXPECT().SomeMethod(3, "third"),
|
||||
// )
|
||||
//
|
||||
// The standard TestReporter most users will pass to `NewController` is a
|
||||
// `*testing.T` from the context of the test. Note that this will use the
|
||||
// standard `t.Error` and `t.Fatal` methods to report what happened in the test.
|
||||
// In some cases this can leave your testing package in a weird state if global
|
||||
// state is used since `t.Fatal` is like calling panic in the middle of a
|
||||
// function. In these cases it is recommended that you pass in your own
|
||||
// `TestReporter`.
|
||||
package gomock
|
|
@ -0,0 +1,447 @@
|
|||
// Copyright 2010 Google Inc.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package gomock
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"reflect"
|
||||
"regexp"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// A Matcher is a representation of a class of values.
|
||||
// It is used to represent the valid or expected arguments to a mocked method.
|
||||
type Matcher interface {
|
||||
// Matches returns whether x is a match.
|
||||
Matches(x any) bool
|
||||
|
||||
// String describes what the matcher matches.
|
||||
String() string
|
||||
}
|
||||
|
||||
// WantFormatter modifies the given Matcher's String() method to the given
|
||||
// Stringer. This allows for control on how the "Want" is formatted when
|
||||
// printing .
|
||||
func WantFormatter(s fmt.Stringer, m Matcher) Matcher {
|
||||
type matcher interface {
|
||||
Matches(x any) bool
|
||||
}
|
||||
|
||||
return struct {
|
||||
matcher
|
||||
fmt.Stringer
|
||||
}{
|
||||
matcher: m,
|
||||
Stringer: s,
|
||||
}
|
||||
}
|
||||
|
||||
// StringerFunc type is an adapter to allow the use of ordinary functions as
|
||||
// a Stringer. If f is a function with the appropriate signature,
|
||||
// StringerFunc(f) is a Stringer that calls f.
|
||||
type StringerFunc func() string
|
||||
|
||||
// String implements fmt.Stringer.
|
||||
func (f StringerFunc) String() string {
|
||||
return f()
|
||||
}
|
||||
|
||||
// GotFormatter is used to better print failure messages. If a matcher
|
||||
// implements GotFormatter, it will use the result from Got when printing
|
||||
// the failure message.
|
||||
type GotFormatter interface {
|
||||
// Got is invoked with the received value. The result is used when
|
||||
// printing the failure message.
|
||||
Got(got any) string
|
||||
}
|
||||
|
||||
// GotFormatterFunc type is an adapter to allow the use of ordinary
|
||||
// functions as a GotFormatter. If f is a function with the appropriate
|
||||
// signature, GotFormatterFunc(f) is a GotFormatter that calls f.
|
||||
type GotFormatterFunc func(got any) string
|
||||
|
||||
// Got implements GotFormatter.
|
||||
func (f GotFormatterFunc) Got(got any) string {
|
||||
return f(got)
|
||||
}
|
||||
|
||||
// GotFormatterAdapter attaches a GotFormatter to a Matcher.
|
||||
func GotFormatterAdapter(s GotFormatter, m Matcher) Matcher {
|
||||
return struct {
|
||||
GotFormatter
|
||||
Matcher
|
||||
}{
|
||||
GotFormatter: s,
|
||||
Matcher: m,
|
||||
}
|
||||
}
|
||||
|
||||
type anyMatcher struct{}
|
||||
|
||||
func (anyMatcher) Matches(any) bool {
|
||||
return true
|
||||
}
|
||||
|
||||
func (anyMatcher) String() string {
|
||||
return "is anything"
|
||||
}
|
||||
|
||||
type condMatcher[T any] struct {
|
||||
fn func(x T) bool
|
||||
}
|
||||
|
||||
func (c condMatcher[T]) Matches(x any) bool {
|
||||
typed, ok := x.(T)
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
return c.fn(typed)
|
||||
}
|
||||
|
||||
func (c condMatcher[T]) String() string {
|
||||
return "adheres to a custom condition"
|
||||
}
|
||||
|
||||
type eqMatcher struct {
|
||||
x any
|
||||
}
|
||||
|
||||
func (e eqMatcher) Matches(x any) bool {
|
||||
// In case, some value is nil
|
||||
if e.x == nil || x == nil {
|
||||
return reflect.DeepEqual(e.x, x)
|
||||
}
|
||||
|
||||
// Check if types assignable and convert them to common type
|
||||
x1Val := reflect.ValueOf(e.x)
|
||||
x2Val := reflect.ValueOf(x)
|
||||
|
||||
if x1Val.Type().AssignableTo(x2Val.Type()) {
|
||||
x1ValConverted := x1Val.Convert(x2Val.Type())
|
||||
return reflect.DeepEqual(x1ValConverted.Interface(), x2Val.Interface())
|
||||
}
|
||||
|
||||
return false
|
||||
}
|
||||
|
||||
func (e eqMatcher) String() string {
|
||||
return fmt.Sprintf("is equal to %s (%T)", getString(e.x), e.x)
|
||||
}
|
||||
|
||||
type nilMatcher struct{}
|
||||
|
||||
func (nilMatcher) Matches(x any) bool {
|
||||
if x == nil {
|
||||
return true
|
||||
}
|
||||
|
||||
v := reflect.ValueOf(x)
|
||||
switch v.Kind() {
|
||||
case reflect.Chan, reflect.Func, reflect.Interface, reflect.Map,
|
||||
reflect.Ptr, reflect.Slice:
|
||||
return v.IsNil()
|
||||
}
|
||||
|
||||
return false
|
||||
}
|
||||
|
||||
func (nilMatcher) String() string {
|
||||
return "is nil"
|
||||
}
|
||||
|
||||
type notMatcher struct {
|
||||
m Matcher
|
||||
}
|
||||
|
||||
func (n notMatcher) Matches(x any) bool {
|
||||
return !n.m.Matches(x)
|
||||
}
|
||||
|
||||
func (n notMatcher) String() string {
|
||||
return "not(" + n.m.String() + ")"
|
||||
}
|
||||
|
||||
type regexMatcher struct {
|
||||
regex *regexp.Regexp
|
||||
}
|
||||
|
||||
func (m regexMatcher) Matches(x any) bool {
|
||||
switch t := x.(type) {
|
||||
case string:
|
||||
return m.regex.MatchString(t)
|
||||
case []byte:
|
||||
return m.regex.Match(t)
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func (m regexMatcher) String() string {
|
||||
return "matches regex " + m.regex.String()
|
||||
}
|
||||
|
||||
type assignableToTypeOfMatcher struct {
|
||||
targetType reflect.Type
|
||||
}
|
||||
|
||||
func (m assignableToTypeOfMatcher) Matches(x any) bool {
|
||||
return reflect.TypeOf(x).AssignableTo(m.targetType)
|
||||
}
|
||||
|
||||
func (m assignableToTypeOfMatcher) String() string {
|
||||
return "is assignable to " + m.targetType.Name()
|
||||
}
|
||||
|
||||
type anyOfMatcher struct {
|
||||
matchers []Matcher
|
||||
}
|
||||
|
||||
func (am anyOfMatcher) Matches(x any) bool {
|
||||
for _, m := range am.matchers {
|
||||
if m.Matches(x) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (am anyOfMatcher) String() string {
|
||||
ss := make([]string, 0, len(am.matchers))
|
||||
for _, matcher := range am.matchers {
|
||||
ss = append(ss, matcher.String())
|
||||
}
|
||||
return strings.Join(ss, " | ")
|
||||
}
|
||||
|
||||
type allMatcher struct {
|
||||
matchers []Matcher
|
||||
}
|
||||
|
||||
func (am allMatcher) Matches(x any) bool {
|
||||
for _, m := range am.matchers {
|
||||
if !m.Matches(x) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func (am allMatcher) String() string {
|
||||
ss := make([]string, 0, len(am.matchers))
|
||||
for _, matcher := range am.matchers {
|
||||
ss = append(ss, matcher.String())
|
||||
}
|
||||
return strings.Join(ss, "; ")
|
||||
}
|
||||
|
||||
type lenMatcher struct {
|
||||
i int
|
||||
}
|
||||
|
||||
func (m lenMatcher) Matches(x any) bool {
|
||||
v := reflect.ValueOf(x)
|
||||
switch v.Kind() {
|
||||
case reflect.Array, reflect.Chan, reflect.Map, reflect.Slice, reflect.String:
|
||||
return v.Len() == m.i
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func (m lenMatcher) String() string {
|
||||
return fmt.Sprintf("has length %d", m.i)
|
||||
}
|
||||
|
||||
type inAnyOrderMatcher struct {
|
||||
x any
|
||||
}
|
||||
|
||||
func (m inAnyOrderMatcher) Matches(x any) bool {
|
||||
given, ok := m.prepareValue(x)
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
wanted, ok := m.prepareValue(m.x)
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
|
||||
if given.Len() != wanted.Len() {
|
||||
return false
|
||||
}
|
||||
|
||||
usedFromGiven := make([]bool, given.Len())
|
||||
foundFromWanted := make([]bool, wanted.Len())
|
||||
for i := 0; i < wanted.Len(); i++ {
|
||||
wantedMatcher := Eq(wanted.Index(i).Interface())
|
||||
for j := 0; j < given.Len(); j++ {
|
||||
if usedFromGiven[j] {
|
||||
continue
|
||||
}
|
||||
if wantedMatcher.Matches(given.Index(j).Interface()) {
|
||||
foundFromWanted[i] = true
|
||||
usedFromGiven[j] = true
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
missingFromWanted := 0
|
||||
for _, found := range foundFromWanted {
|
||||
if !found {
|
||||
missingFromWanted++
|
||||
}
|
||||
}
|
||||
extraInGiven := 0
|
||||
for _, used := range usedFromGiven {
|
||||
if !used {
|
||||
extraInGiven++
|
||||
}
|
||||
}
|
||||
|
||||
return extraInGiven == 0 && missingFromWanted == 0
|
||||
}
|
||||
|
||||
func (m inAnyOrderMatcher) prepareValue(x any) (reflect.Value, bool) {
|
||||
xValue := reflect.ValueOf(x)
|
||||
switch xValue.Kind() {
|
||||
case reflect.Slice, reflect.Array:
|
||||
return xValue, true
|
||||
default:
|
||||
return reflect.Value{}, false
|
||||
}
|
||||
}
|
||||
|
||||
func (m inAnyOrderMatcher) String() string {
|
||||
return fmt.Sprintf("has the same elements as %v", m.x)
|
||||
}
|
||||
|
||||
// Constructors
|
||||
|
||||
// All returns a composite Matcher that returns true if and only all of the
|
||||
// matchers return true.
|
||||
func All(ms ...Matcher) Matcher { return allMatcher{ms} }
|
||||
|
||||
// Any returns a matcher that always matches.
|
||||
func Any() Matcher { return anyMatcher{} }
|
||||
|
||||
// Cond returns a matcher that matches when the given function returns true
|
||||
// after passing it the parameter to the mock function.
|
||||
// This is particularly useful in case you want to match over a field of a custom struct, or dynamic logic.
|
||||
//
|
||||
// Example usage:
|
||||
//
|
||||
// Cond(func(x int){return x == 1}).Matches(1) // returns true
|
||||
// Cond(func(x int){return x == 2}).Matches(1) // returns false
|
||||
func Cond[T any](fn func(x T) bool) Matcher { return condMatcher[T]{fn} }
|
||||
|
||||
// AnyOf returns a composite Matcher that returns true if at least one of the
|
||||
// matchers returns true.
|
||||
//
|
||||
// Example usage:
|
||||
//
|
||||
// AnyOf(1, 2, 3).Matches(2) // returns true
|
||||
// AnyOf(1, 2, 3).Matches(10) // returns false
|
||||
// AnyOf(Nil(), Len(2)).Matches(nil) // returns true
|
||||
// AnyOf(Nil(), Len(2)).Matches("hi") // returns true
|
||||
// AnyOf(Nil(), Len(2)).Matches("hello") // returns false
|
||||
func AnyOf(xs ...any) Matcher {
|
||||
ms := make([]Matcher, 0, len(xs))
|
||||
for _, x := range xs {
|
||||
if m, ok := x.(Matcher); ok {
|
||||
ms = append(ms, m)
|
||||
} else {
|
||||
ms = append(ms, Eq(x))
|
||||
}
|
||||
}
|
||||
return anyOfMatcher{ms}
|
||||
}
|
||||
|
||||
// Eq returns a matcher that matches on equality.
|
||||
//
|
||||
// Example usage:
|
||||
//
|
||||
// Eq(5).Matches(5) // returns true
|
||||
// Eq(5).Matches(4) // returns false
|
||||
func Eq(x any) Matcher { return eqMatcher{x} }
|
||||
|
||||
// Len returns a matcher that matches on length. This matcher returns false if
|
||||
// is compared to a type that is not an array, chan, map, slice, or string.
|
||||
func Len(i int) Matcher {
|
||||
return lenMatcher{i}
|
||||
}
|
||||
|
||||
// Nil returns a matcher that matches if the received value is nil.
|
||||
//
|
||||
// Example usage:
|
||||
//
|
||||
// var x *bytes.Buffer
|
||||
// Nil().Matches(x) // returns true
|
||||
// x = &bytes.Buffer{}
|
||||
// Nil().Matches(x) // returns false
|
||||
func Nil() Matcher { return nilMatcher{} }
|
||||
|
||||
// Not reverses the results of its given child matcher.
|
||||
//
|
||||
// Example usage:
|
||||
//
|
||||
// Not(Eq(5)).Matches(4) // returns true
|
||||
// Not(Eq(5)).Matches(5) // returns false
|
||||
func Not(x any) Matcher {
|
||||
if m, ok := x.(Matcher); ok {
|
||||
return notMatcher{m}
|
||||
}
|
||||
return notMatcher{Eq(x)}
|
||||
}
|
||||
|
||||
// Regex checks whether parameter matches the associated regex.
|
||||
//
|
||||
// Example usage:
|
||||
//
|
||||
// Regex("[0-9]{2}:[0-9]{2}").Matches("23:02") // returns true
|
||||
// Regex("[0-9]{2}:[0-9]{2}").Matches([]byte{'2', '3', ':', '0', '2'}) // returns true
|
||||
// Regex("[0-9]{2}:[0-9]{2}").Matches("hello world") // returns false
|
||||
// Regex("[0-9]{2}").Matches(21) // returns false as it's not a valid type
|
||||
func Regex(regexStr string) Matcher {
|
||||
return regexMatcher{regex: regexp.MustCompile(regexStr)}
|
||||
}
|
||||
|
||||
// AssignableToTypeOf is a Matcher that matches if the parameter to the mock
|
||||
// function is assignable to the type of the parameter to this function.
|
||||
//
|
||||
// Example usage:
|
||||
//
|
||||
// var s fmt.Stringer = &bytes.Buffer{}
|
||||
// AssignableToTypeOf(s).Matches(time.Second) // returns true
|
||||
// AssignableToTypeOf(s).Matches(99) // returns false
|
||||
//
|
||||
// var ctx = reflect.TypeOf((*context.Context)(nil)).Elem()
|
||||
// AssignableToTypeOf(ctx).Matches(context.Background()) // returns true
|
||||
func AssignableToTypeOf(x any) Matcher {
|
||||
if xt, ok := x.(reflect.Type); ok {
|
||||
return assignableToTypeOfMatcher{xt}
|
||||
}
|
||||
return assignableToTypeOfMatcher{reflect.TypeOf(x)}
|
||||
}
|
||||
|
||||
// InAnyOrder is a Matcher that returns true for collections of the same elements ignoring the order.
|
||||
//
|
||||
// Example usage:
|
||||
//
|
||||
// InAnyOrder([]int{1, 2, 3}).Matches([]int{1, 3, 2}) // returns true
|
||||
// InAnyOrder([]int{1, 2, 3}).Matches([]int{1, 2}) // returns false
|
||||
func InAnyOrder(x any) Matcher {
|
||||
return inAnyOrderMatcher{x}
|
||||
}
|
|
@ -0,0 +1,36 @@
|
|||
package gomock
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"reflect"
|
||||
)
|
||||
|
||||
// getString is a safe way to convert a value to a string for printing results
|
||||
// If the value is a a mock, getString avoids calling the mocked String() method,
|
||||
// which avoids potential deadlocks
|
||||
func getString(x any) string {
|
||||
if isGeneratedMock(x) {
|
||||
return fmt.Sprintf("%T", x)
|
||||
}
|
||||
if s, ok := x.(fmt.Stringer); ok {
|
||||
return s.String()
|
||||
}
|
||||
return fmt.Sprintf("%v", x)
|
||||
}
|
||||
|
||||
// isGeneratedMock checks if the given type has a "isgomock" field,
|
||||
// indicating it is a generated mock.
|
||||
func isGeneratedMock(x any) bool {
|
||||
typ := reflect.TypeOf(x)
|
||||
if typ == nil {
|
||||
return false
|
||||
}
|
||||
if typ.Kind() == reflect.Ptr {
|
||||
typ = typ.Elem()
|
||||
}
|
||||
if typ.Kind() != reflect.Struct {
|
||||
return false
|
||||
}
|
||||
_, isgomock := typ.FieldByName("isgomock")
|
||||
return isgomock
|
||||
}
|
|
@ -0,0 +1,41 @@
|
|||
package main
|
||||
|
||||
import (
|
||||
"flag"
|
||||
"log"
|
||||
"os"
|
||||
)
|
||||
|
||||
const (
|
||||
deprecatedFlagProgOnly = "prog_only"
|
||||
deprecatedFlagExecOnly = "exec_only"
|
||||
)
|
||||
|
||||
var (
|
||||
_ = flag.Bool("prog_only", false, "DEPRECATED (reflect mode) Only generate the reflection program; write it to stdout and exit.")
|
||||
_ = flag.String("exec_only", "", "DEPRECATED (reflect mode) If set, execute this reflection program.")
|
||||
)
|
||||
|
||||
// notifyAboutDeprecatedFlags prints a warning message for a deprecated flags if they are set.
|
||||
func notifyAboutDeprecatedFlags() {
|
||||
const resetColorPostfix = "\033[0m"
|
||||
logger := initWarningLogger()
|
||||
|
||||
flag.Visit(func(f *flag.Flag) {
|
||||
switch f.Name {
|
||||
case deprecatedFlagProgOnly:
|
||||
logger.Println("The -prog_only flag is deprecated and has no effect.", resetColorPostfix)
|
||||
case deprecatedFlagExecOnly:
|
||||
logger.Println("The -exec_only flag is deprecated and has no effect.", resetColorPostfix)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func initWarningLogger() *log.Logger {
|
||||
const (
|
||||
yellowColor = "\033[33m"
|
||||
warningPrefix = yellowColor + "WARNING: "
|
||||
)
|
||||
|
||||
return log.New(os.Stdout, warningPrefix, log.Ldate|log.Ltime)
|
||||
}
|
|
@ -5,9 +5,6 @@
|
|||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
//go:build go1.18
|
||||
// +build go1.18
|
||||
|
||||
package main
|
||||
|
||||
import (
|
||||
|
@ -15,7 +12,6 @@ import (
|
|||
"fmt"
|
||||
"go/ast"
|
||||
"go/token"
|
||||
"strings"
|
||||
|
||||
"go.uber.org/mock/mockgen/model"
|
||||
)
|
||||
|
@ -67,29 +63,6 @@ func (p *fileParser) parseGenericType(pkg string, typ ast.Expr, tps map[string]m
|
|||
return nil, nil
|
||||
}
|
||||
|
||||
func getIdentTypeParams(decl any) string {
|
||||
if decl == nil {
|
||||
return ""
|
||||
}
|
||||
ts, ok := decl.(*ast.TypeSpec)
|
||||
if !ok {
|
||||
return ""
|
||||
}
|
||||
if ts.TypeParams == nil || len(ts.TypeParams.List) == 0 {
|
||||
return ""
|
||||
}
|
||||
var sb strings.Builder
|
||||
sb.WriteString("[")
|
||||
for i, v := range ts.TypeParams.List {
|
||||
if i != 0 {
|
||||
sb.WriteString(", ")
|
||||
}
|
||||
sb.WriteString(v.Names[0].Name)
|
||||
}
|
||||
sb.WriteString("]")
|
||||
return sb.String()
|
||||
}
|
||||
|
||||
func (p *fileParser) parseGenericMethod(field *ast.Field, it *namedInterface, iface *model.Interface, pkg string, tps map[string]model.Type) ([]*model.Method, error) {
|
||||
var indices []ast.Expr
|
||||
var typ ast.Expr
|
|
@ -1,41 +0,0 @@
|
|||
// Copyright 2022 Google LLC
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
//go:build !go1.18
|
||||
// +build !go1.18
|
||||
|
||||
package main
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"go/ast"
|
||||
|
||||
"go.uber.org/mock/mockgen/model"
|
||||
)
|
||||
|
||||
func getTypeSpecTypeParams(ts *ast.TypeSpec) []*ast.Field {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (p *fileParser) parseGenericType(pkg string, typ ast.Expr, tps map[string]model.Type) (model.Type, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func getIdentTypeParams(decl any) string {
|
||||
return ""
|
||||
}
|
||||
|
||||
func (p *fileParser) parseGenericMethod(field *ast.Field, it *namedInterface, iface *model.Interface, pkg string, tps map[string]model.Type) ([]*model.Method, error) {
|
||||
return nil, fmt.Errorf("don't know how to mock method of type %T", field.Type)
|
||||
}
|
|
@ -0,0 +1,21 @@
|
|||
package main
|
||||
|
||||
import (
|
||||
"encoding/gob"
|
||||
"os"
|
||||
|
||||
"go.uber.org/mock/mockgen/model"
|
||||
)
|
||||
|
||||
func gobMode(path string) (*model.Package, error) {
|
||||
in, err := os.Open(path)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer in.Close()
|
||||
var pkg model.Package
|
||||
if err := gob.NewDecoder(in).Decode(&pkg); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &pkg, nil
|
||||
}
|
|
@ -59,14 +59,17 @@ var (
|
|||
mockNames = flag.String("mock_names", "", "Comma-separated interfaceName=mockName pairs of explicit mock names to use. Mock names default to 'Mock'+ interfaceName suffix.")
|
||||
packageOut = flag.String("package", "", "Package of the generated code; defaults to the package of the input with a 'mock_' prefix.")
|
||||
selfPackage = flag.String("self_package", "", "The full package import path for the generated code. The purpose of this flag is to prevent import cycles in the generated code by trying to include its own package. This can happen if the mock's package is set to one of its inputs (usually the main one) and the output is stdio so mockgen cannot detect the final output package. Setting this flag will then tell mockgen which import to exclude.")
|
||||
writeCmdComment = flag.Bool("write_command_comment", true, "Writes the command used as a comment if true.")
|
||||
writePkgComment = flag.Bool("write_package_comment", true, "Writes package documentation comment (godoc) if true.")
|
||||
writeSourceComment = flag.Bool("write_source_comment", true, "Writes original file (source mode) or interface names (reflect mode) comment if true.")
|
||||
writeSourceComment = flag.Bool("write_source_comment", true, "Writes original file (source mode) or interface names (package mode) comment if true.")
|
||||
writeGenerateDirective = flag.Bool("write_generate_directive", false, "Add //go:generate directive to regenerate the mock")
|
||||
copyrightFile = flag.String("copyright_file", "", "Copyright file used to add copyright header")
|
||||
buildConstraint = flag.String("build_constraint", "", "If non-empty, added as //go:build <constraint>")
|
||||
typed = flag.Bool("typed", false, "Generate Type-safe 'Return', 'Do', 'DoAndReturn' function")
|
||||
imports = flag.String("imports", "", "(source mode) Comma-separated name=path pairs of explicit imports to use.")
|
||||
auxFiles = flag.String("aux_files", "", "(source mode) Comma-separated pkg=path pairs of auxiliary Go source files.")
|
||||
excludeInterfaces = flag.String("exclude_interfaces", "", "Comma-separated names of interfaces to be excluded")
|
||||
excludeInterfaces = flag.String("exclude_interfaces", "", "(source mode) Comma-separated names of interfaces to be excluded")
|
||||
modelGob = flag.String("model_gob", "", "Skip package/source loading entirely and use the gob encoded model.Package at the given path")
|
||||
|
||||
debugParser = flag.Bool("debug_parser", false, "Print out parser results only.")
|
||||
showVersion = flag.Bool("version", false, "Print version.")
|
||||
|
@ -76,6 +79,8 @@ func main() {
|
|||
flag.Usage = usage
|
||||
flag.Parse()
|
||||
|
||||
notifyAboutDeprecatedFlags()
|
||||
|
||||
if *showVersion {
|
||||
printVersion()
|
||||
return
|
||||
|
@ -84,7 +89,9 @@ func main() {
|
|||
var pkg *model.Package
|
||||
var err error
|
||||
var packageName string
|
||||
if *source != "" {
|
||||
if *modelGob != "" {
|
||||
pkg, err = gobMode(*modelGob)
|
||||
} else if *source != "" {
|
||||
pkg, err = sourceMode(*source)
|
||||
} else {
|
||||
if flag.NArg() != 2 {
|
||||
|
@ -103,7 +110,8 @@ func main() {
|
|||
log.Fatalf("Parse package name failed: %v", err)
|
||||
}
|
||||
}
|
||||
pkg, err = reflectMode(packageName, interfaces)
|
||||
parser := packageModeParser{}
|
||||
pkg, err = parser.parsePackage(packageName, interfaces)
|
||||
}
|
||||
if err != nil {
|
||||
log.Fatalf("Loading input failed: %v", err)
|
||||
|
@ -116,7 +124,7 @@ func main() {
|
|||
|
||||
outputPackageName := *packageOut
|
||||
if outputPackageName == "" {
|
||||
// pkg.Name in reflect mode is the base name of the import path,
|
||||
// pkg.Name in package mode is the base name of the import path,
|
||||
// which might have characters that are illegal to have in package names.
|
||||
outputPackageName = "mock_" + sanitize(pkg.Name)
|
||||
}
|
||||
|
@ -142,7 +150,9 @@ func main() {
|
|||
}
|
||||
}
|
||||
|
||||
g := new(generator)
|
||||
g := &generator{
|
||||
buildConstraint: *buildConstraint,
|
||||
}
|
||||
if *source != "" {
|
||||
g.filename = *source
|
||||
} else {
|
||||
|
@ -225,20 +235,21 @@ func usage() {
|
|||
flag.PrintDefaults()
|
||||
}
|
||||
|
||||
const usageText = `mockgen has two modes of operation: source and reflect.
|
||||
const usageText = `mockgen has two modes of operation: source and package.
|
||||
|
||||
Source mode generates mock interfaces from a source file.
|
||||
It is enabled by using the -source flag. Other flags that
|
||||
may be useful in this mode are -imports and -aux_files.
|
||||
may be useful in this mode are -imports, -aux_files and -exclude_interfaces.
|
||||
Example:
|
||||
mockgen -source=foo.go [other options]
|
||||
|
||||
Reflect mode generates mock interfaces by building a program
|
||||
that uses reflection to understand interfaces. It is enabled
|
||||
by passing two non-flag arguments: an import path, and a
|
||||
comma-separated list of symbols.
|
||||
Package mode works by specifying the package and interface names.
|
||||
It is enabled by passing two non-flag arguments: an import path, and a
|
||||
comma-separated list of symbols.
|
||||
You can use "." to refer to the current path's package.
|
||||
Example:
|
||||
mockgen database/sql/driver Conn,Driver
|
||||
mockgen . SomeInterface
|
||||
|
||||
`
|
||||
|
||||
|
@ -250,12 +261,13 @@ type generator struct {
|
|||
destination string // may be empty
|
||||
srcPackage, srcInterfaces string // may be empty
|
||||
copyrightHeader string
|
||||
buildConstraint string // may be empty
|
||||
|
||||
packageMap map[string]string // map from import path to package name
|
||||
}
|
||||
|
||||
func (g *generator) p(format string, args ...any) {
|
||||
fmt.Fprintf(&g.buf, g.indent+format+"\n", args...)
|
||||
_, _ = fmt.Fprintf(&g.buf, g.indent+format+"\n", args...)
|
||||
}
|
||||
|
||||
func (g *generator) in() {
|
||||
|
@ -305,6 +317,12 @@ func (g *generator) Generate(pkg *model.Package, outputPkgName string, outputPac
|
|||
g.p("")
|
||||
}
|
||||
|
||||
if g.buildConstraint != "" {
|
||||
g.p("//go:build %s", g.buildConstraint)
|
||||
// https://pkg.go.dev/cmd/go#hdr-Build_constraints:~:text=a%20build%20constraint%20should%20be%20followed%20by%20a%20blank%20line
|
||||
g.p("")
|
||||
}
|
||||
|
||||
g.p("// Code generated by MockGen. DO NOT EDIT.")
|
||||
if *writeSourceComment {
|
||||
if g.filename != "" {
|
||||
|
@ -313,16 +331,18 @@ func (g *generator) Generate(pkg *model.Package, outputPkgName string, outputPac
|
|||
g.p("// Source: %v (interfaces: %v)", g.srcPackage, g.srcInterfaces)
|
||||
}
|
||||
}
|
||||
g.p("//")
|
||||
g.p("// Generated by this command:")
|
||||
g.p("//")
|
||||
// only log the name of the executable, not the full path
|
||||
name := filepath.Base(os.Args[0])
|
||||
if runtime.GOOS == "windows" {
|
||||
name = strings.TrimSuffix(name, ".exe")
|
||||
if *writeCmdComment {
|
||||
g.p("//")
|
||||
g.p("// Generated by this command:")
|
||||
g.p("//")
|
||||
// only log the name of the executable, not the full path
|
||||
name := filepath.Base(os.Args[0])
|
||||
if runtime.GOOS == "windows" {
|
||||
name = strings.TrimSuffix(name, ".exe")
|
||||
}
|
||||
g.p("//\t%v", strings.Join(append([]string{name}, os.Args[1:]...), " "))
|
||||
g.p("//")
|
||||
}
|
||||
g.p("//\t%v", strings.Join(append([]string{name}, os.Args[1:]...), " "))
|
||||
g.p("//")
|
||||
|
||||
// Get all required imports, and generate unique names for them all.
|
||||
im := pkg.Imports()
|
||||
|
@ -392,11 +412,13 @@ func (g *generator) Generate(pkg *model.Package, outputPkgName string, outputPac
|
|||
localNames[pkgName] = true
|
||||
}
|
||||
|
||||
if *writePkgComment {
|
||||
// Ensure there's an empty line before the package to follow the recommendations:
|
||||
// https://github.com/golang/go/wiki/CodeReviewComments#package-comments
|
||||
g.p("")
|
||||
// Ensure there is an empty line between “generated by” block and
|
||||
// package documentation comments to follow the recommendations:
|
||||
// https://go.dev/wiki/CodeReviewComments#package-comments
|
||||
// That is, “generated by” should not be a package comment.
|
||||
g.p("")
|
||||
|
||||
if *writePkgComment {
|
||||
g.p("// Package %v is a generated GoMock package.", outputPkgName)
|
||||
}
|
||||
g.p("package %v", outputPkgName)
|
||||
|
@ -472,6 +494,7 @@ func (g *generator) GenerateMockInterface(intf *model.Interface, outputPackagePa
|
|||
g.in()
|
||||
g.p("ctrl *gomock.Controller")
|
||||
g.p("recorder *%vMockRecorder%v", mockType, shortTp)
|
||||
g.p("isgomock struct{}")
|
||||
g.out()
|
||||
g.p("}")
|
||||
g.p("")
|
||||
|
@ -816,7 +839,7 @@ func createPackageMap(importPaths []string) map[string]string {
|
|||
}
|
||||
pkgMap := make(map[string]string)
|
||||
b := bytes.NewBuffer(nil)
|
||||
args := []string{"list", "-json"}
|
||||
args := []string{"list", "-json=ImportPath,Name"}
|
||||
args = append(args, importPaths...)
|
||||
cmd := exec.Command("go", args...)
|
||||
cmd.Stdout = b
|
||||
|
|
|
@ -0,0 +1,358 @@
|
|||
package main
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"flag"
|
||||
"fmt"
|
||||
"go/types"
|
||||
"strings"
|
||||
|
||||
"go.uber.org/mock/mockgen/model"
|
||||
"golang.org/x/tools/go/packages"
|
||||
)
|
||||
|
||||
var (
|
||||
buildFlags = flag.String("build_flags", "", "(package mode) Additional flags for go build.")
|
||||
)
|
||||
|
||||
type packageModeParser struct {
|
||||
pkgName string
|
||||
}
|
||||
|
||||
func (p *packageModeParser) parsePackage(packageName string, ifaces []string) (*model.Package, error) {
|
||||
p.pkgName = packageName
|
||||
|
||||
pkg, err := p.loadPackage(packageName)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("load package: %w", err)
|
||||
}
|
||||
|
||||
interfaces, err := p.extractInterfacesFromPackage(pkg, ifaces)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("extract interfaces from package: %w", err)
|
||||
}
|
||||
|
||||
return &model.Package{
|
||||
Name: pkg.Types.Name(),
|
||||
PkgPath: packageName,
|
||||
Interfaces: interfaces,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (p *packageModeParser) loadPackage(packageName string) (*packages.Package, error) {
|
||||
var buildFlagsSet []string
|
||||
if *buildFlags != "" {
|
||||
buildFlagsSet = strings.Split(*buildFlags, " ")
|
||||
}
|
||||
|
||||
cfg := &packages.Config{
|
||||
Mode: packages.NeedDeps | packages.NeedImports | packages.NeedTypes | packages.NeedTypesInfo | packages.NeedEmbedFiles,
|
||||
BuildFlags: buildFlagsSet,
|
||||
}
|
||||
pkgs, err := packages.Load(cfg, packageName)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("load packages: %w", err)
|
||||
}
|
||||
|
||||
if len(pkgs) != 1 {
|
||||
return nil, fmt.Errorf("packages length must be 1: %d", len(pkgs))
|
||||
}
|
||||
|
||||
if len(pkgs[0].Errors) > 0 {
|
||||
errs := make([]error, len(pkgs[0].Errors))
|
||||
for i, err := range pkgs[0].Errors {
|
||||
errs[i] = err
|
||||
}
|
||||
|
||||
return nil, errors.Join(errs...)
|
||||
}
|
||||
|
||||
return pkgs[0], nil
|
||||
}
|
||||
|
||||
func (p *packageModeParser) extractInterfacesFromPackage(pkg *packages.Package, ifaces []string) ([]*model.Interface, error) {
|
||||
interfaces := make([]*model.Interface, len(ifaces))
|
||||
for i, iface := range ifaces {
|
||||
obj := pkg.Types.Scope().Lookup(iface)
|
||||
if obj == nil {
|
||||
return nil, fmt.Errorf("interface %s does not exist", iface)
|
||||
}
|
||||
|
||||
modelIface, err := p.parseInterface(obj)
|
||||
if err != nil {
|
||||
return nil, newParseTypeError("parse interface", obj.Name(), err)
|
||||
}
|
||||
|
||||
interfaces[i] = modelIface
|
||||
}
|
||||
|
||||
return interfaces, nil
|
||||
}
|
||||
|
||||
func (p *packageModeParser) parseInterface(obj types.Object) (*model.Interface, error) {
|
||||
named, ok := types.Unalias(obj.Type()).(*types.Named)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("%s is not an interface. it is a %T", obj.Name(), obj.Type().Underlying())
|
||||
}
|
||||
|
||||
iface, ok := named.Underlying().(*types.Interface)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("%s is not an interface. it is a %T", obj.Name(), obj.Type().Underlying())
|
||||
}
|
||||
|
||||
if p.isConstraint(iface) {
|
||||
return nil, fmt.Errorf("interface %s is a constraint", obj.Name())
|
||||
}
|
||||
|
||||
methods := make([]*model.Method, iface.NumMethods())
|
||||
for i := range iface.NumMethods() {
|
||||
method := iface.Method(i)
|
||||
typedMethod, ok := method.Type().(*types.Signature)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("method %s is not a signature", method.Name())
|
||||
}
|
||||
|
||||
modelFunc, err := p.parseFunc(typedMethod)
|
||||
if err != nil {
|
||||
return nil, newParseTypeError("parse method", typedMethod.String(), err)
|
||||
}
|
||||
|
||||
methods[i] = &model.Method{
|
||||
Name: method.Name(),
|
||||
In: modelFunc.In,
|
||||
Out: modelFunc.Out,
|
||||
Variadic: modelFunc.Variadic,
|
||||
}
|
||||
}
|
||||
|
||||
if named.TypeParams() == nil {
|
||||
return &model.Interface{Name: obj.Name(), Methods: methods}, nil
|
||||
}
|
||||
|
||||
typeParams := make([]*model.Parameter, named.TypeParams().Len())
|
||||
for i := range named.TypeParams().Len() {
|
||||
param := named.TypeParams().At(i)
|
||||
typeParam, err := p.parseConstraint(param)
|
||||
if err != nil {
|
||||
return nil, newParseTypeError("parse type parameter", param.String(), err)
|
||||
}
|
||||
|
||||
typeParams[i] = &model.Parameter{Name: param.Obj().Name(), Type: typeParam}
|
||||
}
|
||||
|
||||
return &model.Interface{Name: obj.Name(), Methods: methods, TypeParams: typeParams}, nil
|
||||
}
|
||||
|
||||
func (o *packageModeParser) isConstraint(t *types.Interface) bool {
|
||||
for i := range t.NumEmbeddeds() {
|
||||
embed := t.EmbeddedType(i)
|
||||
if _, ok := embed.Underlying().(*types.Interface); !ok {
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
return false
|
||||
}
|
||||
|
||||
func (p *packageModeParser) parseType(t types.Type) (model.Type, error) {
|
||||
switch t := t.(type) {
|
||||
case *types.Array:
|
||||
elementType, err := p.parseType(t.Elem())
|
||||
if err != nil {
|
||||
return nil, newParseTypeError("parse array type", t.Elem().String(), err)
|
||||
}
|
||||
return &model.ArrayType{Len: int(t.Len()), Type: elementType}, nil
|
||||
case *types.Slice:
|
||||
elementType, err := p.parseType(t.Elem())
|
||||
if err != nil {
|
||||
return nil, newParseTypeError("parse slice type", t.Elem().String(), err)
|
||||
}
|
||||
|
||||
return &model.ArrayType{Len: -1, Type: elementType}, nil
|
||||
case *types.Chan:
|
||||
var dir model.ChanDir
|
||||
switch t.Dir() {
|
||||
case types.RecvOnly:
|
||||
dir = model.RecvDir
|
||||
case types.SendOnly:
|
||||
dir = model.SendDir
|
||||
}
|
||||
|
||||
chanType, err := p.parseType(t.Elem())
|
||||
if err != nil {
|
||||
return nil, newParseTypeError("parse chan type", t.Elem().String(), err)
|
||||
}
|
||||
|
||||
return &model.ChanType{Dir: dir, Type: chanType}, nil
|
||||
case *types.Signature:
|
||||
sig, err := p.parseFunc(t)
|
||||
if err != nil {
|
||||
return nil, newParseTypeError("parse signature", t.String(), err)
|
||||
}
|
||||
|
||||
return sig, nil
|
||||
case *types.Named, *types.Alias:
|
||||
object := t.(interface{ Obj() *types.TypeName })
|
||||
var pkg string
|
||||
if object.Obj().Pkg() != nil {
|
||||
pkg = object.Obj().Pkg().Path()
|
||||
}
|
||||
|
||||
// TypeArgs method not available for aliases in go1.22
|
||||
genericType, ok := t.(interface{ TypeArgs() *types.TypeList })
|
||||
if !ok || genericType.TypeArgs() == nil {
|
||||
return &model.NamedType{
|
||||
Package: pkg,
|
||||
Type: object.Obj().Name(),
|
||||
}, nil
|
||||
}
|
||||
|
||||
typeParams := &model.TypeParametersType{TypeParameters: make([]model.Type, genericType.TypeArgs().Len())}
|
||||
for i := range genericType.TypeArgs().Len() {
|
||||
typeParam := genericType.TypeArgs().At(i)
|
||||
typedParam, err := p.parseType(typeParam)
|
||||
if err != nil {
|
||||
return nil, newParseTypeError("parse type parameter", typeParam.String(), err)
|
||||
}
|
||||
|
||||
typeParams.TypeParameters[i] = typedParam
|
||||
}
|
||||
|
||||
return &model.NamedType{
|
||||
Package: pkg,
|
||||
Type: object.Obj().Name(),
|
||||
TypeParams: typeParams,
|
||||
}, nil
|
||||
case *types.Interface:
|
||||
if t.Empty() {
|
||||
return model.PredeclaredType("any"), nil
|
||||
}
|
||||
|
||||
return nil, fmt.Errorf("cannot handle non-empty unnamed interfaces")
|
||||
case *types.Map:
|
||||
key, err := p.parseType(t.Key())
|
||||
if err != nil {
|
||||
return nil, newParseTypeError("parse map key", t.Key().String(), err)
|
||||
}
|
||||
value, err := p.parseType(t.Elem())
|
||||
if err != nil {
|
||||
return nil, newParseTypeError("parse map value", t.Elem().String(), err)
|
||||
}
|
||||
|
||||
return &model.MapType{Key: key, Value: value}, nil
|
||||
case *types.Pointer:
|
||||
valueType, err := p.parseType(t.Elem())
|
||||
if err != nil {
|
||||
return nil, newParseTypeError("parse pointer type", t.Elem().String(), err)
|
||||
}
|
||||
|
||||
return &model.PointerType{Type: valueType}, nil
|
||||
case *types.Struct:
|
||||
if t.NumFields() > 0 {
|
||||
return nil, fmt.Errorf("cannot handle non-empty unnamed structs")
|
||||
}
|
||||
|
||||
return model.PredeclaredType("struct{}"), nil
|
||||
case *types.Basic:
|
||||
return model.PredeclaredType(t.Name()), nil
|
||||
case *types.Tuple:
|
||||
panic("tuple field") // TODO
|
||||
case *types.TypeParam:
|
||||
return &model.NamedType{Type: t.Obj().Name()}, nil
|
||||
default:
|
||||
panic("unknown type") // TODO
|
||||
}
|
||||
}
|
||||
|
||||
func (p *packageModeParser) parseFunc(sig *types.Signature) (*model.FuncType, error) {
|
||||
var variadic *model.Parameter
|
||||
params := make([]*model.Parameter, 0, sig.Params().Len())
|
||||
for i := range sig.Params().Len() {
|
||||
param := sig.Params().At(i)
|
||||
|
||||
isVariadicParam := i == sig.Params().Len()-1 && sig.Variadic()
|
||||
parseType := param.Type()
|
||||
if isVariadicParam {
|
||||
sliceType, ok := param.Type().(*types.Slice)
|
||||
if !ok {
|
||||
return nil, newParseTypeError("variadic parameter is not a slice", param.String(), nil)
|
||||
}
|
||||
|
||||
parseType = sliceType.Elem()
|
||||
}
|
||||
|
||||
paramType, err := p.parseType(parseType)
|
||||
if err != nil {
|
||||
return nil, newParseTypeError("parse parameter type", parseType.String(), err)
|
||||
}
|
||||
|
||||
modelParameter := &model.Parameter{Type: paramType, Name: param.Name()}
|
||||
|
||||
if isVariadicParam {
|
||||
variadic = modelParameter
|
||||
} else {
|
||||
params = append(params, modelParameter)
|
||||
}
|
||||
}
|
||||
|
||||
if len(params) == 0 {
|
||||
params = nil
|
||||
}
|
||||
|
||||
results := make([]*model.Parameter, sig.Results().Len())
|
||||
for i := range sig.Results().Len() {
|
||||
result := sig.Results().At(i)
|
||||
|
||||
resultType, err := p.parseType(result.Type())
|
||||
if err != nil {
|
||||
return nil, newParseTypeError("parse result type", result.Type().String(), err)
|
||||
}
|
||||
|
||||
results[i] = &model.Parameter{Type: resultType, Name: result.Name()}
|
||||
}
|
||||
|
||||
if len(results) == 0 {
|
||||
results = nil
|
||||
}
|
||||
|
||||
return &model.FuncType{
|
||||
In: params,
|
||||
Out: results,
|
||||
Variadic: variadic,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (p *packageModeParser) parseConstraint(t *types.TypeParam) (model.Type, error) {
|
||||
if t == nil {
|
||||
return nil, fmt.Errorf("nil type param")
|
||||
}
|
||||
|
||||
typeParam, err := p.parseType(t.Constraint())
|
||||
if err != nil {
|
||||
return nil, newParseTypeError("parse constraint type", t.Constraint().String(), err)
|
||||
}
|
||||
|
||||
return typeParam, nil
|
||||
}
|
||||
|
||||
type parseTypeError struct {
|
||||
message string
|
||||
typeString string
|
||||
error error
|
||||
}
|
||||
|
||||
func newParseTypeError(message string, typeString string, error error) *parseTypeError {
|
||||
return &parseTypeError{typeString: typeString, error: error, message: message}
|
||||
}
|
||||
|
||||
func (p parseTypeError) Error() string {
|
||||
if p.error != nil {
|
||||
return fmt.Sprintf("%s: error parsing %s: %s", p.message, p.typeString, p.error)
|
||||
}
|
||||
|
||||
return fmt.Sprintf("%s: error parsing type %s", p.message, p.typeString)
|
||||
}
|
||||
|
||||
func (p parseTypeError) Unwrap() error {
|
||||
return p.error
|
||||
}
|
|
@ -1,256 +0,0 @@
|
|||
// Copyright 2012 Google Inc.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package main
|
||||
|
||||
// This file contains the model construction by reflection.
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/gob"
|
||||
"flag"
|
||||
"fmt"
|
||||
"go/build"
|
||||
"io"
|
||||
"log"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"strings"
|
||||
"text/template"
|
||||
|
||||
"go.uber.org/mock/mockgen/model"
|
||||
)
|
||||
|
||||
var (
|
||||
progOnly = flag.Bool("prog_only", false, "(reflect mode) Only generate the reflection program; write it to stdout and exit.")
|
||||
execOnly = flag.String("exec_only", "", "(reflect mode) If set, execute this reflection program.")
|
||||
buildFlags = flag.String("build_flags", "", "(reflect mode) Additional flags for go build.")
|
||||
)
|
||||
|
||||
// reflectMode generates mocks via reflection on an interface.
|
||||
func reflectMode(importPath string, symbols []string) (*model.Package, error) {
|
||||
if *execOnly != "" {
|
||||
return run(*execOnly)
|
||||
}
|
||||
|
||||
program, err := writeProgram(importPath, symbols)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if *progOnly {
|
||||
if _, err := os.Stdout.Write(program); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
os.Exit(0)
|
||||
}
|
||||
|
||||
wd, _ := os.Getwd()
|
||||
|
||||
// Try to run the reflection program in the current working directory.
|
||||
if p, err := runInDir(program, wd); err == nil {
|
||||
return p, nil
|
||||
}
|
||||
|
||||
// Try to run the program in the same directory as the input package.
|
||||
if p, err := build.Import(importPath, wd, build.FindOnly); err == nil {
|
||||
dir := p.Dir
|
||||
if p, err := runInDir(program, dir); err == nil {
|
||||
return p, nil
|
||||
}
|
||||
}
|
||||
|
||||
// Try to run it in a standard temp directory.
|
||||
return runInDir(program, "")
|
||||
}
|
||||
|
||||
func writeProgram(importPath string, symbols []string) ([]byte, error) {
|
||||
var program bytes.Buffer
|
||||
data := reflectData{
|
||||
ImportPath: importPath,
|
||||
Symbols: symbols,
|
||||
}
|
||||
if err := reflectProgram.Execute(&program, &data); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return program.Bytes(), nil
|
||||
}
|
||||
|
||||
// run the given program and parse the output as a model.Package.
|
||||
func run(program string) (*model.Package, error) {
|
||||
f, err := os.CreateTemp("", "")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
filename := f.Name()
|
||||
defer os.Remove(filename)
|
||||
if err := f.Close(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Run the program.
|
||||
cmd := exec.Command(program, "-output", filename)
|
||||
cmd.Stdout = os.Stdout
|
||||
cmd.Stderr = os.Stderr
|
||||
if err := cmd.Run(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
f, err = os.Open(filename)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Process output.
|
||||
var pkg model.Package
|
||||
if err := gob.NewDecoder(f).Decode(&pkg); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if err := f.Close(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &pkg, nil
|
||||
}
|
||||
|
||||
// runInDir writes the given program into the given dir, runs it there, and
|
||||
// parses the output as a model.Package.
|
||||
func runInDir(program []byte, dir string) (*model.Package, error) {
|
||||
// We use TempDir instead of TempFile so we can control the filename.
|
||||
tmpDir, err := os.MkdirTemp(dir, "gomock_reflect_")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer func() {
|
||||
if err := os.RemoveAll(tmpDir); err != nil {
|
||||
log.Printf("failed to remove temp directory: %s", err)
|
||||
}
|
||||
}()
|
||||
const progSource = "prog.go"
|
||||
var progBinary = "prog.bin"
|
||||
if runtime.GOOS == "windows" {
|
||||
// Windows won't execute a program unless it has a ".exe" suffix.
|
||||
progBinary += ".exe"
|
||||
}
|
||||
|
||||
if err := os.WriteFile(filepath.Join(tmpDir, progSource), program, 0600); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
cmdArgs := []string{}
|
||||
cmdArgs = append(cmdArgs, "build")
|
||||
if *buildFlags != "" {
|
||||
cmdArgs = append(cmdArgs, strings.Split(*buildFlags, " ")...)
|
||||
}
|
||||
cmdArgs = append(cmdArgs, "-o", progBinary, progSource)
|
||||
|
||||
// Build the program.
|
||||
buf := bytes.NewBuffer(nil)
|
||||
cmd := exec.Command("go", cmdArgs...)
|
||||
cmd.Dir = tmpDir
|
||||
cmd.Stdout = os.Stdout
|
||||
cmd.Stderr = io.MultiWriter(os.Stderr, buf)
|
||||
if err := cmd.Run(); err != nil {
|
||||
sErr := buf.String()
|
||||
if strings.Contains(sErr, `cannot find package "."`) &&
|
||||
strings.Contains(sErr, "go.uber.org/mock/mockgen/model") {
|
||||
fmt.Fprint(os.Stderr, "Please reference the steps in the README to fix this error:\n\thttps://go.uber.org/mock#reflect-vendoring-error.\n")
|
||||
return nil, err
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return run(filepath.Join(tmpDir, progBinary))
|
||||
}
|
||||
|
||||
type reflectData struct {
|
||||
ImportPath string
|
||||
Symbols []string
|
||||
}
|
||||
|
||||
// This program reflects on an interface value, and prints the
|
||||
// gob encoding of a model.Package to standard output.
|
||||
// JSON doesn't work because of the model.Type interface.
|
||||
var reflectProgram = template.Must(template.New("program").Parse(`
|
||||
// Code generated by MockGen. DO NOT EDIT.
|
||||
package main
|
||||
|
||||
import (
|
||||
"encoding/gob"
|
||||
"flag"
|
||||
"fmt"
|
||||
"os"
|
||||
"path"
|
||||
"reflect"
|
||||
|
||||
"go.uber.org/mock/mockgen/model"
|
||||
|
||||
pkg_ {{printf "%q" .ImportPath}}
|
||||
)
|
||||
|
||||
var output = flag.String("output", "", "The output file name, or empty to use stdout.")
|
||||
|
||||
func main() {
|
||||
flag.Parse()
|
||||
|
||||
its := []struct{
|
||||
sym string
|
||||
typ reflect.Type
|
||||
}{
|
||||
{{range .Symbols}}
|
||||
{ {{printf "%q" .}}, reflect.TypeOf((*pkg_.{{.}})(nil)).Elem()},
|
||||
{{end}}
|
||||
}
|
||||
pkg := &model.Package{
|
||||
// NOTE: This behaves contrary to documented behaviour if the
|
||||
// package name is not the final component of the import path.
|
||||
// The reflect package doesn't expose the package name, though.
|
||||
Name: path.Base({{printf "%q" .ImportPath}}),
|
||||
}
|
||||
|
||||
for _, it := range its {
|
||||
intf, err := model.InterfaceFromInterfaceType(it.typ)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Reflection: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
intf.Name = it.sym
|
||||
pkg.Interfaces = append(pkg.Interfaces, intf)
|
||||
}
|
||||
|
||||
outfile := os.Stdout
|
||||
if len(*output) != 0 {
|
||||
var err error
|
||||
outfile, err = os.Create(*output)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "failed to open output file %q", *output)
|
||||
}
|
||||
defer func() {
|
||||
if err := outfile.Close(); err != nil {
|
||||
fmt.Fprintf(os.Stderr, "failed to close output file %q", *output)
|
||||
os.Exit(1)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
if err := gob.NewEncoder(outfile).Encode(pkg); err != nil {
|
||||
fmt.Fprintf(os.Stderr, "gob encode: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
}
|
||||
`))
|
|
@ -904,6 +904,10 @@ func (k *skECDSAPublicKey) Verify(data []byte, sig *Signature) error {
|
|||
return errors.New("ssh: signature did not verify")
|
||||
}
|
||||
|
||||
func (k *skECDSAPublicKey) CryptoPublicKey() crypto.PublicKey {
|
||||
return &k.PublicKey
|
||||
}
|
||||
|
||||
type skEd25519PublicKey struct {
|
||||
// application is a URL-like string, typically "ssh:" for SSH.
|
||||
// see openssh/PROTOCOL.u2f for details.
|
||||
|
@ -1000,6 +1004,10 @@ func (k *skEd25519PublicKey) Verify(data []byte, sig *Signature) error {
|
|||
return nil
|
||||
}
|
||||
|
||||
func (k *skEd25519PublicKey) CryptoPublicKey() crypto.PublicKey {
|
||||
return k.PublicKey
|
||||
}
|
||||
|
||||
// NewSignerFromKey takes an *rsa.PrivateKey, *dsa.PrivateKey,
|
||||
// *ecdsa.PrivateKey or any other crypto.Signer and returns a
|
||||
// corresponding Signer instance. ECDSA keys must use P-256, P-384 or
|
||||
|
|
|
@ -462,6 +462,24 @@ func (p *PartialSuccessError) Error() string {
|
|||
// It is returned in ServerAuthError.Errors from NewServerConn.
|
||||
var ErrNoAuth = errors.New("ssh: no auth passed yet")
|
||||
|
||||
// BannerError is an error that can be returned by authentication handlers in
|
||||
// ServerConfig to send a banner message to the client.
|
||||
type BannerError struct {
|
||||
Err error
|
||||
Message string
|
||||
}
|
||||
|
||||
func (b *BannerError) Unwrap() error {
|
||||
return b.Err
|
||||
}
|
||||
|
||||
func (b *BannerError) Error() string {
|
||||
if b.Err == nil {
|
||||
return b.Message
|
||||
}
|
||||
return b.Err.Error()
|
||||
}
|
||||
|
||||
func (s *connection) serverAuthenticate(config *ServerConfig) (*Permissions, error) {
|
||||
sessionID := s.transport.getSessionID()
|
||||
var cache pubKeyCache
|
||||
|
@ -734,6 +752,18 @@ userAuthLoop:
|
|||
config.AuthLogCallback(s, userAuthReq.Method, authErr)
|
||||
}
|
||||
|
||||
var bannerErr *BannerError
|
||||
if errors.As(authErr, &bannerErr) {
|
||||
if bannerErr.Message != "" {
|
||||
bannerMsg := &userAuthBannerMsg{
|
||||
Message: bannerErr.Message,
|
||||
}
|
||||
if err := s.transport.writePacket(Marshal(bannerMsg)); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if authErr == nil {
|
||||
break userAuthLoop
|
||||
}
|
||||
|
|
|
@ -38,6 +38,7 @@ type File struct {
|
|||
Module *Module
|
||||
Go *Go
|
||||
Toolchain *Toolchain
|
||||
Godebug []*Godebug
|
||||
Require []*Require
|
||||
Exclude []*Exclude
|
||||
Replace []*Replace
|
||||
|
@ -65,6 +66,13 @@ type Toolchain struct {
|
|||
Syntax *Line
|
||||
}
|
||||
|
||||
// A Godebug is a single godebug key=value statement.
|
||||
type Godebug struct {
|
||||
Key string
|
||||
Value string
|
||||
Syntax *Line
|
||||
}
|
||||
|
||||
// An Exclude is a single exclude statement.
|
||||
type Exclude struct {
|
||||
Mod module.Version
|
||||
|
@ -289,7 +297,7 @@ func parseToFile(file string, data []byte, fix VersionFixer, strict bool) (parse
|
|||
})
|
||||
}
|
||||
continue
|
||||
case "module", "require", "exclude", "replace", "retract":
|
||||
case "module", "godebug", "require", "exclude", "replace", "retract":
|
||||
for _, l := range x.Line {
|
||||
f.add(&errs, x, l, x.Token[0], l.Token, fix, strict)
|
||||
}
|
||||
|
@ -308,7 +316,9 @@ var laxGoVersionRE = lazyregexp.New(`^v?(([1-9][0-9]*)\.(0|[1-9][0-9]*))([^0-9].
|
|||
|
||||
// Toolchains must be named beginning with `go1`,
|
||||
// like "go1.20.3" or "go1.20.3-gccgo". As a special case, "default" is also permitted.
|
||||
// TODO(samthanawalla): Replace regex with https://pkg.go.dev/go/version#IsValid in 1.23+
|
||||
// Note that this regexp is a much looser condition than go/version.IsValid,
|
||||
// for forward compatibility.
|
||||
// (This code has to be work to identify new toolchains even if we tweak the syntax in the future.)
|
||||
var ToolchainRE = lazyregexp.New(`^default$|^go1($|\.)`)
|
||||
|
||||
func (f *File) add(errs *ErrorList, block *LineBlock, line *Line, verb string, args []string, fix VersionFixer, strict bool) {
|
||||
|
@ -384,7 +394,7 @@ func (f *File) add(errs *ErrorList, block *LineBlock, line *Line, verb string, a
|
|||
if len(args) != 1 {
|
||||
errorf("toolchain directive expects exactly one argument")
|
||||
return
|
||||
} else if strict && !ToolchainRE.MatchString(args[0]) {
|
||||
} else if !ToolchainRE.MatchString(args[0]) {
|
||||
errorf("invalid toolchain version '%s': must match format go1.23.0 or default", args[0])
|
||||
return
|
||||
}
|
||||
|
@ -412,6 +422,22 @@ func (f *File) add(errs *ErrorList, block *LineBlock, line *Line, verb string, a
|
|||
}
|
||||
f.Module.Mod = module.Version{Path: s}
|
||||
|
||||
case "godebug":
|
||||
if len(args) != 1 || strings.ContainsAny(args[0], "\"`',") {
|
||||
errorf("usage: godebug key=value")
|
||||
return
|
||||
}
|
||||
key, value, ok := strings.Cut(args[0], "=")
|
||||
if !ok {
|
||||
errorf("usage: godebug key=value")
|
||||
return
|
||||
}
|
||||
f.Godebug = append(f.Godebug, &Godebug{
|
||||
Key: key,
|
||||
Value: value,
|
||||
Syntax: line,
|
||||
})
|
||||
|
||||
case "require", "exclude":
|
||||
if len(args) != 2 {
|
||||
errorf("usage: %s module/path v1.2.3", verb)
|
||||
|
@ -654,6 +680,22 @@ func (f *WorkFile) add(errs *ErrorList, line *Line, verb string, args []string,
|
|||
f.Toolchain = &Toolchain{Syntax: line}
|
||||
f.Toolchain.Name = args[0]
|
||||
|
||||
case "godebug":
|
||||
if len(args) != 1 || strings.ContainsAny(args[0], "\"`',") {
|
||||
errorf("usage: godebug key=value")
|
||||
return
|
||||
}
|
||||
key, value, ok := strings.Cut(args[0], "=")
|
||||
if !ok {
|
||||
errorf("usage: godebug key=value")
|
||||
return
|
||||
}
|
||||
f.Godebug = append(f.Godebug, &Godebug{
|
||||
Key: key,
|
||||
Value: value,
|
||||
Syntax: line,
|
||||
})
|
||||
|
||||
case "use":
|
||||
if len(args) != 1 {
|
||||
errorf("usage: %s local/dir", verb)
|
||||
|
@ -929,6 +971,15 @@ func (f *File) Format() ([]byte, error) {
|
|||
// Cleanup cleans out all the cleared entries.
|
||||
func (f *File) Cleanup() {
|
||||
w := 0
|
||||
for _, g := range f.Godebug {
|
||||
if g.Key != "" {
|
||||
f.Godebug[w] = g
|
||||
w++
|
||||
}
|
||||
}
|
||||
f.Godebug = f.Godebug[:w]
|
||||
|
||||
w = 0
|
||||
for _, r := range f.Require {
|
||||
if r.Mod.Path != "" {
|
||||
f.Require[w] = r
|
||||
|
@ -1027,6 +1078,45 @@ func (f *File) AddToolchainStmt(name string) error {
|
|||
return nil
|
||||
}
|
||||
|
||||
// AddGodebug sets the first godebug line for key to value,
|
||||
// preserving any existing comments for that line and removing all
|
||||
// other godebug lines for key.
|
||||
//
|
||||
// If no line currently exists for key, AddGodebug adds a new line
|
||||
// at the end of the last godebug block.
|
||||
func (f *File) AddGodebug(key, value string) error {
|
||||
need := true
|
||||
for _, g := range f.Godebug {
|
||||
if g.Key == key {
|
||||
if need {
|
||||
g.Value = value
|
||||
f.Syntax.updateLine(g.Syntax, "godebug", key+"="+value)
|
||||
need = false
|
||||
} else {
|
||||
g.Syntax.markRemoved()
|
||||
*g = Godebug{}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if need {
|
||||
f.addNewGodebug(key, value)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// addNewGodebug adds a new godebug key=value line at the end
|
||||
// of the last godebug block, regardless of any existing godebug lines for key.
|
||||
func (f *File) addNewGodebug(key, value string) {
|
||||
line := f.Syntax.addLine(nil, "godebug", key+"="+value)
|
||||
g := &Godebug{
|
||||
Key: key,
|
||||
Value: value,
|
||||
Syntax: line,
|
||||
}
|
||||
f.Godebug = append(f.Godebug, g)
|
||||
}
|
||||
|
||||
// AddRequire sets the first require line for path to version vers,
|
||||
// preserving any existing comments for that line and removing all
|
||||
// other lines for path.
|
||||
|
@ -1334,6 +1424,16 @@ func (f *File) SetRequireSeparateIndirect(req []*Require) {
|
|||
f.SortBlocks()
|
||||
}
|
||||
|
||||
func (f *File) DropGodebug(key string) error {
|
||||
for _, g := range f.Godebug {
|
||||
if g.Key == key {
|
||||
g.Syntax.markRemoved()
|
||||
*g = Godebug{}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (f *File) DropRequire(path string) error {
|
||||
for _, r := range f.Require {
|
||||
if r.Mod.Path == path {
|
||||
|
|
|
@ -14,6 +14,7 @@ import (
|
|||
type WorkFile struct {
|
||||
Go *Go
|
||||
Toolchain *Toolchain
|
||||
Godebug []*Godebug
|
||||
Use []*Use
|
||||
Replace []*Replace
|
||||
|
||||
|
@ -68,7 +69,7 @@ func ParseWork(file string, data []byte, fix VersionFixer) (*WorkFile, error) {
|
|||
Err: fmt.Errorf("unknown block type: %s", strings.Join(x.Token, " ")),
|
||||
})
|
||||
continue
|
||||
case "use", "replace":
|
||||
case "godebug", "use", "replace":
|
||||
for _, l := range x.Line {
|
||||
f.add(&errs, l, x.Token[0], l.Token, fix)
|
||||
}
|
||||
|
@ -184,6 +185,55 @@ func (f *WorkFile) DropToolchainStmt() {
|
|||
}
|
||||
}
|
||||
|
||||
// AddGodebug sets the first godebug line for key to value,
|
||||
// preserving any existing comments for that line and removing all
|
||||
// other godebug lines for key.
|
||||
//
|
||||
// If no line currently exists for key, AddGodebug adds a new line
|
||||
// at the end of the last godebug block.
|
||||
func (f *WorkFile) AddGodebug(key, value string) error {
|
||||
need := true
|
||||
for _, g := range f.Godebug {
|
||||
if g.Key == key {
|
||||
if need {
|
||||
g.Value = value
|
||||
f.Syntax.updateLine(g.Syntax, "godebug", key+"="+value)
|
||||
need = false
|
||||
} else {
|
||||
g.Syntax.markRemoved()
|
||||
*g = Godebug{}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if need {
|
||||
f.addNewGodebug(key, value)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// addNewGodebug adds a new godebug key=value line at the end
|
||||
// of the last godebug block, regardless of any existing godebug lines for key.
|
||||
func (f *WorkFile) addNewGodebug(key, value string) {
|
||||
line := f.Syntax.addLine(nil, "godebug", key+"="+value)
|
||||
g := &Godebug{
|
||||
Key: key,
|
||||
Value: value,
|
||||
Syntax: line,
|
||||
}
|
||||
f.Godebug = append(f.Godebug, g)
|
||||
}
|
||||
|
||||
func (f *WorkFile) DropGodebug(key string) error {
|
||||
for _, g := range f.Godebug {
|
||||
if g.Key == key {
|
||||
g.Syntax.markRemoved()
|
||||
*g = Godebug{}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (f *WorkFile) AddUse(diskPath, modulePath string) error {
|
||||
need := true
|
||||
for _, d := range f.Use {
|
||||
|
|
|
@ -506,6 +506,7 @@ var badWindowsNames = []string{
|
|||
"PRN",
|
||||
"AUX",
|
||||
"NUL",
|
||||
"COM0",
|
||||
"COM1",
|
||||
"COM2",
|
||||
"COM3",
|
||||
|
@ -515,6 +516,7 @@ var badWindowsNames = []string{
|
|||
"COM7",
|
||||
"COM8",
|
||||
"COM9",
|
||||
"LPT0",
|
||||
"LPT1",
|
||||
"LPT2",
|
||||
"LPT3",
|
||||
|
|
|
@ -17,6 +17,7 @@ package http2 // import "golang.org/x/net/http2"
|
|||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"fmt"
|
||||
"io"
|
||||
|
@ -26,6 +27,7 @@ import (
|
|||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"golang.org/x/net/http/httpguts"
|
||||
)
|
||||
|
@ -210,12 +212,6 @@ type stringWriter interface {
|
|||
WriteString(s string) (n int, err error)
|
||||
}
|
||||
|
||||
// A gate lets two goroutines coordinate their activities.
|
||||
type gate chan struct{}
|
||||
|
||||
func (g gate) Done() { g <- struct{}{} }
|
||||
func (g gate) Wait() { <-g }
|
||||
|
||||
// A closeWaiter is like a sync.WaitGroup but only goes 1 to 0 (open to closed).
|
||||
type closeWaiter chan struct{}
|
||||
|
||||
|
@ -383,3 +379,14 @@ func validPseudoPath(v string) bool {
|
|||
// makes that struct also non-comparable, and generally doesn't add
|
||||
// any size (as long as it's first).
|
||||
type incomparable [0]func()
|
||||
|
||||
// synctestGroupInterface is the methods of synctestGroup used by Server and Transport.
|
||||
// It's defined as an interface here to let us keep synctestGroup entirely test-only
|
||||
// and not a part of non-test builds.
|
||||
type synctestGroupInterface interface {
|
||||
Join()
|
||||
Now() time.Time
|
||||
NewTimer(d time.Duration) timer
|
||||
AfterFunc(d time.Duration, f func()) timer
|
||||
ContextWithTimeout(ctx context.Context, d time.Duration) (context.Context, context.CancelFunc)
|
||||
}
|
||||
|
|
|
@ -154,6 +154,39 @@ type Server struct {
|
|||
// so that we don't embed a Mutex in this struct, which will make the
|
||||
// struct non-copyable, which might break some callers.
|
||||
state *serverInternalState
|
||||
|
||||
// Synchronization group used for testing.
|
||||
// Outside of tests, this is nil.
|
||||
group synctestGroupInterface
|
||||
}
|
||||
|
||||
func (s *Server) markNewGoroutine() {
|
||||
if s.group != nil {
|
||||
s.group.Join()
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) now() time.Time {
|
||||
if s.group != nil {
|
||||
return s.group.Now()
|
||||
}
|
||||
return time.Now()
|
||||
}
|
||||
|
||||
// newTimer creates a new time.Timer, or a synthetic timer in tests.
|
||||
func (s *Server) newTimer(d time.Duration) timer {
|
||||
if s.group != nil {
|
||||
return s.group.NewTimer(d)
|
||||
}
|
||||
return timeTimer{time.NewTimer(d)}
|
||||
}
|
||||
|
||||
// afterFunc creates a new time.AfterFunc timer, or a synthetic timer in tests.
|
||||
func (s *Server) afterFunc(d time.Duration, f func()) timer {
|
||||
if s.group != nil {
|
||||
return s.group.AfterFunc(d, f)
|
||||
}
|
||||
return timeTimer{time.AfterFunc(d, f)}
|
||||
}
|
||||
|
||||
func (s *Server) initialConnRecvWindowSize() int32 {
|
||||
|
@ -400,6 +433,10 @@ func (o *ServeConnOpts) handler() http.Handler {
|
|||
//
|
||||
// The opts parameter is optional. If nil, default values are used.
|
||||
func (s *Server) ServeConn(c net.Conn, opts *ServeConnOpts) {
|
||||
s.serveConn(c, opts, nil)
|
||||
}
|
||||
|
||||
func (s *Server) serveConn(c net.Conn, opts *ServeConnOpts, newf func(*serverConn)) {
|
||||
baseCtx, cancel := serverConnBaseContext(c, opts)
|
||||
defer cancel()
|
||||
|
||||
|
@ -426,6 +463,9 @@ func (s *Server) ServeConn(c net.Conn, opts *ServeConnOpts) {
|
|||
pushEnabled: true,
|
||||
sawClientPreface: opts.SawClientPreface,
|
||||
}
|
||||
if newf != nil {
|
||||
newf(sc)
|
||||
}
|
||||
|
||||
s.state.registerConn(sc)
|
||||
defer s.state.unregisterConn(sc)
|
||||
|
@ -599,8 +639,8 @@ type serverConn struct {
|
|||
inFrameScheduleLoop bool // whether we're in the scheduleFrameWrite loop
|
||||
needToSendGoAway bool // we need to schedule a GOAWAY frame write
|
||||
goAwayCode ErrCode
|
||||
shutdownTimer *time.Timer // nil until used
|
||||
idleTimer *time.Timer // nil if unused
|
||||
shutdownTimer timer // nil until used
|
||||
idleTimer timer // nil if unused
|
||||
|
||||
// Owned by the writeFrameAsync goroutine:
|
||||
headerWriteBuf bytes.Buffer
|
||||
|
@ -649,12 +689,12 @@ type stream struct {
|
|||
flow outflow // limits writing from Handler to client
|
||||
inflow inflow // what the client is allowed to POST/etc to us
|
||||
state streamState
|
||||
resetQueued bool // RST_STREAM queued for write; set by sc.resetStream
|
||||
gotTrailerHeader bool // HEADER frame for trailers was seen
|
||||
wroteHeaders bool // whether we wrote headers (not status 100)
|
||||
readDeadline *time.Timer // nil if unused
|
||||
writeDeadline *time.Timer // nil if unused
|
||||
closeErr error // set before cw is closed
|
||||
resetQueued bool // RST_STREAM queued for write; set by sc.resetStream
|
||||
gotTrailerHeader bool // HEADER frame for trailers was seen
|
||||
wroteHeaders bool // whether we wrote headers (not status 100)
|
||||
readDeadline timer // nil if unused
|
||||
writeDeadline timer // nil if unused
|
||||
closeErr error // set before cw is closed
|
||||
|
||||
trailer http.Header // accumulated trailers
|
||||
reqTrailer http.Header // handler's Request.Trailer
|
||||
|
@ -811,8 +851,9 @@ type readFrameResult struct {
|
|||
// consumer is done with the frame.
|
||||
// It's run on its own goroutine.
|
||||
func (sc *serverConn) readFrames() {
|
||||
gate := make(gate)
|
||||
gateDone := gate.Done
|
||||
sc.srv.markNewGoroutine()
|
||||
gate := make(chan struct{})
|
||||
gateDone := func() { gate <- struct{}{} }
|
||||
for {
|
||||
f, err := sc.framer.ReadFrame()
|
||||
select {
|
||||
|
@ -843,6 +884,7 @@ type frameWriteResult struct {
|
|||
// At most one goroutine can be running writeFrameAsync at a time per
|
||||
// serverConn.
|
||||
func (sc *serverConn) writeFrameAsync(wr FrameWriteRequest, wd *writeData) {
|
||||
sc.srv.markNewGoroutine()
|
||||
var err error
|
||||
if wd == nil {
|
||||
err = wr.write.writeFrame(sc)
|
||||
|
@ -922,13 +964,13 @@ func (sc *serverConn) serve() {
|
|||
sc.setConnState(http.StateIdle)
|
||||
|
||||
if sc.srv.IdleTimeout > 0 {
|
||||
sc.idleTimer = time.AfterFunc(sc.srv.IdleTimeout, sc.onIdleTimer)
|
||||
sc.idleTimer = sc.srv.afterFunc(sc.srv.IdleTimeout, sc.onIdleTimer)
|
||||
defer sc.idleTimer.Stop()
|
||||
}
|
||||
|
||||
go sc.readFrames() // closed by defer sc.conn.Close above
|
||||
|
||||
settingsTimer := time.AfterFunc(firstSettingsTimeout, sc.onSettingsTimer)
|
||||
settingsTimer := sc.srv.afterFunc(firstSettingsTimeout, sc.onSettingsTimer)
|
||||
defer settingsTimer.Stop()
|
||||
|
||||
loopNum := 0
|
||||
|
@ -1057,10 +1099,10 @@ func (sc *serverConn) readPreface() error {
|
|||
errc <- nil
|
||||
}
|
||||
}()
|
||||
timer := time.NewTimer(prefaceTimeout) // TODO: configurable on *Server?
|
||||
timer := sc.srv.newTimer(prefaceTimeout) // TODO: configurable on *Server?
|
||||
defer timer.Stop()
|
||||
select {
|
||||
case <-timer.C:
|
||||
case <-timer.C():
|
||||
return errPrefaceTimeout
|
||||
case err := <-errc:
|
||||
if err == nil {
|
||||
|
@ -1425,7 +1467,7 @@ func (sc *serverConn) goAway(code ErrCode) {
|
|||
|
||||
func (sc *serverConn) shutDownIn(d time.Duration) {
|
||||
sc.serveG.check()
|
||||
sc.shutdownTimer = time.AfterFunc(d, sc.onShutdownTimer)
|
||||
sc.shutdownTimer = sc.srv.afterFunc(d, sc.onShutdownTimer)
|
||||
}
|
||||
|
||||
func (sc *serverConn) resetStream(se StreamError) {
|
||||
|
@ -1639,7 +1681,7 @@ func (sc *serverConn) closeStream(st *stream, err error) {
|
|||
delete(sc.streams, st.id)
|
||||
if len(sc.streams) == 0 {
|
||||
sc.setConnState(http.StateIdle)
|
||||
if sc.srv.IdleTimeout > 0 {
|
||||
if sc.srv.IdleTimeout > 0 && sc.idleTimer != nil {
|
||||
sc.idleTimer.Reset(sc.srv.IdleTimeout)
|
||||
}
|
||||
if h1ServerKeepAlivesDisabled(sc.hs) {
|
||||
|
@ -1661,6 +1703,7 @@ func (sc *serverConn) closeStream(st *stream, err error) {
|
|||
}
|
||||
}
|
||||
st.closeErr = err
|
||||
st.cancelCtx()
|
||||
st.cw.Close() // signals Handler's CloseNotifier, unblocks writes, etc
|
||||
sc.writeSched.CloseStream(st.id)
|
||||
}
|
||||
|
@ -2021,7 +2064,7 @@ func (sc *serverConn) processHeaders(f *MetaHeadersFrame) error {
|
|||
// (in Go 1.8), though. That's a more sane option anyway.
|
||||
if sc.hs.ReadTimeout > 0 {
|
||||
sc.conn.SetReadDeadline(time.Time{})
|
||||
st.readDeadline = time.AfterFunc(sc.hs.ReadTimeout, st.onReadTimeout)
|
||||
st.readDeadline = sc.srv.afterFunc(sc.hs.ReadTimeout, st.onReadTimeout)
|
||||
}
|
||||
|
||||
return sc.scheduleHandler(id, rw, req, handler)
|
||||
|
@ -2119,7 +2162,7 @@ func (sc *serverConn) newStream(id, pusherID uint32, state streamState) *stream
|
|||
st.flow.add(sc.initialStreamSendWindowSize)
|
||||
st.inflow.init(sc.srv.initialStreamRecvWindowSize())
|
||||
if sc.hs.WriteTimeout > 0 {
|
||||
st.writeDeadline = time.AfterFunc(sc.hs.WriteTimeout, st.onWriteTimeout)
|
||||
st.writeDeadline = sc.srv.afterFunc(sc.hs.WriteTimeout, st.onWriteTimeout)
|
||||
}
|
||||
|
||||
sc.streams[id] = st
|
||||
|
@ -2343,6 +2386,7 @@ func (sc *serverConn) handlerDone() {
|
|||
|
||||
// Run on its own goroutine.
|
||||
func (sc *serverConn) runHandler(rw *responseWriter, req *http.Request, handler func(http.ResponseWriter, *http.Request)) {
|
||||
sc.srv.markNewGoroutine()
|
||||
defer sc.sendServeMsg(handlerDoneMsg)
|
||||
didPanic := true
|
||||
defer func() {
|
||||
|
@ -2639,7 +2683,7 @@ func (rws *responseWriterState) writeChunk(p []byte) (n int, err error) {
|
|||
var date string
|
||||
if _, ok := rws.snapHeader["Date"]; !ok {
|
||||
// TODO(bradfitz): be faster here, like net/http? measure.
|
||||
date = time.Now().UTC().Format(http.TimeFormat)
|
||||
date = rws.conn.srv.now().UTC().Format(http.TimeFormat)
|
||||
}
|
||||
|
||||
for _, v := range rws.snapHeader["Trailer"] {
|
||||
|
@ -2761,7 +2805,7 @@ func (rws *responseWriterState) promoteUndeclaredTrailers() {
|
|||
|
||||
func (w *responseWriter) SetReadDeadline(deadline time.Time) error {
|
||||
st := w.rws.stream
|
||||
if !deadline.IsZero() && deadline.Before(time.Now()) {
|
||||
if !deadline.IsZero() && deadline.Before(w.rws.conn.srv.now()) {
|
||||
// If we're setting a deadline in the past, reset the stream immediately
|
||||
// so writes after SetWriteDeadline returns will fail.
|
||||
st.onReadTimeout()
|
||||
|
@ -2777,9 +2821,9 @@ func (w *responseWriter) SetReadDeadline(deadline time.Time) error {
|
|||
if deadline.IsZero() {
|
||||
st.readDeadline = nil
|
||||
} else if st.readDeadline == nil {
|
||||
st.readDeadline = time.AfterFunc(deadline.Sub(time.Now()), st.onReadTimeout)
|
||||
st.readDeadline = sc.srv.afterFunc(deadline.Sub(sc.srv.now()), st.onReadTimeout)
|
||||
} else {
|
||||
st.readDeadline.Reset(deadline.Sub(time.Now()))
|
||||
st.readDeadline.Reset(deadline.Sub(sc.srv.now()))
|
||||
}
|
||||
})
|
||||
return nil
|
||||
|
@ -2787,7 +2831,7 @@ func (w *responseWriter) SetReadDeadline(deadline time.Time) error {
|
|||
|
||||
func (w *responseWriter) SetWriteDeadline(deadline time.Time) error {
|
||||
st := w.rws.stream
|
||||
if !deadline.IsZero() && deadline.Before(time.Now()) {
|
||||
if !deadline.IsZero() && deadline.Before(w.rws.conn.srv.now()) {
|
||||
// If we're setting a deadline in the past, reset the stream immediately
|
||||
// so writes after SetWriteDeadline returns will fail.
|
||||
st.onWriteTimeout()
|
||||
|
@ -2803,9 +2847,9 @@ func (w *responseWriter) SetWriteDeadline(deadline time.Time) error {
|
|||
if deadline.IsZero() {
|
||||
st.writeDeadline = nil
|
||||
} else if st.writeDeadline == nil {
|
||||
st.writeDeadline = time.AfterFunc(deadline.Sub(time.Now()), st.onWriteTimeout)
|
||||
st.writeDeadline = sc.srv.afterFunc(deadline.Sub(sc.srv.now()), st.onWriteTimeout)
|
||||
} else {
|
||||
st.writeDeadline.Reset(deadline.Sub(time.Now()))
|
||||
st.writeDeadline.Reset(deadline.Sub(sc.srv.now()))
|
||||
}
|
||||
})
|
||||
return nil
|
||||
|
|
|
@ -1,331 +0,0 @@
|
|||
// Copyright 2024 The Go Authors. All rights reserved.
|
||||
// Use of this source code is governed by a BSD-style
|
||||
// license that can be found in the LICENSE file.
|
||||
package http2
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// testSyncHooks coordinates goroutines in tests.
|
||||
//
|
||||
// For example, a call to ClientConn.RoundTrip involves several goroutines, including:
|
||||
// - the goroutine running RoundTrip;
|
||||
// - the clientStream.doRequest goroutine, which writes the request; and
|
||||
// - the clientStream.readLoop goroutine, which reads the response.
|
||||
//
|
||||
// Using testSyncHooks, a test can start a RoundTrip and identify when all these goroutines
|
||||
// are blocked waiting for some condition such as reading the Request.Body or waiting for
|
||||
// flow control to become available.
|
||||
//
|
||||
// The testSyncHooks also manage timers and synthetic time in tests.
|
||||
// This permits us to, for example, start a request and cause it to time out waiting for
|
||||
// response headers without resorting to time.Sleep calls.
|
||||
type testSyncHooks struct {
|
||||
// active/inactive act as a mutex and condition variable.
|
||||
//
|
||||
// - neither chan contains a value: testSyncHooks is locked.
|
||||
// - active contains a value: unlocked, and at least one goroutine is not blocked
|
||||
// - inactive contains a value: unlocked, and all goroutines are blocked
|
||||
active chan struct{}
|
||||
inactive chan struct{}
|
||||
|
||||
// goroutine counts
|
||||
total int // total goroutines
|
||||
condwait map[*sync.Cond]int // blocked in sync.Cond.Wait
|
||||
blocked []*testBlockedGoroutine // otherwise blocked
|
||||
|
||||
// fake time
|
||||
now time.Time
|
||||
timers []*fakeTimer
|
||||
|
||||
// Transport testing: Report various events.
|
||||
newclientconn func(*ClientConn)
|
||||
newstream func(*clientStream)
|
||||
}
|
||||
|
||||
// testBlockedGoroutine is a blocked goroutine.
|
||||
type testBlockedGoroutine struct {
|
||||
f func() bool // blocked until f returns true
|
||||
ch chan struct{} // closed when unblocked
|
||||
}
|
||||
|
||||
func newTestSyncHooks() *testSyncHooks {
|
||||
h := &testSyncHooks{
|
||||
active: make(chan struct{}, 1),
|
||||
inactive: make(chan struct{}, 1),
|
||||
condwait: map[*sync.Cond]int{},
|
||||
}
|
||||
h.inactive <- struct{}{}
|
||||
h.now = time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC)
|
||||
return h
|
||||
}
|
||||
|
||||
// lock acquires the testSyncHooks mutex.
|
||||
func (h *testSyncHooks) lock() {
|
||||
select {
|
||||
case <-h.active:
|
||||
case <-h.inactive:
|
||||
}
|
||||
}
|
||||
|
||||
// waitInactive waits for all goroutines to become inactive.
|
||||
func (h *testSyncHooks) waitInactive() {
|
||||
for {
|
||||
<-h.inactive
|
||||
if !h.unlock() {
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// unlock releases the testSyncHooks mutex.
|
||||
// It reports whether any goroutines are active.
|
||||
func (h *testSyncHooks) unlock() (active bool) {
|
||||
// Look for a blocked goroutine which can be unblocked.
|
||||
blocked := h.blocked[:0]
|
||||
unblocked := false
|
||||
for _, b := range h.blocked {
|
||||
if !unblocked && b.f() {
|
||||
unblocked = true
|
||||
close(b.ch)
|
||||
} else {
|
||||
blocked = append(blocked, b)
|
||||
}
|
||||
}
|
||||
h.blocked = blocked
|
||||
|
||||
// Count goroutines blocked on condition variables.
|
||||
condwait := 0
|
||||
for _, count := range h.condwait {
|
||||
condwait += count
|
||||
}
|
||||
|
||||
if h.total > condwait+len(blocked) {
|
||||
h.active <- struct{}{}
|
||||
return true
|
||||
} else {
|
||||
h.inactive <- struct{}{}
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// goRun starts a new goroutine.
|
||||
func (h *testSyncHooks) goRun(f func()) {
|
||||
h.lock()
|
||||
h.total++
|
||||
h.unlock()
|
||||
go func() {
|
||||
defer func() {
|
||||
h.lock()
|
||||
h.total--
|
||||
h.unlock()
|
||||
}()
|
||||
f()
|
||||
}()
|
||||
}
|
||||
|
||||
// blockUntil indicates that a goroutine is blocked waiting for some condition to become true.
|
||||
// It waits until f returns true before proceeding.
|
||||
//
|
||||
// Example usage:
|
||||
//
|
||||
// h.blockUntil(func() bool {
|
||||
// // Is the context done yet?
|
||||
// select {
|
||||
// case <-ctx.Done():
|
||||
// default:
|
||||
// return false
|
||||
// }
|
||||
// return true
|
||||
// })
|
||||
// // Wait for the context to become done.
|
||||
// <-ctx.Done()
|
||||
//
|
||||
// The function f passed to blockUntil must be non-blocking and idempotent.
|
||||
func (h *testSyncHooks) blockUntil(f func() bool) {
|
||||
if f() {
|
||||
return
|
||||
}
|
||||
ch := make(chan struct{})
|
||||
h.lock()
|
||||
h.blocked = append(h.blocked, &testBlockedGoroutine{
|
||||
f: f,
|
||||
ch: ch,
|
||||
})
|
||||
h.unlock()
|
||||
<-ch
|
||||
}
|
||||
|
||||
// broadcast is sync.Cond.Broadcast.
|
||||
func (h *testSyncHooks) condBroadcast(cond *sync.Cond) {
|
||||
h.lock()
|
||||
delete(h.condwait, cond)
|
||||
h.unlock()
|
||||
cond.Broadcast()
|
||||
}
|
||||
|
||||
// broadcast is sync.Cond.Wait.
|
||||
func (h *testSyncHooks) condWait(cond *sync.Cond) {
|
||||
h.lock()
|
||||
h.condwait[cond]++
|
||||
h.unlock()
|
||||
}
|
||||
|
||||
// newTimer creates a new fake timer.
|
||||
func (h *testSyncHooks) newTimer(d time.Duration) timer {
|
||||
h.lock()
|
||||
defer h.unlock()
|
||||
t := &fakeTimer{
|
||||
hooks: h,
|
||||
when: h.now.Add(d),
|
||||
c: make(chan time.Time),
|
||||
}
|
||||
h.timers = append(h.timers, t)
|
||||
return t
|
||||
}
|
||||
|
||||
// afterFunc creates a new fake AfterFunc timer.
|
||||
func (h *testSyncHooks) afterFunc(d time.Duration, f func()) timer {
|
||||
h.lock()
|
||||
defer h.unlock()
|
||||
t := &fakeTimer{
|
||||
hooks: h,
|
||||
when: h.now.Add(d),
|
||||
f: f,
|
||||
}
|
||||
h.timers = append(h.timers, t)
|
||||
return t
|
||||
}
|
||||
|
||||
func (h *testSyncHooks) contextWithTimeout(ctx context.Context, d time.Duration) (context.Context, context.CancelFunc) {
|
||||
ctx, cancel := context.WithCancel(ctx)
|
||||
t := h.afterFunc(d, cancel)
|
||||
return ctx, func() {
|
||||
t.Stop()
|
||||
cancel()
|
||||
}
|
||||
}
|
||||
|
||||
func (h *testSyncHooks) timeUntilEvent() time.Duration {
|
||||
h.lock()
|
||||
defer h.unlock()
|
||||
var next time.Time
|
||||
for _, t := range h.timers {
|
||||
if next.IsZero() || t.when.Before(next) {
|
||||
next = t.when
|
||||
}
|
||||
}
|
||||
if d := next.Sub(h.now); d > 0 {
|
||||
return d
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// advance advances time and causes synthetic timers to fire.
|
||||
func (h *testSyncHooks) advance(d time.Duration) {
|
||||
h.lock()
|
||||
defer h.unlock()
|
||||
h.now = h.now.Add(d)
|
||||
timers := h.timers[:0]
|
||||
for _, t := range h.timers {
|
||||
t := t // remove after go.mod depends on go1.22
|
||||
t.mu.Lock()
|
||||
switch {
|
||||
case t.when.After(h.now):
|
||||
timers = append(timers, t)
|
||||
case t.when.IsZero():
|
||||
// stopped timer
|
||||
default:
|
||||
t.when = time.Time{}
|
||||
if t.c != nil {
|
||||
close(t.c)
|
||||
}
|
||||
if t.f != nil {
|
||||
h.total++
|
||||
go func() {
|
||||
defer func() {
|
||||
h.lock()
|
||||
h.total--
|
||||
h.unlock()
|
||||
}()
|
||||
t.f()
|
||||
}()
|
||||
}
|
||||
}
|
||||
t.mu.Unlock()
|
||||
}
|
||||
h.timers = timers
|
||||
}
|
||||
|
||||
// A timer wraps a time.Timer, or a synthetic equivalent in tests.
|
||||
// Unlike time.Timer, timer is single-use: The timer channel is closed when the timer expires.
|
||||
type timer interface {
|
||||
C() <-chan time.Time
|
||||
Stop() bool
|
||||
Reset(d time.Duration) bool
|
||||
}
|
||||
|
||||
// timeTimer implements timer using real time.
|
||||
type timeTimer struct {
|
||||
t *time.Timer
|
||||
c chan time.Time
|
||||
}
|
||||
|
||||
// newTimeTimer creates a new timer using real time.
|
||||
func newTimeTimer(d time.Duration) timer {
|
||||
ch := make(chan time.Time)
|
||||
t := time.AfterFunc(d, func() {
|
||||
close(ch)
|
||||
})
|
||||
return &timeTimer{t, ch}
|
||||
}
|
||||
|
||||
// newTimeAfterFunc creates an AfterFunc timer using real time.
|
||||
func newTimeAfterFunc(d time.Duration, f func()) timer {
|
||||
return &timeTimer{
|
||||
t: time.AfterFunc(d, f),
|
||||
}
|
||||
}
|
||||
|
||||
func (t timeTimer) C() <-chan time.Time { return t.c }
|
||||
func (t timeTimer) Stop() bool { return t.t.Stop() }
|
||||
func (t timeTimer) Reset(d time.Duration) bool { return t.t.Reset(d) }
|
||||
|
||||
// fakeTimer implements timer using fake time.
|
||||
type fakeTimer struct {
|
||||
hooks *testSyncHooks
|
||||
|
||||
mu sync.Mutex
|
||||
when time.Time // when the timer will fire
|
||||
c chan time.Time // closed when the timer fires; mutually exclusive with f
|
||||
f func() // called when the timer fires; mutually exclusive with c
|
||||
}
|
||||
|
||||
func (t *fakeTimer) C() <-chan time.Time { return t.c }
|
||||
|
||||
func (t *fakeTimer) Stop() bool {
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
stopped := t.when.IsZero()
|
||||
t.when = time.Time{}
|
||||
return stopped
|
||||
}
|
||||
|
||||
func (t *fakeTimer) Reset(d time.Duration) bool {
|
||||
if t.c != nil || t.f == nil {
|
||||
panic("fakeTimer only supports Reset on AfterFunc timers")
|
||||
}
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
t.hooks.lock()
|
||||
defer t.hooks.unlock()
|
||||
active := !t.when.IsZero()
|
||||
t.when = t.hooks.now.Add(d)
|
||||
if !active {
|
||||
t.hooks.timers = append(t.hooks.timers, t)
|
||||
}
|
||||
return active
|
||||
}
|
|
@ -0,0 +1,20 @@
|
|||
// Copyright 2024 The Go Authors. All rights reserved.
|
||||
// Use of this source code is governed by a BSD-style
|
||||
// license that can be found in the LICENSE file.
|
||||
package http2
|
||||
|
||||
import "time"
|
||||
|
||||
// A timer is a time.Timer, as an interface which can be replaced in tests.
|
||||
type timer = interface {
|
||||
C() <-chan time.Time
|
||||
Reset(d time.Duration) bool
|
||||
Stop() bool
|
||||
}
|
||||
|
||||
// timeTimer adapts a time.Timer to the timer interface.
|
||||
type timeTimer struct {
|
||||
*time.Timer
|
||||
}
|
||||
|
||||
func (t timeTimer) C() <-chan time.Time { return t.Timer.C }
|
|
@ -185,7 +185,45 @@ type Transport struct {
|
|||
connPoolOnce sync.Once
|
||||
connPoolOrDef ClientConnPool // non-nil version of ConnPool
|
||||
|
||||
syncHooks *testSyncHooks
|
||||
*transportTestHooks
|
||||
}
|
||||
|
||||
// Hook points used for testing.
|
||||
// Outside of tests, t.transportTestHooks is nil and these all have minimal implementations.
|
||||
// Inside tests, see the testSyncHooks function docs.
|
||||
|
||||
type transportTestHooks struct {
|
||||
newclientconn func(*ClientConn)
|
||||
group synctestGroupInterface
|
||||
}
|
||||
|
||||
func (t *Transport) markNewGoroutine() {
|
||||
if t != nil && t.transportTestHooks != nil {
|
||||
t.transportTestHooks.group.Join()
|
||||
}
|
||||
}
|
||||
|
||||
// newTimer creates a new time.Timer, or a synthetic timer in tests.
|
||||
func (t *Transport) newTimer(d time.Duration) timer {
|
||||
if t.transportTestHooks != nil {
|
||||
return t.transportTestHooks.group.NewTimer(d)
|
||||
}
|
||||
return timeTimer{time.NewTimer(d)}
|
||||
}
|
||||
|
||||
// afterFunc creates a new time.AfterFunc timer, or a synthetic timer in tests.
|
||||
func (t *Transport) afterFunc(d time.Duration, f func()) timer {
|
||||
if t.transportTestHooks != nil {
|
||||
return t.transportTestHooks.group.AfterFunc(d, f)
|
||||
}
|
||||
return timeTimer{time.AfterFunc(d, f)}
|
||||
}
|
||||
|
||||
func (t *Transport) contextWithTimeout(ctx context.Context, d time.Duration) (context.Context, context.CancelFunc) {
|
||||
if t.transportTestHooks != nil {
|
||||
return t.transportTestHooks.group.ContextWithTimeout(ctx, d)
|
||||
}
|
||||
return context.WithTimeout(ctx, d)
|
||||
}
|
||||
|
||||
func (t *Transport) maxHeaderListSize() uint32 {
|
||||
|
@ -352,60 +390,6 @@ type ClientConn struct {
|
|||
werr error // first write error that has occurred
|
||||
hbuf bytes.Buffer // HPACK encoder writes into this
|
||||
henc *hpack.Encoder
|
||||
|
||||
syncHooks *testSyncHooks // can be nil
|
||||
}
|
||||
|
||||
// Hook points used for testing.
|
||||
// Outside of tests, cc.syncHooks is nil and these all have minimal implementations.
|
||||
// Inside tests, see the testSyncHooks function docs.
|
||||
|
||||
// goRun starts a new goroutine.
|
||||
func (cc *ClientConn) goRun(f func()) {
|
||||
if cc.syncHooks != nil {
|
||||
cc.syncHooks.goRun(f)
|
||||
return
|
||||
}
|
||||
go f()
|
||||
}
|
||||
|
||||
// condBroadcast is cc.cond.Broadcast.
|
||||
func (cc *ClientConn) condBroadcast() {
|
||||
if cc.syncHooks != nil {
|
||||
cc.syncHooks.condBroadcast(cc.cond)
|
||||
}
|
||||
cc.cond.Broadcast()
|
||||
}
|
||||
|
||||
// condWait is cc.cond.Wait.
|
||||
func (cc *ClientConn) condWait() {
|
||||
if cc.syncHooks != nil {
|
||||
cc.syncHooks.condWait(cc.cond)
|
||||
}
|
||||
cc.cond.Wait()
|
||||
}
|
||||
|
||||
// newTimer creates a new time.Timer, or a synthetic timer in tests.
|
||||
func (cc *ClientConn) newTimer(d time.Duration) timer {
|
||||
if cc.syncHooks != nil {
|
||||
return cc.syncHooks.newTimer(d)
|
||||
}
|
||||
return newTimeTimer(d)
|
||||
}
|
||||
|
||||
// afterFunc creates a new time.AfterFunc timer, or a synthetic timer in tests.
|
||||
func (cc *ClientConn) afterFunc(d time.Duration, f func()) timer {
|
||||
if cc.syncHooks != nil {
|
||||
return cc.syncHooks.afterFunc(d, f)
|
||||
}
|
||||
return newTimeAfterFunc(d, f)
|
||||
}
|
||||
|
||||
func (cc *ClientConn) contextWithTimeout(ctx context.Context, d time.Duration) (context.Context, context.CancelFunc) {
|
||||
if cc.syncHooks != nil {
|
||||
return cc.syncHooks.contextWithTimeout(ctx, d)
|
||||
}
|
||||
return context.WithTimeout(ctx, d)
|
||||
}
|
||||
|
||||
// clientStream is the state for a single HTTP/2 stream. One of these
|
||||
|
@ -487,7 +471,7 @@ func (cs *clientStream) abortStreamLocked(err error) {
|
|||
// TODO(dneil): Clean up tests where cs.cc.cond is nil.
|
||||
if cs.cc.cond != nil {
|
||||
// Wake up writeRequestBody if it is waiting on flow control.
|
||||
cs.cc.condBroadcast()
|
||||
cs.cc.cond.Broadcast()
|
||||
}
|
||||
}
|
||||
|
||||
|
@ -497,7 +481,7 @@ func (cs *clientStream) abortRequestBodyWrite() {
|
|||
defer cc.mu.Unlock()
|
||||
if cs.reqBody != nil && cs.reqBodyClosed == nil {
|
||||
cs.closeReqBodyLocked()
|
||||
cc.condBroadcast()
|
||||
cc.cond.Broadcast()
|
||||
}
|
||||
}
|
||||
|
||||
|
@ -507,10 +491,11 @@ func (cs *clientStream) closeReqBodyLocked() {
|
|||
}
|
||||
cs.reqBodyClosed = make(chan struct{})
|
||||
reqBodyClosed := cs.reqBodyClosed
|
||||
cs.cc.goRun(func() {
|
||||
go func() {
|
||||
cs.cc.t.markNewGoroutine()
|
||||
cs.reqBody.Close()
|
||||
close(reqBodyClosed)
|
||||
})
|
||||
}()
|
||||
}
|
||||
|
||||
type stickyErrWriter struct {
|
||||
|
@ -626,21 +611,7 @@ func (t *Transport) RoundTripOpt(req *http.Request, opt RoundTripOpt) (*http.Res
|
|||
backoff := float64(uint(1) << (uint(retry) - 1))
|
||||
backoff += backoff * (0.1 * mathrand.Float64())
|
||||
d := time.Second * time.Duration(backoff)
|
||||
var tm timer
|
||||
if t.syncHooks != nil {
|
||||
tm = t.syncHooks.newTimer(d)
|
||||
t.syncHooks.blockUntil(func() bool {
|
||||
select {
|
||||
case <-tm.C():
|
||||
case <-req.Context().Done():
|
||||
default:
|
||||
return false
|
||||
}
|
||||
return true
|
||||
})
|
||||
} else {
|
||||
tm = newTimeTimer(d)
|
||||
}
|
||||
tm := t.newTimer(d)
|
||||
select {
|
||||
case <-tm.C():
|
||||
t.vlogf("RoundTrip retrying after failure: %v", roundTripErr)
|
||||
|
@ -725,8 +696,8 @@ func canRetryError(err error) bool {
|
|||
}
|
||||
|
||||
func (t *Transport) dialClientConn(ctx context.Context, addr string, singleUse bool) (*ClientConn, error) {
|
||||
if t.syncHooks != nil {
|
||||
return t.newClientConn(nil, singleUse, t.syncHooks)
|
||||
if t.transportTestHooks != nil {
|
||||
return t.newClientConn(nil, singleUse)
|
||||
}
|
||||
host, _, err := net.SplitHostPort(addr)
|
||||
if err != nil {
|
||||
|
@ -736,7 +707,7 @@ func (t *Transport) dialClientConn(ctx context.Context, addr string, singleUse b
|
|||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return t.newClientConn(tconn, singleUse, nil)
|
||||
return t.newClientConn(tconn, singleUse)
|
||||
}
|
||||
|
||||
func (t *Transport) newTLSConfig(host string) *tls.Config {
|
||||
|
@ -802,10 +773,10 @@ func (t *Transport) maxEncoderHeaderTableSize() uint32 {
|
|||
}
|
||||
|
||||
func (t *Transport) NewClientConn(c net.Conn) (*ClientConn, error) {
|
||||
return t.newClientConn(c, t.disableKeepAlives(), nil)
|
||||
return t.newClientConn(c, t.disableKeepAlives())
|
||||
}
|
||||
|
||||
func (t *Transport) newClientConn(c net.Conn, singleUse bool, hooks *testSyncHooks) (*ClientConn, error) {
|
||||
func (t *Transport) newClientConn(c net.Conn, singleUse bool) (*ClientConn, error) {
|
||||
cc := &ClientConn{
|
||||
t: t,
|
||||
tconn: c,
|
||||
|
@ -820,16 +791,12 @@ func (t *Transport) newClientConn(c net.Conn, singleUse bool, hooks *testSyncHoo
|
|||
wantSettingsAck: true,
|
||||
pings: make(map[[8]byte]chan struct{}),
|
||||
reqHeaderMu: make(chan struct{}, 1),
|
||||
syncHooks: hooks,
|
||||
}
|
||||
if hooks != nil {
|
||||
hooks.newclientconn(cc)
|
||||
if t.transportTestHooks != nil {
|
||||
t.markNewGoroutine()
|
||||
t.transportTestHooks.newclientconn(cc)
|
||||
c = cc.tconn
|
||||
}
|
||||
if d := t.idleConnTimeout(); d != 0 {
|
||||
cc.idleTimeout = d
|
||||
cc.idleTimer = cc.afterFunc(d, cc.onIdleTimeout)
|
||||
}
|
||||
if VerboseLogs {
|
||||
t.vlogf("http2: Transport creating client conn %p to %v", cc, c.RemoteAddr())
|
||||
}
|
||||
|
@ -893,7 +860,13 @@ func (t *Transport) newClientConn(c net.Conn, singleUse bool, hooks *testSyncHoo
|
|||
return nil, cc.werr
|
||||
}
|
||||
|
||||
cc.goRun(cc.readLoop)
|
||||
// Start the idle timer after the connection is fully initialized.
|
||||
if d := t.idleConnTimeout(); d != 0 {
|
||||
cc.idleTimeout = d
|
||||
cc.idleTimer = t.afterFunc(d, cc.onIdleTimeout)
|
||||
}
|
||||
|
||||
go cc.readLoop()
|
||||
return cc, nil
|
||||
}
|
||||
|
||||
|
@ -901,7 +874,7 @@ func (cc *ClientConn) healthCheck() {
|
|||
pingTimeout := cc.t.pingTimeout()
|
||||
// We don't need to periodically ping in the health check, because the readLoop of ClientConn will
|
||||
// trigger the healthCheck again if there is no frame received.
|
||||
ctx, cancel := cc.contextWithTimeout(context.Background(), pingTimeout)
|
||||
ctx, cancel := cc.t.contextWithTimeout(context.Background(), pingTimeout)
|
||||
defer cancel()
|
||||
cc.vlogf("http2: Transport sending health check")
|
||||
err := cc.Ping(ctx)
|
||||
|
@ -1144,7 +1117,8 @@ func (cc *ClientConn) Shutdown(ctx context.Context) error {
|
|||
// Wait for all in-flight streams to complete or connection to close
|
||||
done := make(chan struct{})
|
||||
cancelled := false // guarded by cc.mu
|
||||
cc.goRun(func() {
|
||||
go func() {
|
||||
cc.t.markNewGoroutine()
|
||||
cc.mu.Lock()
|
||||
defer cc.mu.Unlock()
|
||||
for {
|
||||
|
@ -1156,9 +1130,9 @@ func (cc *ClientConn) Shutdown(ctx context.Context) error {
|
|||
if cancelled {
|
||||
break
|
||||
}
|
||||
cc.condWait()
|
||||
cc.cond.Wait()
|
||||
}
|
||||
})
|
||||
}()
|
||||
shutdownEnterWaitStateHook()
|
||||
select {
|
||||
case <-done:
|
||||
|
@ -1168,7 +1142,7 @@ func (cc *ClientConn) Shutdown(ctx context.Context) error {
|
|||
cc.mu.Lock()
|
||||
// Free the goroutine above
|
||||
cancelled = true
|
||||
cc.condBroadcast()
|
||||
cc.cond.Broadcast()
|
||||
cc.mu.Unlock()
|
||||
return ctx.Err()
|
||||
}
|
||||
|
@ -1206,7 +1180,7 @@ func (cc *ClientConn) closeForError(err error) {
|
|||
for _, cs := range cc.streams {
|
||||
cs.abortStreamLocked(err)
|
||||
}
|
||||
cc.condBroadcast()
|
||||
cc.cond.Broadcast()
|
||||
cc.mu.Unlock()
|
||||
cc.closeConn()
|
||||
}
|
||||
|
@ -1321,23 +1295,30 @@ func (cc *ClientConn) roundTrip(req *http.Request, streamf func(*clientStream))
|
|||
respHeaderRecv: make(chan struct{}),
|
||||
donec: make(chan struct{}),
|
||||
}
|
||||
cc.goRun(func() {
|
||||
cs.doRequest(req)
|
||||
})
|
||||
|
||||
// TODO(bradfitz): this is a copy of the logic in net/http. Unify somewhere?
|
||||
if !cc.t.disableCompression() &&
|
||||
req.Header.Get("Accept-Encoding") == "" &&
|
||||
req.Header.Get("Range") == "" &&
|
||||
!cs.isHead {
|
||||
// Request gzip only, not deflate. Deflate is ambiguous and
|
||||
// not as universally supported anyway.
|
||||
// See: https://zlib.net/zlib_faq.html#faq39
|
||||
//
|
||||
// Note that we don't request this for HEAD requests,
|
||||
// due to a bug in nginx:
|
||||
// http://trac.nginx.org/nginx/ticket/358
|
||||
// https://golang.org/issue/5522
|
||||
//
|
||||
// We don't request gzip if the request is for a range, since
|
||||
// auto-decoding a portion of a gzipped document will just fail
|
||||
// anyway. See https://golang.org/issue/8923
|
||||
cs.requestedGzip = true
|
||||
}
|
||||
|
||||
go cs.doRequest(req, streamf)
|
||||
|
||||
waitDone := func() error {
|
||||
if cc.syncHooks != nil {
|
||||
cc.syncHooks.blockUntil(func() bool {
|
||||
select {
|
||||
case <-cs.donec:
|
||||
case <-ctx.Done():
|
||||
case <-cs.reqCancel:
|
||||
default:
|
||||
return false
|
||||
}
|
||||
return true
|
||||
})
|
||||
}
|
||||
select {
|
||||
case <-cs.donec:
|
||||
return nil
|
||||
|
@ -1398,24 +1379,7 @@ func (cc *ClientConn) roundTrip(req *http.Request, streamf func(*clientStream))
|
|||
return err
|
||||
}
|
||||
|
||||
if streamf != nil {
|
||||
streamf(cs)
|
||||
}
|
||||
|
||||
for {
|
||||
if cc.syncHooks != nil {
|
||||
cc.syncHooks.blockUntil(func() bool {
|
||||
select {
|
||||
case <-cs.respHeaderRecv:
|
||||
case <-cs.abort:
|
||||
case <-ctx.Done():
|
||||
case <-cs.reqCancel:
|
||||
default:
|
||||
return false
|
||||
}
|
||||
return true
|
||||
})
|
||||
}
|
||||
select {
|
||||
case <-cs.respHeaderRecv:
|
||||
return handleResponseHeaders()
|
||||
|
@ -1445,8 +1409,9 @@ func (cc *ClientConn) roundTrip(req *http.Request, streamf func(*clientStream))
|
|||
// doRequest runs for the duration of the request lifetime.
|
||||
//
|
||||
// It sends the request and performs post-request cleanup (closing Request.Body, etc.).
|
||||
func (cs *clientStream) doRequest(req *http.Request) {
|
||||
err := cs.writeRequest(req)
|
||||
func (cs *clientStream) doRequest(req *http.Request, streamf func(*clientStream)) {
|
||||
cs.cc.t.markNewGoroutine()
|
||||
err := cs.writeRequest(req, streamf)
|
||||
cs.cleanupWriteRequest(err)
|
||||
}
|
||||
|
||||
|
@ -1457,7 +1422,7 @@ func (cs *clientStream) doRequest(req *http.Request) {
|
|||
//
|
||||
// It returns non-nil if the request ends otherwise.
|
||||
// If the returned error is StreamError, the error Code may be used in resetting the stream.
|
||||
func (cs *clientStream) writeRequest(req *http.Request) (err error) {
|
||||
func (cs *clientStream) writeRequest(req *http.Request, streamf func(*clientStream)) (err error) {
|
||||
cc := cs.cc
|
||||
ctx := cs.ctx
|
||||
|
||||
|
@ -1471,21 +1436,6 @@ func (cs *clientStream) writeRequest(req *http.Request) (err error) {
|
|||
if cc.reqHeaderMu == nil {
|
||||
panic("RoundTrip on uninitialized ClientConn") // for tests
|
||||
}
|
||||
var newStreamHook func(*clientStream)
|
||||
if cc.syncHooks != nil {
|
||||
newStreamHook = cc.syncHooks.newstream
|
||||
cc.syncHooks.blockUntil(func() bool {
|
||||
select {
|
||||
case cc.reqHeaderMu <- struct{}{}:
|
||||
<-cc.reqHeaderMu
|
||||
case <-cs.reqCancel:
|
||||
case <-ctx.Done():
|
||||
default:
|
||||
return false
|
||||
}
|
||||
return true
|
||||
})
|
||||
}
|
||||
select {
|
||||
case cc.reqHeaderMu <- struct{}{}:
|
||||
case <-cs.reqCancel:
|
||||
|
@ -1510,28 +1460,8 @@ func (cs *clientStream) writeRequest(req *http.Request) (err error) {
|
|||
}
|
||||
cc.mu.Unlock()
|
||||
|
||||
if newStreamHook != nil {
|
||||
newStreamHook(cs)
|
||||
}
|
||||
|
||||
// TODO(bradfitz): this is a copy of the logic in net/http. Unify somewhere?
|
||||
if !cc.t.disableCompression() &&
|
||||
req.Header.Get("Accept-Encoding") == "" &&
|
||||
req.Header.Get("Range") == "" &&
|
||||
!cs.isHead {
|
||||
// Request gzip only, not deflate. Deflate is ambiguous and
|
||||
// not as universally supported anyway.
|
||||
// See: https://zlib.net/zlib_faq.html#faq39
|
||||
//
|
||||
// Note that we don't request this for HEAD requests,
|
||||
// due to a bug in nginx:
|
||||
// http://trac.nginx.org/nginx/ticket/358
|
||||
// https://golang.org/issue/5522
|
||||
//
|
||||
// We don't request gzip if the request is for a range, since
|
||||
// auto-decoding a portion of a gzipped document will just fail
|
||||
// anyway. See https://golang.org/issue/8923
|
||||
cs.requestedGzip = true
|
||||
if streamf != nil {
|
||||
streamf(cs)
|
||||
}
|
||||
|
||||
continueTimeout := cc.t.expectContinueTimeout()
|
||||
|
@ -1594,7 +1524,7 @@ func (cs *clientStream) writeRequest(req *http.Request) (err error) {
|
|||
var respHeaderTimer <-chan time.Time
|
||||
var respHeaderRecv chan struct{}
|
||||
if d := cc.responseHeaderTimeout(); d != 0 {
|
||||
timer := cc.newTimer(d)
|
||||
timer := cc.t.newTimer(d)
|
||||
defer timer.Stop()
|
||||
respHeaderTimer = timer.C()
|
||||
respHeaderRecv = cs.respHeaderRecv
|
||||
|
@ -1603,21 +1533,6 @@ func (cs *clientStream) writeRequest(req *http.Request) (err error) {
|
|||
// or until the request is aborted (via context, error, or otherwise),
|
||||
// whichever comes first.
|
||||
for {
|
||||
if cc.syncHooks != nil {
|
||||
cc.syncHooks.blockUntil(func() bool {
|
||||
select {
|
||||
case <-cs.peerClosed:
|
||||
case <-respHeaderTimer:
|
||||
case <-respHeaderRecv:
|
||||
case <-cs.abort:
|
||||
case <-ctx.Done():
|
||||
case <-cs.reqCancel:
|
||||
default:
|
||||
return false
|
||||
}
|
||||
return true
|
||||
})
|
||||
}
|
||||
select {
|
||||
case <-cs.peerClosed:
|
||||
return nil
|
||||
|
@ -1766,7 +1681,7 @@ func (cc *ClientConn) awaitOpenSlotForStreamLocked(cs *clientStream) error {
|
|||
return nil
|
||||
}
|
||||
cc.pendingRequests++
|
||||
cc.condWait()
|
||||
cc.cond.Wait()
|
||||
cc.pendingRequests--
|
||||
select {
|
||||
case <-cs.abort:
|
||||
|
@ -2028,7 +1943,7 @@ func (cs *clientStream) awaitFlowControl(maxBytes int) (taken int32, err error)
|
|||
cs.flow.take(take)
|
||||
return take, nil
|
||||
}
|
||||
cc.condWait()
|
||||
cc.cond.Wait()
|
||||
}
|
||||
}
|
||||
|
||||
|
@ -2311,7 +2226,7 @@ func (cc *ClientConn) forgetStreamID(id uint32) {
|
|||
}
|
||||
// Wake up writeRequestBody via clientStream.awaitFlowControl and
|
||||
// wake up RoundTrip if there is a pending request.
|
||||
cc.condBroadcast()
|
||||
cc.cond.Broadcast()
|
||||
|
||||
closeOnIdle := cc.singleUse || cc.doNotReuse || cc.t.disableKeepAlives() || cc.goAway != nil
|
||||
if closeOnIdle && cc.streamsReserved == 0 && len(cc.streams) == 0 {
|
||||
|
@ -2333,6 +2248,7 @@ type clientConnReadLoop struct {
|
|||
|
||||
// readLoop runs in its own goroutine and reads and dispatches frames.
|
||||
func (cc *ClientConn) readLoop() {
|
||||
cc.t.markNewGoroutine()
|
||||
rl := &clientConnReadLoop{cc: cc}
|
||||
defer rl.cleanup()
|
||||
cc.readerErr = rl.run()
|
||||
|
@ -2399,7 +2315,7 @@ func (rl *clientConnReadLoop) cleanup() {
|
|||
cs.abortStreamLocked(err)
|
||||
}
|
||||
}
|
||||
cc.condBroadcast()
|
||||
cc.cond.Broadcast()
|
||||
cc.mu.Unlock()
|
||||
}
|
||||
|
||||
|
@ -2436,7 +2352,7 @@ func (rl *clientConnReadLoop) run() error {
|
|||
readIdleTimeout := cc.t.ReadIdleTimeout
|
||||
var t timer
|
||||
if readIdleTimeout != 0 {
|
||||
t = cc.afterFunc(readIdleTimeout, cc.healthCheck)
|
||||
t = cc.t.afterFunc(readIdleTimeout, cc.healthCheck)
|
||||
}
|
||||
for {
|
||||
f, err := cc.fr.ReadFrame()
|
||||
|
@ -3034,7 +2950,7 @@ func (rl *clientConnReadLoop) processSettingsNoWrite(f *SettingsFrame) error {
|
|||
for _, cs := range cc.streams {
|
||||
cs.flow.add(delta)
|
||||
}
|
||||
cc.condBroadcast()
|
||||
cc.cond.Broadcast()
|
||||
|
||||
cc.initialWindowSize = s.Val
|
||||
case SettingHeaderTableSize:
|
||||
|
@ -3089,7 +3005,7 @@ func (rl *clientConnReadLoop) processWindowUpdate(f *WindowUpdateFrame) error {
|
|||
|
||||
return ConnectionError(ErrCodeFlowControl)
|
||||
}
|
||||
cc.condBroadcast()
|
||||
cc.cond.Broadcast()
|
||||
return nil
|
||||
}
|
||||
|
||||
|
@ -3133,7 +3049,8 @@ func (cc *ClientConn) Ping(ctx context.Context) error {
|
|||
}
|
||||
var pingError error
|
||||
errc := make(chan struct{})
|
||||
cc.goRun(func() {
|
||||
go func() {
|
||||
cc.t.markNewGoroutine()
|
||||
cc.wmu.Lock()
|
||||
defer cc.wmu.Unlock()
|
||||
if pingError = cc.fr.WritePing(false, p); pingError != nil {
|
||||
|
@ -3144,20 +3061,7 @@ func (cc *ClientConn) Ping(ctx context.Context) error {
|
|||
close(errc)
|
||||
return
|
||||
}
|
||||
})
|
||||
if cc.syncHooks != nil {
|
||||
cc.syncHooks.blockUntil(func() bool {
|
||||
select {
|
||||
case <-c:
|
||||
case <-errc:
|
||||
case <-ctx.Done():
|
||||
case <-cc.readerDone:
|
||||
default:
|
||||
return false
|
||||
}
|
||||
return true
|
||||
})
|
||||
}
|
||||
}()
|
||||
select {
|
||||
case <-c:
|
||||
return nil
|
||||
|
|
|
@ -443,8 +443,8 @@ func (ws *priorityWriteScheduler) addClosedOrIdleNode(list *[]*priorityNode, max
|
|||
}
|
||||
|
||||
func (ws *priorityWriteScheduler) removeNode(n *priorityNode) {
|
||||
for k := n.kids; k != nil; k = k.next {
|
||||
k.setParent(n.parent)
|
||||
for n.kids != nil {
|
||||
n.kids.setParent(n.parent)
|
||||
}
|
||||
n.setParent(nil)
|
||||
delete(ws.nodes, n.id)
|
||||
|
|
|
@ -8,7 +8,6 @@ import (
|
|||
"bytes"
|
||||
"encoding/binary"
|
||||
"io"
|
||||
"io/ioutil"
|
||||
"math/rand"
|
||||
"net"
|
||||
"runtime"
|
||||
|
@ -173,7 +172,7 @@ func testRacyRead(t *testing.T, c1, c2 net.Conn) {
|
|||
// testRacyWrite tests that it is safe to mutate the input Write buffer
|
||||
// immediately after cancelation has occurred.
|
||||
func testRacyWrite(t *testing.T, c1, c2 net.Conn) {
|
||||
go chunkedCopy(ioutil.Discard, c2)
|
||||
go chunkedCopy(io.Discard, c2)
|
||||
|
||||
var wg sync.WaitGroup
|
||||
defer wg.Wait()
|
||||
|
@ -200,7 +199,7 @@ func testRacyWrite(t *testing.T, c1, c2 net.Conn) {
|
|||
|
||||
// testReadTimeout tests that Read timeouts do not affect Write.
|
||||
func testReadTimeout(t *testing.T, c1, c2 net.Conn) {
|
||||
go chunkedCopy(ioutil.Discard, c2)
|
||||
go chunkedCopy(io.Discard, c2)
|
||||
|
||||
c1.SetReadDeadline(aLongTimeAgo)
|
||||
_, err := c1.Read(make([]byte, 1024))
|
||||
|
|
|
@ -8,7 +8,6 @@ package nettest
|
|||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"io/ioutil"
|
||||
"net"
|
||||
"os"
|
||||
"os/exec"
|
||||
|
@ -226,7 +225,7 @@ func LocalPath() (string, error) {
|
|||
if runtime.GOOS == "darwin" {
|
||||
dir = "/tmp"
|
||||
}
|
||||
f, err := ioutil.TempFile(dir, "go-nettest")
|
||||
f, err := os.CreateTemp(dir, "go-nettest")
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
|
|
@ -137,9 +137,7 @@ func (p *PerHost) AddNetwork(net *net.IPNet) {
|
|||
// AddZone specifies a DNS suffix that will use the bypass proxy. A zone of
|
||||
// "example.com" matches "example.com" and all of its subdomains.
|
||||
func (p *PerHost) AddZone(zone string) {
|
||||
if strings.HasSuffix(zone, ".") {
|
||||
zone = zone[:len(zone)-1]
|
||||
}
|
||||
zone = strings.TrimSuffix(zone, ".")
|
||||
if !strings.HasPrefix(zone, ".") {
|
||||
zone = "." + zone
|
||||
}
|
||||
|
@ -148,8 +146,6 @@ func (p *PerHost) AddZone(zone string) {
|
|||
|
||||
// AddHost specifies a host name that will use the bypass proxy.
|
||||
func (p *PerHost) AddHost(host string) {
|
||||
if strings.HasSuffix(host, ".") {
|
||||
host = host[:len(host)-1]
|
||||
}
|
||||
host = strings.TrimSuffix(host, ".")
|
||||
p.bypassHosts = append(p.bypassHosts, host)
|
||||
}
|
||||
|
|
|
@ -16,7 +16,6 @@ import (
|
|||
"encoding/binary"
|
||||
"fmt"
|
||||
"io"
|
||||
"io/ioutil"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
|
@ -279,7 +278,7 @@ func (handler *hybiFrameHandler) HandleFrame(frame frameReader) (frameReader, er
|
|||
}
|
||||
}
|
||||
if header := frame.HeaderReader(); header != nil {
|
||||
io.Copy(ioutil.Discard, header)
|
||||
io.Copy(io.Discard, header)
|
||||
}
|
||||
switch frame.PayloadType() {
|
||||
case ContinuationFrame:
|
||||
|
@ -294,7 +293,7 @@ func (handler *hybiFrameHandler) HandleFrame(frame frameReader) (frameReader, er
|
|||
if err != nil && err != io.EOF && err != io.ErrUnexpectedEOF {
|
||||
return nil, err
|
||||
}
|
||||
io.Copy(ioutil.Discard, frame)
|
||||
io.Copy(io.Discard, frame)
|
||||
if frame.PayloadType() == PingFrame {
|
||||
if _, err := handler.WritePong(b[:n]); err != nil {
|
||||
return nil, err
|
||||
|
|
|
@ -17,7 +17,6 @@ import (
|
|||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"io/ioutil"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
|
@ -208,7 +207,7 @@ again:
|
|||
n, err = ws.frameReader.Read(msg)
|
||||
if err == io.EOF {
|
||||
if trailer := ws.frameReader.TrailerReader(); trailer != nil {
|
||||
io.Copy(ioutil.Discard, trailer)
|
||||
io.Copy(io.Discard, trailer)
|
||||
}
|
||||
ws.frameReader = nil
|
||||
goto again
|
||||
|
@ -330,7 +329,7 @@ func (cd Codec) Receive(ws *Conn, v interface{}) (err error) {
|
|||
ws.rio.Lock()
|
||||
defer ws.rio.Unlock()
|
||||
if ws.frameReader != nil {
|
||||
_, err = io.Copy(ioutil.Discard, ws.frameReader)
|
||||
_, err = io.Copy(io.Discard, ws.frameReader)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
@ -362,7 +361,7 @@ again:
|
|||
return ErrFrameTooLarge
|
||||
}
|
||||
payloadType := frame.PayloadType()
|
||||
data, err := ioutil.ReadAll(frame)
|
||||
data, err := io.ReadAll(frame)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
|
|
@ -263,6 +263,7 @@ struct ltchars {
|
|||
#include <linux/sched.h>
|
||||
#include <linux/seccomp.h>
|
||||
#include <linux/serial.h>
|
||||
#include <linux/sock_diag.h>
|
||||
#include <linux/sockios.h>
|
||||
#include <linux/taskstats.h>
|
||||
#include <linux/tipc.h>
|
||||
|
@ -549,6 +550,7 @@ ccflags="$@"
|
|||
$2 !~ "NLA_TYPE_MASK" &&
|
||||
$2 !~ /^RTC_VL_(ACCURACY|BACKUP|DATA)/ &&
|
||||
$2 ~ /^(NETLINK|NLM|NLMSG|NLA|IFA|IFAN|RT|RTC|RTCF|RTN|RTPROT|RTNH|ARPHRD|ETH_P|NETNSA)_/ ||
|
||||
$2 ~ /^SOCK_|SK_DIAG_|SKNLGRP_$/ ||
|
||||
$2 ~ /^FIORDCHK$/ ||
|
||||
$2 ~ /^SIOC/ ||
|
||||
$2 ~ /^TIOC/ ||
|
||||
|
|
|
@ -502,6 +502,7 @@ const (
|
|||
BPF_IMM = 0x0
|
||||
BPF_IND = 0x40
|
||||
BPF_JA = 0x0
|
||||
BPF_JCOND = 0xe0
|
||||
BPF_JEQ = 0x10
|
||||
BPF_JGE = 0x30
|
||||
BPF_JGT = 0x20
|
||||
|
@ -657,6 +658,9 @@ const (
|
|||
CAN_NPROTO = 0x8
|
||||
CAN_RAW = 0x1
|
||||
CAN_RAW_FILTER_MAX = 0x200
|
||||
CAN_RAW_XL_VCID_RX_FILTER = 0x4
|
||||
CAN_RAW_XL_VCID_TX_PASS = 0x2
|
||||
CAN_RAW_XL_VCID_TX_SET = 0x1
|
||||
CAN_RTR_FLAG = 0x40000000
|
||||
CAN_SFF_ID_BITS = 0xb
|
||||
CAN_SFF_MASK = 0x7ff
|
||||
|
@ -1339,6 +1343,7 @@ const (
|
|||
F_OFD_SETLK = 0x25
|
||||
F_OFD_SETLKW = 0x26
|
||||
F_OK = 0x0
|
||||
F_SEAL_EXEC = 0x20
|
||||
F_SEAL_FUTURE_WRITE = 0x10
|
||||
F_SEAL_GROW = 0x4
|
||||
F_SEAL_SEAL = 0x1
|
||||
|
@ -1627,6 +1632,7 @@ const (
|
|||
IP_FREEBIND = 0xf
|
||||
IP_HDRINCL = 0x3
|
||||
IP_IPSEC_POLICY = 0x10
|
||||
IP_LOCAL_PORT_RANGE = 0x33
|
||||
IP_MAXPACKET = 0xffff
|
||||
IP_MAX_MEMBERSHIPS = 0x14
|
||||
IP_MF = 0x2000
|
||||
|
@ -1653,6 +1659,7 @@ const (
|
|||
IP_PMTUDISC_OMIT = 0x5
|
||||
IP_PMTUDISC_PROBE = 0x3
|
||||
IP_PMTUDISC_WANT = 0x1
|
||||
IP_PROTOCOL = 0x34
|
||||
IP_RECVERR = 0xb
|
||||
IP_RECVERR_RFC4884 = 0x1a
|
||||
IP_RECVFRAGSIZE = 0x19
|
||||
|
@ -2169,7 +2176,7 @@ const (
|
|||
NFT_SECMARK_CTX_MAXLEN = 0x100
|
||||
NFT_SET_MAXNAMELEN = 0x100
|
||||
NFT_SOCKET_MAX = 0x3
|
||||
NFT_TABLE_F_MASK = 0x3
|
||||
NFT_TABLE_F_MASK = 0x7
|
||||
NFT_TABLE_MAXNAMELEN = 0x100
|
||||
NFT_TRACETYPE_MAX = 0x3
|
||||
NFT_TUNNEL_F_MASK = 0x7
|
||||
|
@ -2403,6 +2410,7 @@ const (
|
|||
PERF_RECORD_MISC_USER = 0x2
|
||||
PERF_SAMPLE_BRANCH_PLM_ALL = 0x7
|
||||
PERF_SAMPLE_WEIGHT_TYPE = 0x1004000
|
||||
PID_FS_MAGIC = 0x50494446
|
||||
PIPEFS_MAGIC = 0x50495045
|
||||
PPPIOCGNPMODE = 0xc008744c
|
||||
PPPIOCNEWUNIT = 0xc004743e
|
||||
|
@ -2896,8 +2904,9 @@ const (
|
|||
RWF_APPEND = 0x10
|
||||
RWF_DSYNC = 0x2
|
||||
RWF_HIPRI = 0x1
|
||||
RWF_NOAPPEND = 0x20
|
||||
RWF_NOWAIT = 0x8
|
||||
RWF_SUPPORTED = 0x1f
|
||||
RWF_SUPPORTED = 0x3f
|
||||
RWF_SYNC = 0x4
|
||||
RWF_WRITE_LIFE_NOT_SET = 0x0
|
||||
SCHED_BATCH = 0x3
|
||||
|
@ -2918,7 +2927,9 @@ const (
|
|||
SCHED_RESET_ON_FORK = 0x40000000
|
||||
SCHED_RR = 0x2
|
||||
SCM_CREDENTIALS = 0x2
|
||||
SCM_PIDFD = 0x4
|
||||
SCM_RIGHTS = 0x1
|
||||
SCM_SECURITY = 0x3
|
||||
SCM_TIMESTAMP = 0x1d
|
||||
SC_LOG_FLUSH = 0x100000
|
||||
SECCOMP_ADDFD_FLAG_SEND = 0x2
|
||||
|
@ -3051,6 +3062,8 @@ const (
|
|||
SIOCSMIIREG = 0x8949
|
||||
SIOCSRARP = 0x8962
|
||||
SIOCWANDEV = 0x894a
|
||||
SK_DIAG_BPF_STORAGE_MAX = 0x3
|
||||
SK_DIAG_BPF_STORAGE_REQ_MAX = 0x1
|
||||
SMACK_MAGIC = 0x43415d53
|
||||
SMART_AUTOSAVE = 0xd2
|
||||
SMART_AUTO_OFFLINE = 0xdb
|
||||
|
@ -3071,6 +3084,8 @@ const (
|
|||
SOCKFS_MAGIC = 0x534f434b
|
||||
SOCK_BUF_LOCK_MASK = 0x3
|
||||
SOCK_DCCP = 0x6
|
||||
SOCK_DESTROY = 0x15
|
||||
SOCK_DIAG_BY_FAMILY = 0x14
|
||||
SOCK_IOC_TYPE = 0x89
|
||||
SOCK_PACKET = 0xa
|
||||
SOCK_RAW = 0x3
|
||||
|
@ -3260,6 +3275,7 @@ const (
|
|||
TCP_MAX_WINSHIFT = 0xe
|
||||
TCP_MD5SIG = 0xe
|
||||
TCP_MD5SIG_EXT = 0x20
|
||||
TCP_MD5SIG_FLAG_IFINDEX = 0x2
|
||||
TCP_MD5SIG_FLAG_PREFIX = 0x1
|
||||
TCP_MD5SIG_MAXKEYLEN = 0x50
|
||||
TCP_MSS = 0x200
|
||||
|
|
|
@ -118,6 +118,7 @@ const (
|
|||
IXOFF = 0x1000
|
||||
IXON = 0x400
|
||||
MAP_32BIT = 0x40
|
||||
MAP_ABOVE4G = 0x80
|
||||
MAP_ANON = 0x20
|
||||
MAP_ANONYMOUS = 0x20
|
||||
MAP_DENYWRITE = 0x800
|
||||
|
|
|
@ -118,6 +118,7 @@ const (
|
|||
IXOFF = 0x1000
|
||||
IXON = 0x400
|
||||
MAP_32BIT = 0x40
|
||||
MAP_ABOVE4G = 0x80
|
||||
MAP_ANON = 0x20
|
||||
MAP_ANONYMOUS = 0x20
|
||||
MAP_DENYWRITE = 0x800
|
||||
|
|
|
@ -87,6 +87,7 @@ const (
|
|||
FICLONE = 0x40049409
|
||||
FICLONERANGE = 0x4020940d
|
||||
FLUSHO = 0x1000
|
||||
FPMR_MAGIC = 0x46504d52
|
||||
FPSIMD_MAGIC = 0x46508001
|
||||
FS_IOC_ENABLE_VERITY = 0x40806685
|
||||
FS_IOC_GETFLAGS = 0x80086601
|
||||
|
|
|
@ -4605,7 +4605,7 @@ const (
|
|||
NL80211_ATTR_MAC_HINT = 0xc8
|
||||
NL80211_ATTR_MAC_MASK = 0xd7
|
||||
NL80211_ATTR_MAX_AP_ASSOC_STA = 0xca
|
||||
NL80211_ATTR_MAX = 0x149
|
||||
NL80211_ATTR_MAX = 0x14a
|
||||
NL80211_ATTR_MAX_CRIT_PROT_DURATION = 0xb4
|
||||
NL80211_ATTR_MAX_CSA_COUNTERS = 0xce
|
||||
NL80211_ATTR_MAX_MATCH_SETS = 0x85
|
||||
|
@ -5209,7 +5209,7 @@ const (
|
|||
NL80211_FREQUENCY_ATTR_GO_CONCURRENT = 0xf
|
||||
NL80211_FREQUENCY_ATTR_INDOOR_ONLY = 0xe
|
||||
NL80211_FREQUENCY_ATTR_IR_CONCURRENT = 0xf
|
||||
NL80211_FREQUENCY_ATTR_MAX = 0x1f
|
||||
NL80211_FREQUENCY_ATTR_MAX = 0x20
|
||||
NL80211_FREQUENCY_ATTR_MAX_TX_POWER = 0x6
|
||||
NL80211_FREQUENCY_ATTR_NO_10MHZ = 0x11
|
||||
NL80211_FREQUENCY_ATTR_NO_160MHZ = 0xc
|
||||
|
@ -5703,7 +5703,7 @@ const (
|
|||
NL80211_STA_FLAG_ASSOCIATED = 0x7
|
||||
NL80211_STA_FLAG_AUTHENTICATED = 0x5
|
||||
NL80211_STA_FLAG_AUTHORIZED = 0x1
|
||||
NL80211_STA_FLAG_MAX = 0x7
|
||||
NL80211_STA_FLAG_MAX = 0x8
|
||||
NL80211_STA_FLAG_MAX_OLD_API = 0x6
|
||||
NL80211_STA_FLAG_MFP = 0x4
|
||||
NL80211_STA_FLAG_SHORT_PREAMBLE = 0x2
|
||||
|
@ -6001,3 +6001,34 @@ type CachestatRange struct {
|
|||
Off uint64
|
||||
Len uint64
|
||||
}
|
||||
|
||||
const (
|
||||
SK_MEMINFO_RMEM_ALLOC = 0x0
|
||||
SK_MEMINFO_RCVBUF = 0x1
|
||||
SK_MEMINFO_WMEM_ALLOC = 0x2
|
||||
SK_MEMINFO_SNDBUF = 0x3
|
||||
SK_MEMINFO_FWD_ALLOC = 0x4
|
||||
SK_MEMINFO_WMEM_QUEUED = 0x5
|
||||
SK_MEMINFO_OPTMEM = 0x6
|
||||
SK_MEMINFO_BACKLOG = 0x7
|
||||
SK_MEMINFO_DROPS = 0x8
|
||||
SK_MEMINFO_VARS = 0x9
|
||||
SKNLGRP_NONE = 0x0
|
||||
SKNLGRP_INET_TCP_DESTROY = 0x1
|
||||
SKNLGRP_INET_UDP_DESTROY = 0x2
|
||||
SKNLGRP_INET6_TCP_DESTROY = 0x3
|
||||
SKNLGRP_INET6_UDP_DESTROY = 0x4
|
||||
SK_DIAG_BPF_STORAGE_REQ_NONE = 0x0
|
||||
SK_DIAG_BPF_STORAGE_REQ_MAP_FD = 0x1
|
||||
SK_DIAG_BPF_STORAGE_REP_NONE = 0x0
|
||||
SK_DIAG_BPF_STORAGE = 0x1
|
||||
SK_DIAG_BPF_STORAGE_NONE = 0x0
|
||||
SK_DIAG_BPF_STORAGE_PAD = 0x1
|
||||
SK_DIAG_BPF_STORAGE_MAP_ID = 0x2
|
||||
SK_DIAG_BPF_STORAGE_MAP_VALUE = 0x3
|
||||
)
|
||||
|
||||
type SockDiagReq struct {
|
||||
Family uint8
|
||||
Protocol uint8
|
||||
}
|
||||
|
|
|
@ -68,6 +68,7 @@ type UserInfo10 struct {
|
|||
//sys NetUserGetInfo(serverName *uint16, userName *uint16, level uint32, buf **byte) (neterr error) = netapi32.NetUserGetInfo
|
||||
//sys NetGetJoinInformation(server *uint16, name **uint16, bufType *uint32) (neterr error) = netapi32.NetGetJoinInformation
|
||||
//sys NetApiBufferFree(buf *byte) (neterr error) = netapi32.NetApiBufferFree
|
||||
//sys NetUserEnum(serverName *uint16, level uint32, filter uint32, buf **byte, prefMaxLen uint32, entriesRead *uint32, totalEntries *uint32, resumeHandle *uint32) (neterr error) = netapi32.NetUserEnum
|
||||
|
||||
const (
|
||||
// do not reorder
|
||||
|
|
|
@ -401,6 +401,7 @@ var (
|
|||
procTransmitFile = modmswsock.NewProc("TransmitFile")
|
||||
procNetApiBufferFree = modnetapi32.NewProc("NetApiBufferFree")
|
||||
procNetGetJoinInformation = modnetapi32.NewProc("NetGetJoinInformation")
|
||||
procNetUserEnum = modnetapi32.NewProc("NetUserEnum")
|
||||
procNetUserGetInfo = modnetapi32.NewProc("NetUserGetInfo")
|
||||
procNtCreateFile = modntdll.NewProc("NtCreateFile")
|
||||
procNtCreateNamedPipeFile = modntdll.NewProc("NtCreateNamedPipeFile")
|
||||
|
@ -3486,6 +3487,14 @@ func NetGetJoinInformation(server *uint16, name **uint16, bufType *uint32) (nete
|
|||
return
|
||||
}
|
||||
|
||||
func NetUserEnum(serverName *uint16, level uint32, filter uint32, buf **byte, prefMaxLen uint32, entriesRead *uint32, totalEntries *uint32, resumeHandle *uint32) (neterr error) {
|
||||
r0, _, _ := syscall.Syscall9(procNetUserEnum.Addr(), 8, uintptr(unsafe.Pointer(serverName)), uintptr(level), uintptr(filter), uintptr(unsafe.Pointer(buf)), uintptr(prefMaxLen), uintptr(unsafe.Pointer(entriesRead)), uintptr(unsafe.Pointer(totalEntries)), uintptr(unsafe.Pointer(resumeHandle)), 0)
|
||||
if r0 != 0 {
|
||||
neterr = syscall.Errno(r0)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func NetUserGetInfo(serverName *uint16, userName *uint16, level uint32, buf **byte) (neterr error) {
|
||||
r0, _, _ := syscall.Syscall6(procNetUserGetInfo.Addr(), 4, uintptr(unsafe.Pointer(serverName)), uintptr(unsafe.Pointer(userName)), uintptr(level), uintptr(unsafe.Pointer(buf)), 0, 0)
|
||||
if r0 != 0 {
|
||||
|
|
|
@ -13,6 +13,7 @@ import (
|
|||
"golang.org/x/tools/internal/gocommand"
|
||||
)
|
||||
|
||||
// TODO(adonovan): move back into go/packages.
|
||||
func GetSizesForArgsGolist(ctx context.Context, inv gocommand.Invocation, gocmdRunner *gocommand.Runner) (string, string, error) {
|
||||
inv.Verb = "list"
|
||||
inv.Args = []string{"-f", "{{context.GOARCH}} {{context.Compiler}}", "--", "unsafe"}
|
||||
|
|
|
@ -198,14 +198,6 @@ Instead, ssadump no longer requests the runtime package,
|
|||
but seeks it among the dependencies of the user-specified packages,
|
||||
and emits an error if it is not found.
|
||||
|
||||
Overlays: The Overlay field in the Config allows providing alternate contents
|
||||
for Go source files, by providing a mapping from file path to contents.
|
||||
go/packages will pull in new imports added in overlay files when go/packages
|
||||
is run in LoadImports mode or greater.
|
||||
Overlay support for the go list driver isn't complete yet: if the file doesn't
|
||||
exist on disk, it will only be recognized in an overlay if it is a non-test file
|
||||
and the package would be reported even without the overlay.
|
||||
|
||||
Questions & Tasks
|
||||
|
||||
- Add GOARCH/GOOS?
|
||||
|
|
|
@ -34,8 +34,8 @@ type DriverRequest struct {
|
|||
// Tests specifies whether the patterns should also return test packages.
|
||||
Tests bool `json:"tests"`
|
||||
|
||||
// Overlay maps file paths (relative to the driver's working directory) to the byte contents
|
||||
// of overlay files.
|
||||
// Overlay maps file paths (relative to the driver's working directory)
|
||||
// to the contents of overlay files (see Config.Overlay).
|
||||
Overlay map[string][]byte `json:"overlay"`
|
||||
}
|
||||
|
||||
|
@ -119,7 +119,19 @@ func findExternalDriver(cfg *Config) driver {
|
|||
stderr := new(bytes.Buffer)
|
||||
cmd := exec.CommandContext(cfg.Context, tool, words...)
|
||||
cmd.Dir = cfg.Dir
|
||||
cmd.Env = cfg.Env
|
||||
// The cwd gets resolved to the real path. On Darwin, where
|
||||
// /tmp is a symlink, this breaks anything that expects the
|
||||
// working directory to keep the original path, including the
|
||||
// go command when dealing with modules.
|
||||
//
|
||||
// os.Getwd stdlib has a special feature where if the
|
||||
// cwd and the PWD are the same node then it trusts
|
||||
// the PWD, so by setting it in the env for the child
|
||||
// process we fix up all the paths returned by the go
|
||||
// command.
|
||||
//
|
||||
// (See similar trick in Invocation.run in ../../internal/gocommand/invoke.go)
|
||||
cmd.Env = append(slicesClip(cfg.Env), "PWD="+cfg.Dir)
|
||||
cmd.Stdin = bytes.NewReader(req)
|
||||
cmd.Stdout = buf
|
||||
cmd.Stderr = stderr
|
||||
|
@ -138,3 +150,7 @@ func findExternalDriver(cfg *Config) driver {
|
|||
return &response, nil
|
||||
}
|
||||
}
|
||||
|
||||
// slicesClip removes unused capacity from the slice, returning s[:len(s):len(s)].
|
||||
// TODO(adonovan): use go1.21 slices.Clip.
|
||||
func slicesClip[S ~[]E, E any](s S) S { return s[:len(s):len(s)] }
|
||||
|
|
|
@ -841,6 +841,7 @@ func (state *golistState) cfgInvocation() gocommand.Invocation {
|
|||
Env: cfg.Env,
|
||||
Logf: cfg.Logf,
|
||||
WorkingDir: cfg.Dir,
|
||||
Overlay: cfg.goListOverlayFile,
|
||||
}
|
||||
}
|
||||
|
||||
|
@ -849,26 +850,6 @@ func (state *golistState) invokeGo(verb string, args ...string) (*bytes.Buffer,
|
|||
cfg := state.cfg
|
||||
|
||||
inv := state.cfgInvocation()
|
||||
|
||||
// For Go versions 1.16 and above, `go list` accepts overlays directly via
|
||||
// the -overlay flag. Set it, if it's available.
|
||||
//
|
||||
// The check for "list" is not necessarily required, but we should avoid
|
||||
// getting the go version if possible.
|
||||
if verb == "list" {
|
||||
goVersion, err := state.getGoVersion()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if goVersion >= 16 {
|
||||
filename, cleanup, err := state.writeOverlays()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer cleanup()
|
||||
inv.Overlay = filename
|
||||
}
|
||||
}
|
||||
inv.Verb = verb
|
||||
inv.Args = args
|
||||
gocmdRunner := cfg.gocmdRunner
|
||||
|
@ -1015,67 +996,6 @@ func (state *golistState) invokeGo(verb string, args ...string) (*bytes.Buffer,
|
|||
return stdout, nil
|
||||
}
|
||||
|
||||
// OverlayJSON is the format overlay files are expected to be in.
|
||||
// The Replace map maps from overlaid paths to replacement paths:
|
||||
// the Go command will forward all reads trying to open
|
||||
// each overlaid path to its replacement path, or consider the overlaid
|
||||
// path not to exist if the replacement path is empty.
|
||||
//
|
||||
// From golang/go#39958.
|
||||
type OverlayJSON struct {
|
||||
Replace map[string]string `json:"replace,omitempty"`
|
||||
}
|
||||
|
||||
// writeOverlays writes out files for go list's -overlay flag, as described
|
||||
// above.
|
||||
func (state *golistState) writeOverlays() (filename string, cleanup func(), err error) {
|
||||
// Do nothing if there are no overlays in the config.
|
||||
if len(state.cfg.Overlay) == 0 {
|
||||
return "", func() {}, nil
|
||||
}
|
||||
dir, err := os.MkdirTemp("", "gopackages-*")
|
||||
if err != nil {
|
||||
return "", nil, err
|
||||
}
|
||||
// The caller must clean up this directory, unless this function returns an
|
||||
// error.
|
||||
cleanup = func() {
|
||||
os.RemoveAll(dir)
|
||||
}
|
||||
defer func() {
|
||||
if err != nil {
|
||||
cleanup()
|
||||
}
|
||||
}()
|
||||
overlays := map[string]string{}
|
||||
for k, v := range state.cfg.Overlay {
|
||||
// Create a unique filename for the overlaid files, to avoid
|
||||
// creating nested directories.
|
||||
noSeparator := strings.Join(strings.Split(filepath.ToSlash(k), "/"), "")
|
||||
f, err := os.CreateTemp(dir, fmt.Sprintf("*-%s", noSeparator))
|
||||
if err != nil {
|
||||
return "", func() {}, err
|
||||
}
|
||||
if _, err := f.Write(v); err != nil {
|
||||
return "", func() {}, err
|
||||
}
|
||||
if err := f.Close(); err != nil {
|
||||
return "", func() {}, err
|
||||
}
|
||||
overlays[k] = f.Name()
|
||||
}
|
||||
b, err := json.Marshal(OverlayJSON{Replace: overlays})
|
||||
if err != nil {
|
||||
return "", func() {}, err
|
||||
}
|
||||
// Write out the overlay file that contains the filepath mappings.
|
||||
filename = filepath.Join(dir, "overlay.json")
|
||||
if err := os.WriteFile(filename, b, 0665); err != nil {
|
||||
return "", func() {}, err
|
||||
}
|
||||
return filename, cleanup, nil
|
||||
}
|
||||
|
||||
func containsGoFile(s []string) bool {
|
||||
for _, f := range s {
|
||||
if strings.HasSuffix(f, ".go") {
|
||||
|
|
|
@ -37,10 +37,20 @@ import (
|
|||
// A LoadMode controls the amount of detail to return when loading.
|
||||
// The bits below can be combined to specify which fields should be
|
||||
// filled in the result packages.
|
||||
//
|
||||
// The zero value is a special case, equivalent to combining
|
||||
// the NeedName, NeedFiles, and NeedCompiledGoFiles bits.
|
||||
//
|
||||
// ID and Errors (if present) will always be filled.
|
||||
// Load may return more information than requested.
|
||||
// [Load] may return more information than requested.
|
||||
//
|
||||
// Unfortunately there are a number of open bugs related to
|
||||
// interactions among the LoadMode bits:
|
||||
// - https://github.com/golang/go/issues/48226
|
||||
// - https://github.com/golang/go/issues/56633
|
||||
// - https://github.com/golang/go/issues/56677
|
||||
// - https://github.com/golang/go/issues/58726
|
||||
// - https://github.com/golang/go/issues/63517
|
||||
type LoadMode int
|
||||
|
||||
const (
|
||||
|
@ -123,7 +133,14 @@ const (
|
|||
|
||||
// A Config specifies details about how packages should be loaded.
|
||||
// The zero value is a valid configuration.
|
||||
//
|
||||
// Calls to Load do not modify this struct.
|
||||
//
|
||||
// TODO(adonovan): #67702: this is currently false: in fact,
|
||||
// calls to [Load] do not modify the public fields of this struct, but
|
||||
// may modify hidden fields, so concurrent calls to [Load] must not
|
||||
// use the same Config. But perhaps we should reestablish the
|
||||
// documented invariant.
|
||||
type Config struct {
|
||||
// Mode controls the level of information returned for each package.
|
||||
Mode LoadMode
|
||||
|
@ -199,13 +216,23 @@ type Config struct {
|
|||
// setting Tests may have no effect.
|
||||
Tests bool
|
||||
|
||||
// Overlay provides a mapping of absolute file paths to file contents.
|
||||
// If the file with the given path already exists, the parser will use the
|
||||
// alternative file contents provided by the map.
|
||||
// Overlay is a mapping from absolute file paths to file contents.
|
||||
//
|
||||
// Overlays provide incomplete support for when a given file doesn't
|
||||
// already exist on disk. See the package doc above for more details.
|
||||
// For each map entry, [Load] uses the alternative file
|
||||
// contents provided by the overlay mapping instead of reading
|
||||
// from the file system. This mechanism can be used to enable
|
||||
// editor-integrated tools to correctly analyze the contents
|
||||
// of modified but unsaved buffers, for example.
|
||||
//
|
||||
// The overlay mapping is passed to the build system's driver
|
||||
// (see "The driver protocol") so that it too can report
|
||||
// consistent package metadata about unsaved files. However,
|
||||
// drivers may vary in their level of support for overlays.
|
||||
Overlay map[string][]byte
|
||||
|
||||
// goListOverlayFile is the JSON file that encodes the Overlay
|
||||
// mapping, used by 'go list -overlay=...'
|
||||
goListOverlayFile string
|
||||
}
|
||||
|
||||
// Load loads and returns the Go packages named by the given patterns.
|
||||
|
@ -213,6 +240,20 @@ type Config struct {
|
|||
// Config specifies loading options;
|
||||
// nil behaves the same as an empty Config.
|
||||
//
|
||||
// The [Config.Mode] field is a set of bits that determine what kinds
|
||||
// of information should be computed and returned. Modes that require
|
||||
// more information tend to be slower. See [LoadMode] for details
|
||||
// and important caveats. Its zero value is equivalent to
|
||||
// NeedName | NeedFiles | NeedCompiledGoFiles.
|
||||
//
|
||||
// Each call to Load returns a new set of [Package] instances.
|
||||
// The Packages and their Imports form a directed acyclic graph.
|
||||
//
|
||||
// If the [NeedTypes] mode flag was set, each call to Load uses a new
|
||||
// [types.Importer], so [types.Object] and [types.Type] values from
|
||||
// different calls to Load must not be mixed as they will have
|
||||
// inconsistent notions of type identity.
|
||||
//
|
||||
// If any of the patterns was invalid as defined by the
|
||||
// underlying build system, Load returns an error.
|
||||
// It may return an empty list of packages without an error,
|
||||
|
@ -286,6 +327,17 @@ func defaultDriver(cfg *Config, patterns ...string) (*DriverResponse, bool, erro
|
|||
// (fall through)
|
||||
}
|
||||
|
||||
// go list fallback
|
||||
//
|
||||
// Write overlays once, as there are many calls
|
||||
// to 'go list' (one per chunk plus others too).
|
||||
overlay, cleanupOverlay, err := gocommand.WriteOverlays(cfg.Overlay)
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
defer cleanupOverlay()
|
||||
cfg.goListOverlayFile = overlay
|
||||
|
||||
response, err := callDriverOnChunks(goListDriver, cfg, chunks)
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
|
@ -365,6 +417,9 @@ func mergeResponses(responses ...*DriverResponse) *DriverResponse {
|
|||
}
|
||||
|
||||
// A Package describes a loaded Go package.
|
||||
//
|
||||
// It also defines part of the JSON schema of [DriverResponse].
|
||||
// See the package documentation for an overview.
|
||||
type Package struct {
|
||||
// ID is a unique identifier for a package,
|
||||
// in a syntax provided by the underlying build system.
|
||||
|
@ -423,6 +478,13 @@ type Package struct {
|
|||
// to corresponding loaded Packages.
|
||||
Imports map[string]*Package
|
||||
|
||||
// Module is the module information for the package if it exists.
|
||||
//
|
||||
// Note: it may be missing for std and cmd; see Go issue #65816.
|
||||
Module *Module
|
||||
|
||||
// -- The following fields are not part of the driver JSON schema. --
|
||||
|
||||
// Types provides type information for the package.
|
||||
// The NeedTypes LoadMode bit sets this field for packages matching the
|
||||
// patterns; type information for dependencies may be missing or incomplete,
|
||||
|
@ -431,15 +493,15 @@ type Package struct {
|
|||
// Each call to [Load] returns a consistent set of type
|
||||
// symbols, as defined by the comment at [types.Identical].
|
||||
// Avoid mixing type information from two or more calls to [Load].
|
||||
Types *types.Package
|
||||
Types *types.Package `json:"-"`
|
||||
|
||||
// Fset provides position information for Types, TypesInfo, and Syntax.
|
||||
// It is set only when Types is set.
|
||||
Fset *token.FileSet
|
||||
Fset *token.FileSet `json:"-"`
|
||||
|
||||
// IllTyped indicates whether the package or any dependency contains errors.
|
||||
// It is set only when Types is set.
|
||||
IllTyped bool
|
||||
IllTyped bool `json:"-"`
|
||||
|
||||
// Syntax is the package's syntax trees, for the files listed in CompiledGoFiles.
|
||||
//
|
||||
|
@ -449,26 +511,28 @@ type Package struct {
|
|||
//
|
||||
// Syntax is kept in the same order as CompiledGoFiles, with the caveat that nils are
|
||||
// removed. If parsing returned nil, Syntax may be shorter than CompiledGoFiles.
|
||||
Syntax []*ast.File
|
||||
Syntax []*ast.File `json:"-"`
|
||||
|
||||
// TypesInfo provides type information about the package's syntax trees.
|
||||
// It is set only when Syntax is set.
|
||||
TypesInfo *types.Info
|
||||
TypesInfo *types.Info `json:"-"`
|
||||
|
||||
// TypesSizes provides the effective size function for types in TypesInfo.
|
||||
TypesSizes types.Sizes
|
||||
TypesSizes types.Sizes `json:"-"`
|
||||
|
||||
// -- internal --
|
||||
|
||||
// forTest is the package under test, if any.
|
||||
forTest string
|
||||
|
||||
// depsErrors is the DepsErrors field from the go list response, if any.
|
||||
depsErrors []*packagesinternal.PackageError
|
||||
|
||||
// module is the module information for the package if it exists.
|
||||
Module *Module
|
||||
}
|
||||
|
||||
// Module provides module information for a package.
|
||||
//
|
||||
// It also defines part of the JSON schema of [DriverResponse].
|
||||
// See the package documentation for an overview.
|
||||
type Module struct {
|
||||
Path string // module path
|
||||
Version string // module version
|
||||
|
@ -601,6 +665,7 @@ func (p *Package) UnmarshalJSON(b []byte) error {
|
|||
OtherFiles: flat.OtherFiles,
|
||||
EmbedFiles: flat.EmbedFiles,
|
||||
EmbedPatterns: flat.EmbedPatterns,
|
||||
IgnoredFiles: flat.IgnoredFiles,
|
||||
ExportFile: flat.ExportFile,
|
||||
}
|
||||
if len(flat.Imports) > 0 {
|
||||
|
|
|
@ -8,12 +8,14 @@ package gocommand
|
|||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"reflect"
|
||||
"regexp"
|
||||
"runtime"
|
||||
|
@ -167,7 +169,9 @@ type Invocation struct {
|
|||
// TODO(rfindley): remove, in favor of Args.
|
||||
ModFile string
|
||||
|
||||
// If Overlay is set, the go command is invoked with -overlay=Overlay.
|
||||
// Overlay is the name of the JSON overlay file that describes
|
||||
// unsaved editor buffers; see [WriteOverlays].
|
||||
// If set, the go command is invoked with -overlay=Overlay.
|
||||
// TODO(rfindley): remove, in favor of Args.
|
||||
Overlay string
|
||||
|
||||
|
@ -255,12 +259,15 @@ func (i *Invocation) run(ctx context.Context, stdout, stderr io.Writer) error {
|
|||
waitDelay.Set(reflect.ValueOf(30 * time.Second))
|
||||
}
|
||||
|
||||
// On darwin the cwd gets resolved to the real path, which breaks anything that
|
||||
// expects the working directory to keep the original path, including the
|
||||
// The cwd gets resolved to the real path. On Darwin, where
|
||||
// /tmp is a symlink, this breaks anything that expects the
|
||||
// working directory to keep the original path, including the
|
||||
// go command when dealing with modules.
|
||||
// The Go stdlib has a special feature where if the cwd and the PWD are the
|
||||
// same node then it trusts the PWD, so by setting it in the env for the child
|
||||
// process we fix up all the paths returned by the go command.
|
||||
//
|
||||
// os.Getwd has a special feature where if the cwd and the PWD
|
||||
// are the same node then it trusts the PWD, so by setting it
|
||||
// in the env for the child process we fix up all the paths
|
||||
// returned by the go command.
|
||||
if !i.CleanEnv {
|
||||
cmd.Env = os.Environ()
|
||||
}
|
||||
|
@ -351,6 +358,7 @@ func runCmdContext(ctx context.Context, cmd *exec.Cmd) (err error) {
|
|||
}
|
||||
}
|
||||
|
||||
startTime := time.Now()
|
||||
err = cmd.Start()
|
||||
if stdoutW != nil {
|
||||
// The child process has inherited the pipe file,
|
||||
|
@ -377,7 +385,7 @@ func runCmdContext(ctx context.Context, cmd *exec.Cmd) (err error) {
|
|||
case err := <-resChan:
|
||||
return err
|
||||
case <-timer.C:
|
||||
HandleHangingGoCommand(cmd.Process)
|
||||
HandleHangingGoCommand(startTime, cmd)
|
||||
case <-ctx.Done():
|
||||
}
|
||||
} else {
|
||||
|
@ -411,7 +419,7 @@ func runCmdContext(ctx context.Context, cmd *exec.Cmd) (err error) {
|
|||
return <-resChan
|
||||
}
|
||||
|
||||
func HandleHangingGoCommand(proc *os.Process) {
|
||||
func HandleHangingGoCommand(start time.Time, cmd *exec.Cmd) {
|
||||
switch runtime.GOOS {
|
||||
case "linux", "darwin", "freebsd", "netbsd":
|
||||
fmt.Fprintln(os.Stderr, `DETECTED A HANGING GO COMMAND
|
||||
|
@ -444,7 +452,7 @@ See golang/go#54461 for more details.`)
|
|||
panic(fmt.Sprintf("running %s: %v", listFiles, err))
|
||||
}
|
||||
}
|
||||
panic(fmt.Sprintf("detected hanging go command (pid %d): see golang/go#54461 for more details", proc.Pid))
|
||||
panic(fmt.Sprintf("detected hanging go command (golang/go#54461); waited %s\n\tcommand:%s\n\tpid:%d", time.Since(start), cmd, cmd.Process.Pid))
|
||||
}
|
||||
|
||||
func cmdDebugStr(cmd *exec.Cmd) string {
|
||||
|
@ -468,3 +476,73 @@ func cmdDebugStr(cmd *exec.Cmd) string {
|
|||
}
|
||||
return fmt.Sprintf("GOROOT=%v GOPATH=%v GO111MODULE=%v GOPROXY=%v PWD=%v %v", env["GOROOT"], env["GOPATH"], env["GO111MODULE"], env["GOPROXY"], env["PWD"], strings.Join(args, " "))
|
||||
}
|
||||
|
||||
// WriteOverlays writes each value in the overlay (see the Overlay
|
||||
// field of go/packages.Config) to a temporary file and returns the name
|
||||
// of a JSON file describing the mapping that is suitable for the "go
|
||||
// list -overlay" flag.
|
||||
//
|
||||
// On success, the caller must call the cleanup function exactly once
|
||||
// when the files are no longer needed.
|
||||
func WriteOverlays(overlay map[string][]byte) (filename string, cleanup func(), err error) {
|
||||
// Do nothing if there are no overlays in the config.
|
||||
if len(overlay) == 0 {
|
||||
return "", func() {}, nil
|
||||
}
|
||||
|
||||
dir, err := os.MkdirTemp("", "gocommand-*")
|
||||
if err != nil {
|
||||
return "", nil, err
|
||||
}
|
||||
|
||||
// The caller must clean up this directory,
|
||||
// unless this function returns an error.
|
||||
// (The cleanup operand of each return
|
||||
// statement below is ignored.)
|
||||
defer func() {
|
||||
cleanup = func() {
|
||||
os.RemoveAll(dir)
|
||||
}
|
||||
if err != nil {
|
||||
cleanup()
|
||||
cleanup = nil
|
||||
}
|
||||
}()
|
||||
|
||||
// Write each map entry to a temporary file.
|
||||
overlays := make(map[string]string)
|
||||
for k, v := range overlay {
|
||||
// Use a unique basename for each file (001-foo.go),
|
||||
// to avoid creating nested directories.
|
||||
base := fmt.Sprintf("%d-%s.go", 1+len(overlays), filepath.Base(k))
|
||||
filename := filepath.Join(dir, base)
|
||||
err := os.WriteFile(filename, v, 0666)
|
||||
if err != nil {
|
||||
return "", nil, err
|
||||
}
|
||||
overlays[k] = filename
|
||||
}
|
||||
|
||||
// Write the JSON overlay file that maps logical file names to temp files.
|
||||
//
|
||||
// OverlayJSON is the format overlay files are expected to be in.
|
||||
// The Replace map maps from overlaid paths to replacement paths:
|
||||
// the Go command will forward all reads trying to open
|
||||
// each overlaid path to its replacement path, or consider the overlaid
|
||||
// path not to exist if the replacement path is empty.
|
||||
//
|
||||
// From golang/go#39958.
|
||||
type OverlayJSON struct {
|
||||
Replace map[string]string `json:"replace,omitempty"`
|
||||
}
|
||||
b, err := json.Marshal(OverlayJSON{Replace: overlays})
|
||||
if err != nil {
|
||||
return "", nil, err
|
||||
}
|
||||
filename = filepath.Join(dir, "overlay.json")
|
||||
if err := os.WriteFile(filename, b, 0666); err != nil {
|
||||
return "", nil, err
|
||||
}
|
||||
|
||||
return filename, nil, nil
|
||||
}
|
||||
|
|
|
@ -104,7 +104,10 @@ type packageInfo struct {
|
|||
|
||||
// parseOtherFiles parses all the Go files in srcDir except filename, including
|
||||
// test files if filename looks like a test.
|
||||
func parseOtherFiles(fset *token.FileSet, srcDir, filename string) []*ast.File {
|
||||
//
|
||||
// It returns an error only if ctx is cancelled. Files with parse errors are
|
||||
// ignored.
|
||||
func parseOtherFiles(ctx context.Context, fset *token.FileSet, srcDir, filename string) ([]*ast.File, error) {
|
||||
// This could use go/packages but it doesn't buy much, and it fails
|
||||
// with https://golang.org/issue/26296 in LoadFiles mode in some cases.
|
||||
considerTests := strings.HasSuffix(filename, "_test.go")
|
||||
|
@ -112,11 +115,14 @@ func parseOtherFiles(fset *token.FileSet, srcDir, filename string) []*ast.File {
|
|||
fileBase := filepath.Base(filename)
|
||||
packageFileInfos, err := os.ReadDir(srcDir)
|
||||
if err != nil {
|
||||
return nil
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
|
||||
var files []*ast.File
|
||||
for _, fi := range packageFileInfos {
|
||||
if ctx.Err() != nil {
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
if fi.Name() == fileBase || !strings.HasSuffix(fi.Name(), ".go") {
|
||||
continue
|
||||
}
|
||||
|
@ -132,7 +138,7 @@ func parseOtherFiles(fset *token.FileSet, srcDir, filename string) []*ast.File {
|
|||
files = append(files, f)
|
||||
}
|
||||
|
||||
return files
|
||||
return files, ctx.Err()
|
||||
}
|
||||
|
||||
// addGlobals puts the names of package vars into the provided map.
|
||||
|
@ -557,12 +563,7 @@ func (p *pass) addCandidate(imp *ImportInfo, pkg *packageInfo) {
|
|||
|
||||
// fixImports adds and removes imports from f so that all its references are
|
||||
// satisfied and there are no unused imports.
|
||||
//
|
||||
// This is declared as a variable rather than a function so goimports can
|
||||
// easily be extended by adding a file with an init function.
|
||||
var fixImports = fixImportsDefault
|
||||
|
||||
func fixImportsDefault(fset *token.FileSet, f *ast.File, filename string, env *ProcessEnv) error {
|
||||
func fixImports(fset *token.FileSet, f *ast.File, filename string, env *ProcessEnv) error {
|
||||
fixes, err := getFixes(context.Background(), fset, f, filename, env)
|
||||
if err != nil {
|
||||
return err
|
||||
|
@ -592,7 +593,10 @@ func getFixes(ctx context.Context, fset *token.FileSet, f *ast.File, filename st
|
|||
return fixes, nil
|
||||
}
|
||||
|
||||
otherFiles := parseOtherFiles(fset, srcDir, filename)
|
||||
otherFiles, err := parseOtherFiles(ctx, fset, srcDir, filename)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Second pass: add information from other files in the same package,
|
||||
// like their package vars and imports.
|
||||
|
@ -1192,7 +1196,7 @@ func addExternalCandidates(ctx context.Context, pass *pass, refs references, fil
|
|||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err = resolver.scan(context.Background(), callback); err != nil {
|
||||
if err = resolver.scan(ctx, callback); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
|
@ -1203,7 +1207,7 @@ func addExternalCandidates(ctx context.Context, pass *pass, refs references, fil
|
|||
}
|
||||
results := make(chan result, len(refs))
|
||||
|
||||
ctx, cancel := context.WithCancel(context.TODO())
|
||||
ctx, cancel := context.WithCancel(ctx)
|
||||
var wg sync.WaitGroup
|
||||
defer func() {
|
||||
cancel()
|
||||
|
|
|
@ -12,7 +12,7 @@ import (
|
|||
"go/types"
|
||||
)
|
||||
|
||||
// FileVersions returns a file's Go version.
|
||||
// FileVersion returns a file's Go version.
|
||||
// The reported version is an unknown Future version if a
|
||||
// version cannot be determined.
|
||||
func FileVersion(info *types.Info, file *ast.File) string {
|
||||
|
|
|
@ -304,11 +304,12 @@ go.opentelemetry.io/proto/otlp/trace/v1
|
|||
go.uber.org/automaxprocs/internal/cgroups
|
||||
go.uber.org/automaxprocs/internal/runtime
|
||||
go.uber.org/automaxprocs/maxprocs
|
||||
# go.uber.org/mock v0.4.0
|
||||
## explicit; go 1.20
|
||||
# go.uber.org/mock v0.5.0
|
||||
## explicit; go 1.22
|
||||
go.uber.org/mock/gomock
|
||||
go.uber.org/mock/mockgen
|
||||
go.uber.org/mock/mockgen/model
|
||||
# golang.org/x/crypto v0.23.0
|
||||
# golang.org/x/crypto v0.24.0
|
||||
## explicit; go 1.18
|
||||
golang.org/x/crypto/blake2b
|
||||
golang.org/x/crypto/blowfish
|
||||
|
@ -328,13 +329,13 @@ golang.org/x/crypto/ssh/internal/bcrypt_pbkdf
|
|||
# golang.org/x/exp v0.0.0-20240506185415-9bf2ced13842
|
||||
## explicit; go 1.20
|
||||
golang.org/x/exp/rand
|
||||
# golang.org/x/mod v0.17.0
|
||||
# golang.org/x/mod v0.18.0
|
||||
## explicit; go 1.18
|
||||
golang.org/x/mod/internal/lazyregexp
|
||||
golang.org/x/mod/modfile
|
||||
golang.org/x/mod/module
|
||||
golang.org/x/mod/semver
|
||||
# golang.org/x/net v0.25.0
|
||||
# golang.org/x/net v0.26.0
|
||||
## explicit; go 1.18
|
||||
golang.org/x/net/bpf
|
||||
golang.org/x/net/context
|
||||
|
@ -360,7 +361,7 @@ golang.org/x/oauth2/internal
|
|||
# golang.org/x/sync v0.7.0
|
||||
## explicit; go 1.18
|
||||
golang.org/x/sync/errgroup
|
||||
# golang.org/x/sys v0.20.0
|
||||
# golang.org/x/sys v0.21.0
|
||||
## explicit; go 1.18
|
||||
golang.org/x/sys/cpu
|
||||
golang.org/x/sys/execabs
|
||||
|
@ -371,10 +372,10 @@ golang.org/x/sys/windows/registry
|
|||
golang.org/x/sys/windows/svc
|
||||
golang.org/x/sys/windows/svc/eventlog
|
||||
golang.org/x/sys/windows/svc/mgr
|
||||
# golang.org/x/term v0.20.0
|
||||
# golang.org/x/term v0.21.0
|
||||
## explicit; go 1.18
|
||||
golang.org/x/term
|
||||
# golang.org/x/text v0.15.0
|
||||
# golang.org/x/text v0.16.0
|
||||
## explicit; go 1.18
|
||||
golang.org/x/text/cases
|
||||
golang.org/x/text/internal
|
||||
|
@ -386,7 +387,7 @@ golang.org/x/text/secure/bidirule
|
|||
golang.org/x/text/transform
|
||||
golang.org/x/text/unicode/bidi
|
||||
golang.org/x/text/unicode/norm
|
||||
# golang.org/x/tools v0.21.0
|
||||
# golang.org/x/tools v0.22.0
|
||||
## explicit; go 1.19
|
||||
golang.org/x/tools/go/ast/astutil
|
||||
golang.org/x/tools/go/ast/inspector
|
||||
|
|
Loading…
Reference in New Issue