Skip to content

Commit 64807e1

Browse files
EmyrkEmyrk
authored
chore: apply the 4mb max limit on drpc protocol message size (#17771)
Respect the 4mb max limit on proto messages
1 parent a1c03b6 commit 64807e1

File tree

10 files changed

+121
-10
lines changed

10 files changed

+121
-10
lines changed

‎agent/agenttest/client.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,7 @@ func NewClient(t testing.TB,
6060
err = agentproto.DRPCRegisterAgent(mux, fakeAAPI)
6161
require.NoError(t, err)
6262
server := drpcserver.NewWithOptions(mux, drpcserver.Options{
63+
Manager: drpcsdk.DefaultDRPCOptions(nil),
6364
Log: func(err error) {
6465
if xerrors.Is(err, io.EOF) {
6566
return

‎coderd/agentapi/api.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ import (
3030
"github.com/coder/coder/v2/coderd/wspubsub"
3131
"github.com/coder/coder/v2/codersdk"
3232
"github.com/coder/coder/v2/codersdk/agentsdk"
33+
"github.com/coder/coder/v2/codersdk/drpcsdk"
3334
"github.com/coder/coder/v2/tailnet"
3435
tailnetproto "github.com/coder/coder/v2/tailnet/proto"
3536
"github.com/coder/quartz"
@@ -209,6 +210,7 @@ func (a *API) Server(ctx context.Context) (*drpcserver.Server, error) {
209210

210211
return drpcserver.NewWithOptions(&tracing.DRPCHandler{Handler: mux},
211212
drpcserver.Options{
213+
Manager: drpcsdk.DefaultDRPCOptions(nil),
212214
Log: func(err error) {
213215
if xerrors.Is(err, io.EOF) {
214216
return

‎coderd/coderd.go

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ import (
3838
"tailscale.com/util/singleflight"
3939

4040
"cdr.dev/slog"
41+
"github.com/coder/coder/v2/codersdk/drpcsdk"
4142
"github.com/coder/quartz"
4243
"github.com/coder/serpent"
4344

@@ -84,7 +85,6 @@ import (
8485
"github.com/coder/coder/v2/coderd/workspaceapps"
8586
"github.com/coder/coder/v2/coderd/workspacestats"
8687
"github.com/coder/coder/v2/codersdk"
87-
"github.com/coder/coder/v2/codersdk/drpcsdk"
8888
"github.com/coder/coder/v2/codersdk/healthsdk"
8989
"github.com/coder/coder/v2/provisionerd/proto"
9090
"github.com/coder/coder/v2/provisionersdk"
@@ -1803,6 +1803,7 @@ func (api *API) CreateInMemoryTaggedProvisionerDaemon(dialCtx context.Context, n
18031803
}
18041804
server := drpcserver.NewWithOptions(&tracing.DRPCHandler{Handler: mux},
18051805
drpcserver.Options{
1806+
Manager: drpcsdk.DefaultDRPCOptions(nil),
18061807
Log: func(err error) {
18071808
if xerrors.Is(err, io.EOF) {
18081809
return

‎codersdk/drpcsdk/transport.go

Lines changed: 24 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import (
99
"github.com/valyala/fasthttp/fasthttputil"
1010
"storj.io/drpc"
1111
"storj.io/drpc/drpcconn"
12+
"storj.io/drpc/drpcmanager"
1213

1314
"github.com/coder/coder/v2/coderd/tracing"
1415
)
@@ -19,6 +20,17 @@ const (
1920
MaxMessageSize = 4 << 20
2021
)
2122

23+
func DefaultDRPCOptions(options *drpcmanager.Options) drpcmanager.Options {
24+
if options == nil {
25+
options = &drpcmanager.Options{}
26+
}
27+
28+
if options.Reader.MaximumBufferSize == 0 {
29+
options.Reader.MaximumBufferSize = MaxMessageSize
30+
}
31+
return *options
32+
}
33+
2234
// MultiplexedConn returns a multiplexed dRPC connection from a yamux Session.
2335
func MultiplexedConn(session *yamux.Session) drpc.Conn {
2436
return &multiplexedDRPC{session}
@@ -43,7 +55,9 @@ func (m *multiplexedDRPC) Invoke(ctx context.Context, rpc string, enc drpc.Encod
4355
if err != nil {
4456
return err
4557
}
46-
dConn := drpcconn.New(conn)
58+
dConn := drpcconn.NewWithOptions(conn, drpcconn.Options{
59+
Manager: DefaultDRPCOptions(nil),
60+
})
4761
defer func() {
4862
_ = dConn.Close()
4963
}()
@@ -55,7 +69,9 @@ func (m *multiplexedDRPC) NewStream(ctx context.Context, rpc string, enc drpc.En
5569
if err != nil {
5670
return nil, err
5771
}
58-
dConn := drpcconn.New(conn)
72+
dConn := drpcconn.NewWithOptions(conn, drpcconn.Options{
73+
Manager: DefaultDRPCOptions(nil),
74+
})
5975
stream, err := dConn.NewStream(ctx, rpc, enc)
6076
if err == nil {
6177
go func() {
@@ -97,7 +113,9 @@ func (m *memDRPC) Invoke(ctx context.Context, rpc string, enc drpc.Encoding, inM
97113
return err
98114
}
99115

100-
dConn := &tracing.DRPCConn{Conn: drpcconn.New(conn)}
116+
dConn := &tracing.DRPCConn{Conn: drpcconn.NewWithOptions(conn, drpcconn.Options{
117+
Manager: DefaultDRPCOptions(nil),
118+
})}
101119
defer func() {
102120
_ = dConn.Close()
103121
_ = conn.Close()
@@ -110,7 +128,9 @@ func (m *memDRPC) NewStream(ctx context.Context, rpc string, enc drpc.Encoding)
110128
if err != nil {
111129
return nil, err
112130
}
113-
dConn := &tracing.DRPCConn{Conn: drpcconn.New(conn)}
131+
dConn := &tracing.DRPCConn{Conn: drpcconn.NewWithOptions(conn, drpcconn.Options{
132+
Manager: DefaultDRPCOptions(nil),
133+
})}
114134
stream, err := dConn.NewStream(ctx, rpc, enc)
115135
if err != nil {
116136
_ = dConn.Close()

‎enterprise/coderd/provisionerdaemons.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ import (
3131
"github.com/coder/coder/v2/coderd/telemetry"
3232
"github.com/coder/coder/v2/coderd/util/ptr"
3333
"github.com/coder/coder/v2/codersdk"
34+
"github.com/coder/coder/v2/codersdk/drpcsdk"
3435
"github.com/coder/coder/v2/provisionerd/proto"
3536
"github.com/coder/coder/v2/provisionersdk"
3637
"github.com/coder/websocket"
@@ -370,6 +371,7 @@ func (api *API) provisionerDaemonServe(rw http.ResponseWriter, r *http.Request)
370371
return
371372
}
372373
server := drpcserver.NewWithOptions(mux, drpcserver.Options{
374+
Manager: drpcsdk.DefaultDRPCOptions(nil),
373375
Log: func(err error) {
374376
if xerrors.Is(err, io.EOF) {
375377
return

‎enterprise/provisionerd/remoteprovisioners.go

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ import (
2727

2828
"cdr.dev/slog"
2929
"github.com/coder/coder/v2/coderd/database"
30+
"github.com/coder/coder/v2/codersdk/drpcsdk"
3031
"github.com/coder/coder/v2/provisioner/echo"
3132
agpl "github.com/coder/coder/v2/provisionerd"
3233
"github.com/coder/coder/v2/provisionerd/proto"
@@ -188,8 +189,10 @@ func (r *remoteConnector) handleConn(conn net.Conn) {
188189
logger.Info(r.ctx, "provisioner connected")
189190
closeConn = false // we're passing the conn over the channel
190191
w.respCh <- agpl.ConnectResponse{
191-
Job: w.job,
192-
Client: sdkproto.NewDRPCProvisionerClient(drpcconn.New(tlsConn)),
192+
Job: w.job,
193+
Client: sdkproto.NewDRPCProvisionerClient(drpcconn.NewWithOptions(tlsConn, drpcconn.Options{
194+
Manager: drpcsdk.DefaultDRPCOptions(nil),
195+
})),
193196
}
194197
}
195198

‎provisionerd/provisionerd_test.go

Lines changed: 76 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -178,6 +178,79 @@ func TestProvisionerd(t *testing.T) {
178178
require.NoError(t, closer.Close())
179179
})
180180

181+
// LargePayloads sends a 3mb tar file to the provisioner. The provisioner also
182+
// returns large payload messages back. The limit should be 4mb, so all
183+
// these messages should work.
184+
t.Run("LargePayloads", func(t *testing.T) {
185+
t.Parallel()
186+
done := make(chan struct{})
187+
t.Cleanup(func() {
188+
close(done)
189+
})
190+
var (
191+
largeSize = 3 * 1024 * 1024
192+
completeChan = make(chan struct{})
193+
completeOnce sync.Once
194+
acq = newAcquireOne(t, &proto.AcquiredJob{
195+
JobId: "test",
196+
Provisioner: "someprovisioner",
197+
TemplateSourceArchive: testutil.CreateTar(t, map[string]string{
198+
"toolarge.txt": string(make([]byte, largeSize)),
199+
}),
200+
Type: &proto.AcquiredJob_TemplateImport_{
201+
TemplateImport: &proto.AcquiredJob_TemplateImport{
202+
Metadata: &sdkproto.Metadata{},
203+
},
204+
},
205+
})
206+
)
207+
208+
closer := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
209+
return createProvisionerDaemonClient(t, done, provisionerDaemonTestServer{
210+
acquireJobWithCancel: acq.acquireWithCancel,
211+
updateJob: noopUpdateJob,
212+
completeJob: func(ctx context.Context, job *proto.CompletedJob) (*proto.Empty, error) {
213+
completeOnce.Do(func() { close(completeChan) })
214+
return &proto.Empty{}, nil
215+
},
216+
}), nil
217+
}, provisionerd.LocalProvisioners{
218+
"someprovisioner": createProvisionerClient(t, done, provisionerTestServer{
219+
parse: func(
220+
s *provisionersdk.Session,
221+
_ *sdkproto.ParseRequest,
222+
cancelOrComplete <-chan struct{},
223+
) *sdkproto.ParseComplete {
224+
return &sdkproto.ParseComplete{
225+
// 6mb readme
226+
Readme: make([]byte, largeSize),
227+
}
228+
},
229+
plan: func(
230+
_ *provisionersdk.Session,
231+
_ *sdkproto.PlanRequest,
232+
_ <-chan struct{},
233+
) *sdkproto.PlanComplete {
234+
return &sdkproto.PlanComplete{
235+
Resources: []*sdkproto.Resource{},
236+
Plan: make([]byte, largeSize),
237+
}
238+
},
239+
apply: func(
240+
_ *provisionersdk.Session,
241+
_ *sdkproto.ApplyRequest,
242+
_ <-chan struct{},
243+
) *sdkproto.ApplyComplete {
244+
return &sdkproto.ApplyComplete{
245+
State: make([]byte, largeSize),
246+
}
247+
},
248+
}),
249+
})
250+
require.Condition(t, closedWithin(completeChan, testutil.WaitShort))
251+
require.NoError(t, closer.Close())
252+
})
253+
181254
t.Run("RunningPeriodicUpdate", func(t *testing.T) {
182255
t.Parallel()
183256
done := make(chan struct{})
@@ -1115,7 +1188,9 @@ func createProvisionerDaemonClient(t *testing.T, done <-chan struct{}, server pr
11151188
mux := drpcmux.New()
11161189
err := proto.DRPCRegisterProvisionerDaemon(mux, &server)
11171190
require.NoError(t, err)
1118-
srv := drpcserver.New(mux)
1191+
srv := drpcserver.NewWithOptions(mux, drpcserver.Options{
1192+
Manager: drpcsdk.DefaultDRPCOptions(nil),
1193+
})
11191194
ctx, cancelFunc := context.WithCancel(context.Background())
11201195
closed := make(chan struct{})
11211196
go func() {

‎provisionersdk/serve.go

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import (
1515
"storj.io/drpc/drpcserver"
1616

1717
"cdr.dev/slog"
18+
"github.com/coder/coder/v2/codersdk/drpcsdk"
1819

1920
"github.com/coder/coder/v2/coderd/tracing"
2021
"github.com/coder/coder/v2/provisionersdk/proto"
@@ -81,7 +82,9 @@ func Serve(ctx context.Context, server Server, options *ServeOptions) error {
8182
if err != nil {
8283
return xerrors.Errorf("register provisioner: %w", err)
8384
}
84-
srv := drpcserver.New(&tracing.DRPCHandler{Handler: mux})
85+
srv := drpcserver.NewWithOptions(&tracing.DRPCHandler{Handler: mux}, drpcserver.Options{
86+
Manager: drpcsdk.DefaultDRPCOptions(nil),
87+
})
8588

8689
if options.Listener != nil {
8790
err = srv.Serve(ctx, options.Listener)

‎provisionersdk/serve_test.go

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -94,7 +94,9 @@ func TestProvisionerSDK(t *testing.T) {
9494
srvErr <- err
9595
}()
9696

97-
api := proto.NewDRPCProvisionerClient(drpcconn.New(client))
97+
api := proto.NewDRPCProvisionerClient(drpcconn.NewWithOptions(client, drpcconn.Options{
98+
Manager: drpcsdk.DefaultDRPCOptions(nil),
99+
}))
98100
s, err := api.Session(ctx)
99101
require.NoError(t, err)
100102
err = s.Send(&proto.Request{Type: &proto.Request_Config{Config: &proto.Config{}}})

‎tailnet/service.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import (
1717

1818
"cdr.dev/slog"
1919
"github.com/coder/coder/v2/apiversion"
20+
"github.com/coder/coder/v2/codersdk/drpcsdk"
2021
"github.com/coder/coder/v2/tailnet/proto"
2122
"github.com/coder/quartz"
2223
)
@@ -92,6 +93,7 @@ func NewClientService(options ClientServiceOptions) (
9293
return nil, xerrors.Errorf("register DRPC service: %w", err)
9394
}
9495
server := drpcserver.NewWithOptions(mux, drpcserver.Options{
96+
Manager: drpcsdk.DefaultDRPCOptions(nil),
9597
Log: func(err error) {
9698
if xerrors.Is(err, io.EOF) ||
9799
xerrors.Is(err, context.Canceled) ||

0 commit comments

Comments
 (0)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy