Skip to content

Commit b55d73e

Browse files
committed
feat(agent,pkg): add configuration to tunnel listener and yamux server
1 parent c89b3c1 commit b55d73e

File tree

2 files changed

+60
-9
lines changed

2 files changed

+60
-9
lines changed

agent/agent.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -631,7 +631,7 @@ func (a *Agent) Listen(ctx context.Context) error {
631631

632632
a.listening <- true
633633

634-
if err := a.tunnel.Listen(conn); err != nil {
634+
if err := a.tunnel.Listen(conn, &tunnel.DefaultConfig); err != nil {
635635
logger.WithError(err).Error("Tunnel listener exited with error")
636636
}
637637

agent/pkg/tunnel/tunnel.go

Lines changed: 59 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -37,15 +37,66 @@ func (t *Tunnel) Handle(protocol string, handler Handler) {
3737
// ErrTunnelDisconnect is returned when the tunnel connection is closed.
3838
var ErrTunnelDisconnect = errors.New("tunnel disconnected")
3939

40-
func (t *Tunnel) Listen(conn net.Conn) error {
40+
type Config struct {
41+
// AcceptBacklog is used to limit how many streams may be
42+
// waiting an accept.
43+
AcceptBacklog int `json:"accept_backlog"`
44+
45+
// EnableKeepalive is used to do a period keep alive
46+
// messages using a ping.
47+
EnableKeepAlive bool `json:"enable_keep_alive"`
48+
49+
// KeepAliveInterval is how often to perform the keep alive
50+
KeepAliveInterval time.Duration `json:"keep_alive_interval"`
51+
52+
// ConnectionWriteTimeout is meant to be a "safety valve" timeout after
53+
// we which will suspect a problem with the underlying connection and
54+
// close it. This is only applied to writes, where's there's generally
55+
// an expectation that things will move along quickly.
56+
ConnectionWriteTimeout time.Duration `json:"connection_write_timeout"`
57+
58+
// MaxStreamWindowSize is used to control the maximum
59+
// window size that we allow for a stream.
60+
MaxStreamWindowSize uint32 `json:"max_stream_window_size"`
61+
62+
// StreamOpenTimeout is the maximum amount of time that a stream will
63+
// be allowed to remain in pending state while waiting for an ack from the peer.
64+
// Once the timeout is reached the session will be gracefully closed.
65+
// A zero value disables the StreamOpenTimeout allowing unbounded
66+
// blocking on OpenStream calls.
67+
StreamOpenTimeout time.Duration `json:"stream_open_timeout"`
68+
69+
// StreamCloseTimeout is the maximum time that a stream will allowed to
70+
// be in a half-closed state when `Close` is called before forcibly
71+
// closing the connection. Forcibly closed connections will empty the
72+
// receive buffer, drop any future packets received for that stream,
73+
// and send a RST to the remote side.
74+
StreamCloseTimeout time.Duration `json:"stream_close_timeout"`
75+
}
76+
77+
var DefaultConfig = Config{
78+
AcceptBacklog: 256,
79+
EnableKeepAlive: true,
80+
KeepAliveInterval: 35 * time.Second,
81+
ConnectionWriteTimeout: 15 * time.Second,
82+
MaxStreamWindowSize: 256 * 1024,
83+
StreamCloseTimeout: 5 * time.Minute,
84+
StreamOpenTimeout: 75 * time.Second,
85+
}
86+
87+
func (t *Tunnel) Listen(conn net.Conn, cfg *Config) error {
88+
if cfg == nil {
89+
cfg = &DefaultConfig
90+
}
91+
4192
session, err := yamux.Server(conn, &yamux.Config{
42-
AcceptBacklog: 256,
43-
EnableKeepAlive: true,
44-
KeepAliveInterval: 35 * time.Second,
45-
ConnectionWriteTimeout: 15 * time.Second,
46-
MaxStreamWindowSize: 256 * 1024,
47-
StreamCloseTimeout: 5 * time.Minute,
48-
StreamOpenTimeout: 75 * time.Second,
93+
AcceptBacklog: cfg.AcceptBacklog,
94+
EnableKeepAlive: cfg.EnableKeepAlive,
95+
KeepAliveInterval: cfg.KeepAliveInterval,
96+
ConnectionWriteTimeout: cfg.ConnectionWriteTimeout,
97+
MaxStreamWindowSize: cfg.MaxStreamWindowSize,
98+
StreamCloseTimeout: cfg.StreamCloseTimeout,
99+
StreamOpenTimeout: cfg.StreamOpenTimeout,
49100
LogOutput: os.Stderr,
50101
})
51102
if err != nil {

0 commit comments

Comments
 (0)