From 73bb720ec6b2f2c54b8aea2131488cf74aee91c9 Mon Sep 17 00:00:00 2001 From: iRedTea Date: Wed, 14 Oct 2026 21:04:13 +0300 Subject: [PATCH] feat: agents connect by WebSocket, falling back to long polling and trying again each minute; auto, ws or lp per machine (0.1.0) --- VERSION.txt | 2 +- deploy/nginx.conf.example | 18 ++- go.mod | 11 ++ go.sum | 60 ++++++- internal/agent/client.go | 104 +++++++++++- internal/agent/session.go | 33 +++- internal/agent/wslink.go | 172 ++++++++++++++++++++ internal/proto/proto.go | 26 +++ internal/server/db.go | 11 +- internal/server/handlers_agent.go | 7 +- internal/server/handlers_panel.go | 18 +++ internal/server/hub.go | 12 ++ internal/server/server.go | 3 + internal/server/web/templates/edit.html | 19 +++ internal/server/wsagent.go | 202 ++++++++++++++++++++++++ 15 files changed, 676 insertions(+), 22 deletions(-) create mode 100644 internal/agent/wslink.go create mode 100644 internal/server/wsagent.go diff --git a/VERSION.txt b/VERSION.txt index d169b2f..6e8bf73 100644 --- a/VERSION.txt +++ b/VERSION.txt @@ -1 +1 @@ -0.0.8 +0.1.0 diff --git a/deploy/nginx.conf.example b/deploy/nginx.conf.example index e929012..0891fe9 100644 --- a/deploy/nginx.conf.example +++ b/deploy/nginx.conf.example @@ -4,14 +4,21 @@ # Put this in /etc/nginx/sites-available/lp_agent and symlink it into # sites-enabled, then: sudo nginx -t && sudo systemctl reload nginx # -# Everything (agent + browser terminal) is plain HTTP long-polling, so NO -# WebSocket config is needed. The two settings that matter: -# - proxy_buffering off -> deliver long-poll responses immediately -# - proxy_read_timeout 120s -> longer than the 25s long-poll hold +# Agents connect by WebSocket (from 0.1.0), with HTTP long polling to fall +# back on; the browser uses plain HTTP. The settings that matter: +# - Upgrade / Connection headers -> let the agents' WebSocket through +# - proxy_buffering off -> deliver long-poll responses immediately +# - proxy_read_timeout 300s -> longer than the 25s long-poll hold # # NOTE: proxy_pass MUST start with "http://". A bare "proxy_pass 127.0.0.1:8091;" # causes nginx error: [emerg] invalid URL prefix. +# "Connection: upgrade" for a WebSocket, "close" for everything else. +map $http_upgrade $connection_upgrade { + default upgrade; + '' close; +} + server { listen 8445 ssl; server_name REDACTED; @@ -33,6 +40,9 @@ server { proxy_set_header X-Real-IP $remote_addr; proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for; proxy_set_header X-Forwarded-Proto $scheme; + # Agents' WebSocket. + proxy_set_header Upgrade $http_upgrade; + proxy_set_header Connection $connection_upgrade; # Long-poll behaviour: no buffering, generous read timeout. proxy_buffering off; diff --git a/go.mod b/go.mod index 97d7dd6..2d22caa 100644 --- a/go.mod +++ b/go.mod @@ -3,6 +3,7 @@ module github.com/x-grk/longpolling-remote-access go 1.26 require ( + github.com/coder/websocket v1.8.12 github.com/creack/pty v1.1.21 golang.org/x/crypto v0.31.0 golang.org/x/sys v0.28.0 @@ -12,9 +13,19 @@ require ( require ( github.com/dustin/go-humanize v1.0.1 // indirect github.com/google/uuid v1.6.0 // indirect + github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51 // indirect github.com/mattn/go-isatty v0.0.16 // indirect github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect + golang.org/x/mod v0.3.0 // indirect + golang.org/x/tools v0.0.0-20201124115921-2c860bdd6e78 // indirect + golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1 // indirect + lukechampine.com/uint128 v1.2.0 // indirect + modernc.org/cc/v3 v3.40.0 // indirect + modernc.org/ccgo/v3 v3.16.13 // indirect modernc.org/libc v1.22.5 // indirect modernc.org/mathutil v1.5.0 // indirect modernc.org/memory v1.5.0 // indirect + modernc.org/opt v0.1.3 // indirect + modernc.org/strutil v1.1.3 // indirect + modernc.org/token v1.0.1 // indirect ) diff --git a/go.sum b/go.sum index dbcce56..e75b7a1 100644 --- a/go.sum +++ b/go.sum @@ -1,28 +1,80 @@ +github.com/coder/websocket v1.8.12 h1:5bUXkEPPIbewrnkU8LTCLVaxi4N4J8ahufH2vlo4NAo= +github.com/coder/websocket v1.8.12/go.mod h1:LNVeNrXQZfe5qhS9ALED3uA+l5pPqvwXg3CKoDBB2gs= github.com/creack/pty v1.1.21 h1:1/QdRyBaHHJP61QkWMXlOIBfsgdDeeKfK8SYVUWJKf0= github.com/creack/pty v1.1.21/go.mod h1:MOBLtS5ELjhRRrroQr9kyvTxUAFNvYEK993ew/Vr4O4= github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= +github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38= +github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= +github.com/google/pprof v0.0.0-20221118152302-e6195bd50e26 h1:Xim43kblpZXfIBQsbuBVKCudVG457BR2GZFIz3uw3hQ= +github.com/google/pprof v0.0.0-20221118152302-e6195bd50e26/go.mod h1:dDKJzRmX4S37WGHujM7tX//fmj1uioxKzKxz3lo4HJo= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= -github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg= -github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= +github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51 h1:Z9n2FFNUXsshfwJMBgNA0RU6/i7WVaAegv3PtuIHPMs= +github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51/go.mod h1:CzGEWj7cYgsdH8dAjBGEr58BoE7ScuLd+fwFZ44+/x8= github.com/mattn/go-isatty v0.0.16 h1:bq3VjFmv/sOjHtdEhmkEV4x1AJtvUvOJ2PFAZ5+peKQ= github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM= +github.com/mattn/go-sqlite3 v1.14.16 h1:yOQRA0RpS5PFz/oikGwBEqvAWhWg5ufRz4ETLjwpU1Y= +github.com/mattn/go-sqlite3 v1.14.16/go.mod h1:2eHXhiwb8IkHr+BDWZGa96P6+rkvnG63S2DGjv9HUNg= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/remyoudompheng/bigfft v0.0.0-20200410134404-eec4a21b6bb0/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo= github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE= github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo= -golang.org/x/crypto v0.28.0 h1:GBDwsMXVQi34v5CCYUm2jkJvu4cbtru2U4TN2PSyQnw= -golang.org/x/crypto v0.28.0/go.mod h1:rmgy+3RHxRZMyY0jjAJShp2zgEdOqj2AO7U0pYmeQ7U= +github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= +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-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto= golang.org/x/crypto v0.31.0 h1:ihbySMvVjLAeSH1IbfcRTkD/iNscyz8rGzjF/E5hV6U= golang.org/x/crypto v0.31.0/go.mod h1:kDsLvtWBEx7MV9tJOj9bnXsPbxwJQ6csT/x4KIN4Ssk= +golang.org/x/mod v0.3.0 h1:RM4zey1++hCTbCVQfnWeKs9/IEsaBLA8vTkd0WVtmH4= +golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= +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-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= +golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= +golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.28.0 h1:Fksou7UEQUWlKvIdsqzJmUmCX3cZuD2+P3XyyzwMhlA= golang.org/x/sys v0.28.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= +golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= +golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= +golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= +golang.org/x/tools v0.0.0-20201124115921-2c860bdd6e78 h1:M8tBwCtWD/cZV9DZpFYRUgaymAYAr+aIUTWzDaM3uPs= +golang.org/x/tools v0.0.0-20201124115921-2c860bdd6e78/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA= +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-20200804184101-5ec99f83aff1 h1:go1bK/D/BFZV2I8cIQd1NKEZ+0owSTG1fDTci4IqFcE= +golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +lukechampine.com/uint128 v1.2.0 h1:mBi/5l91vocEN8otkC5bDLhi2KdCticRiwbdB0O+rjI= +lukechampine.com/uint128 v1.2.0/go.mod h1:c4eWIwlEGaxC/+H1VguhU4PHXNWDCDMUlWdIWl2j1gk= +modernc.org/cc/v3 v3.40.0 h1:P3g79IUS/93SYhtoeaHW+kRCIrYaxJ27MFPv+7kaTOw= +modernc.org/cc/v3 v3.40.0/go.mod h1:/bTg4dnWkSXowUO6ssQKnOV0yMVxDYNIsIrzqTFDGH0= +modernc.org/ccgo/v3 v3.16.13 h1:Mkgdzl46i5F/CNR/Kj80Ri59hC8TKAhZrYSaqvkwzUw= +modernc.org/ccgo/v3 v3.16.13/go.mod h1:2Quk+5YgpImhPjv2Qsob1DnZ/4som1lJTodubIcoUkY= +modernc.org/ccorpus v1.11.6 h1:J16RXiiqiCgua6+ZvQot4yUuUy8zxgqbqEEUuGPlISk= +modernc.org/ccorpus v1.11.6/go.mod h1:2gEUTrWqdpH2pXsmTM1ZkjeSrUWDpjMu2T6m29L/ErQ= +modernc.org/httpfs v1.0.6 h1:AAgIpFZRXuYnkjftxTAZwMIiwEqAfk8aVB2/oA6nAeM= +modernc.org/httpfs v1.0.6/go.mod h1:7dosgurJGp0sPaRanU53W4xZYKh14wfzX420oZADeHM= modernc.org/libc v1.22.5 h1:91BNch/e5B0uPbJFgqbxXuOnxBQjlS//icfQEGmvyjE= modernc.org/libc v1.22.5/go.mod h1:jj+Z7dTNX8fBScMVNRAYZ/jF91K8fdT2hYMThc3YjBY= modernc.org/mathutil v1.5.0 h1:rV0Ko/6SfM+8G+yKiyI830l3Wuz1zRutdslNoQ0kfiQ= modernc.org/mathutil v1.5.0/go.mod h1:mZW8CKdRPY1v87qxC/wUdX5O1qDzXMP5TH3wjfpga6E= modernc.org/memory v1.5.0 h1:N+/8c5rE6EqugZwHii4IFsaJ7MUhoWX07J5tC/iI5Ds= modernc.org/memory v1.5.0/go.mod h1:PkUhL0Mugw21sHPeskwZW4D6VscE/GQJOnIpCnW6pSU= +modernc.org/opt v0.1.3 h1:3XOZf2yznlhC+ibLltsDGzABUGVx8J6pnFMS3E4dcq4= +modernc.org/opt v0.1.3/go.mod h1:WdSiB5evDcignE70guQKxYUl14mgWtbClRi5wmkkTX0= modernc.org/sqlite v1.24.0 h1:EsClRIWHGhLTCX44p+Ri/JLD+vFGo0QGjasg2/F9TlI= modernc.org/sqlite v1.24.0/go.mod h1:OrDj17Mggn6MhE+iPbBNf7RGKODDE9NFT0f3EwDzJqk= +modernc.org/strutil v1.1.3 h1:fNMm+oJklMGYfU9Ylcywl0CO5O6nTfaowNsh2wpPjzY= +modernc.org/strutil v1.1.3/go.mod h1:MEHNA7PdEnEwLvspRMtWTNnp2nnyvMfkimT1NKNAGbw= +modernc.org/tcl v1.15.2 h1:C4ybAYCGJw968e+Me18oW55kD/FexcHbqH2xak1ROSY= +modernc.org/tcl v1.15.2/go.mod h1:3+k/ZaEbKrC8ePv8zJWPtBSW0V7Gg9g8rkmhI1Kfs3c= +modernc.org/token v1.0.1 h1:A3qvTqOwexpfZZeyI0FeGPDlSWX5pjZu9hF4lU+EKWg= +modernc.org/token v1.0.1/go.mod h1:UGzOrNV1mAFSEB63lOFHIpNRUVMvYTc6yu1SMY/XTDM= +modernc.org/z v1.7.3 h1:zDJf6iHjrnB+WRD88stbXokugjyc0/pB91ri1gO6LZY= +modernc.org/z v1.7.3/go.mod h1:Ipv4tsdxZRbQyLq9Q1M6gdbkxYzdlrciF2Hi/lS7nWE= diff --git a/internal/agent/client.go b/internal/agent/client.go index f88dc34..a80d946 100644 --- a/internal/agent/client.go +++ b/internal/agent/client.go @@ -18,6 +18,7 @@ import ( "path/filepath" "runtime" "sync" + "sync/atomic" "time" lpra "github.com/x-grk/longpolling-remote-access" @@ -53,6 +54,8 @@ type Agent struct { pollCount int // Direct tunnel currently running (see setTunnelAddr). + transport atomic.Value // string: auto, ws or lp, as the server says + link atomic.Pointer[wsLink] tunMu sync.Mutex tunAddr string tunCancel context.CancelFunc @@ -107,10 +110,22 @@ func (a *Agent) Run(ctx context.Context) error { a.setTunnelAddr(ctx, a.cfg.TunnelAddr) } + // The control channel: a WebSocket, or long polls. In "auto" a failed or + // dropped WebSocket gives way to long polling for a minute, then is tried + // again; "ws" keeps trying the WebSocket; "lp" only polls. backoff := time.Second for ctx.Err() == nil { - if err := a.pollOnce(ctx); err != nil { - log.Printf("poll: %v (retry in %s)", err, backoff) + if a.mode() == proto.TransportLP { + a.longPoll(ctx, 0) + continue + } + err := a.runWS(ctx) + if ctx.Err() != nil { + return nil + } + switch a.mode() { + case proto.TransportWS: + log.Printf("websocket: %v (retry in %s)", err, backoff) select { case <-time.After(backoff): case <-ctx.Done(): @@ -120,10 +135,71 @@ func (a *Agent) Run(ctx context.Context) error { backoff *= 2 } continue + case proto.TransportAuto: + log.Printf("websocket: %v - long polling, websocket again in a minute", err) + } + backoff = time.Second + a.longPoll(ctx, time.Minute) + } + return nil +} + +// mode is the transport the server set for this agent. +func (a *Agent) mode() string { + if v, ok := a.transport.Load().(string); ok { + return v + } + return proto.TransportAuto +} + +func (a *Agent) applyTransport(v string) { + switch v { + case proto.TransportAuto, proto.TransportWS, proto.TransportLP: + if v != a.mode() { + log.Printf("transport: %s", v) + a.transport.Store(v) + } + } +} + +// longPoll runs control polls: for d, or with d zero until the transport is +// no longer long polling. +func (a *Agent) longPoll(ctx context.Context, d time.Duration) { + until := time.Now().Add(d) + backoff := time.Second + for ctx.Err() == nil { + if d > 0 && (time.Now().After(until) || a.mode() == proto.TransportWS) { + return + } + if d == 0 && a.mode() != proto.TransportLP { + return + } + // In "auto" the minute is a minute: a poll still held then is cut + // short (the server keeps its job for the next one). + reqCtx, cancel := ctx, context.CancelFunc(func() {}) + if d > 0 { + reqCtx, cancel = context.WithDeadline(ctx, until) + } + err := a.pollOnce(ctx, reqCtx) + cut := reqCtx.Err() == context.DeadlineExceeded + cancel() + if err != nil && cut && ctx.Err() == nil { + return + } + if err != nil { + log.Printf("poll: %v (retry in %s)", err, backoff) + select { + case <-time.After(backoff): + case <-ctx.Done(): + return + } + if backoff < 15*time.Second { + backoff *= 2 + } + continue } backoff = time.Second } - return nil } func (a *Agent) enroll(ctx context.Context) error { @@ -157,8 +233,10 @@ func (a *Agent) enroll(ctx context.Context) error { } // pollOnce performs one control long-poll and dispatches any instruction. -func (a *Agent) pollOnce(ctx context.Context) error { - req, _ := http.NewRequestWithContext(ctx, http.MethodGet, a.cfg.Host+"/api/agent/poll", nil) +// The poll itself runs on reqCtx, which may end early; what it starts runs on +// ctx. +func (a *Agent) pollOnce(ctx, reqCtx context.Context) error { + req, _ := http.NewRequestWithContext(reqCtx, http.MethodGet, a.cfg.Host+"/api/agent/poll", nil) a.authHeaders(req) // Report the version and our OS user occasionally (not every poll), so the // panel can flag outdated agents without extra traffic. @@ -183,7 +261,14 @@ func (a *Agent) pollOnce(ctx context.Context) error { return err } a.applyTunnelHeader(ctx, resp.Header) - // Dispatch in the background so we keep polling for more work. + a.applyTransport(resp.Header.Get(proto.TransportHeader)) + a.dispatch(ctx, msg) + return nil +} + +// dispatch starts what a control message asks for, in the background, so the +// control channel keeps listening. +func (a *Agent) dispatch(ctx context.Context, msg proto.ControlMessage) { switch msg.Type { case proto.CtrlOpenTerminal: go a.runTerminal(ctx, msg.SessionID, msg.User) @@ -192,7 +277,6 @@ func (a *Agent) pollOnce(ctx context.Context) error { case proto.CtrlTask: go a.runTask(ctx, msg.TaskID) } - return nil } // fetchConfig asks the server for the settings it hands to agents and reports @@ -207,13 +291,17 @@ func (a *Agent) fetchConfig(ctx context.Context) bool { return false } defer resp.Body.Close() + a.applyTransport(resp.Header.Get(proto.TransportHeader)) return a.applyTunnelHeader(ctx, resp.Header) } // applyTunnelHeader follows the tunnel address the server advertises as // "host:port"; an empty host means the host we already reach the server on. func (a *Agent) applyTunnelHeader(ctx context.Context, h http.Header) bool { - v := h.Get("X-Tunnel-Addr") + return a.applyTunnelAddr(ctx, h.Get("X-Tunnel-Addr")) +} + +func (a *Agent) applyTunnelAddr(ctx context.Context, v string) bool { if v == "" || a.cfg.TunnelPinned { return false } diff --git a/internal/agent/session.go b/internal/agent/session.go index bc12ec0..7f4b77f 100644 --- a/internal/agent/session.go +++ b/internal/agent/session.go @@ -54,6 +54,11 @@ func (a *Agent) bridgeStream(parent context.Context, sid string, rw io.ReadWrite ctx, cancel := context.WithCancel(parent) defer cancel() defer rw.Close() + defer func() { + if l := a.link.Load(); l != nil { + l.unsubscribe(sid) + } + }() // A session closed from the panel (a tab reloaded, a window shut) ends // the stream at once, not on its next output: an idle shell would live // on until the agent stopped. @@ -88,6 +93,21 @@ func (a *Agent) bridgeStream(parent context.Context, sid string, rw io.ReadWrite // pollIn performs one downstream long-poll and applies the frames. // Returns false when the session is gone and the loop should stop. func (a *Agent) pollIn(ctx context.Context, sid string, w io.Writer, onResize func(cols, rows uint16)) bool { + // On the WebSocket the input comes pushed; if it drops, by a long poll. + if l := a.link.Load(); l != nil { + if ch := l.subscribe(sid); ch != nil { + select { + case in := <-ch: + if in.gone { + return false + } + return applyFrames(in.frames, w, onResize) + case <-l.done: + case <-ctx.Done(): + return false + } + } + } req, _ := http.NewRequestWithContext(ctx, http.MethodGet, a.cfg.Host+"/api/agent/term/"+sid+"/in", nil) a.authHeaders(req) resp, err := a.http.Do(req) @@ -105,7 +125,13 @@ func (a *Agent) pollIn(ctx context.Context, sid string, w io.Writer, onResize fu if err := json.NewDecoder(resp.Body).Decode(&batch); err != nil { return true } - for _, f := range batch.Frames { + return applyFrames(batch.Frames, w, onResize) +} + +// applyFrames writes a session's input to its stream; false once the panel +// closed it. +func applyFrames(frames []proto.InFrame, w io.Writer, onResize func(cols, rows uint16)) bool { + for _, f := range frames { switch f.Type { case proto.FrameStdin: if data, err := b64.DecodeString(f.Data); err == nil { @@ -124,6 +150,11 @@ func (a *Agent) pollIn(ctx context.Context, sid string, w io.Writer, onResize fu // postOut sends one upstream message. Returns false if the session is gone. func (a *Agent) postOut(ctx context.Context, sid string, msg proto.OutMessage) bool { + if l := a.link.Load(); l != nil { + if l.send(proto.WSMsg{T: "out", SID: sid, Out: &msg}) == nil { + return true + } + } body, _ := json.Marshal(msg) req, _ := http.NewRequestWithContext(ctx, http.MethodPost, a.cfg.Host+"/api/agent/term/"+sid+"/out", bytes.NewReader(body)) a.authHeaders(req) diff --git a/internal/agent/wslink.go b/internal/agent/wslink.go new file mode 100644 index 0000000..55d672b --- /dev/null +++ b/internal/agent/wslink.go @@ -0,0 +1,172 @@ +package agent + +import ( + "context" + "errors" + "fmt" + "log" + "net/http" + "strings" + "sync" + "time" + + "github.com/coder/websocket" + "github.com/coder/websocket/wsjson" + + "github.com/x-grk/longpolling-remote-access/internal/proto" +) + +// The WebSocket to the server: control messages arrive on it as soon as they +// are queued, and terminal and http-forward sessions send and take their +// bytes over it instead of a request each. When it drops, sessions carry on +// over plain requests. + +const wsPing = 20 * time.Second + +type wsLink struct { + c *websocket.Conn + ctx context.Context + mu sync.Mutex + subs map[string]chan wsIn + asked map[string]bool // sessions whose input was asked for + done chan struct{} +} + +// wsIn is what the server sent for a session: its input, or that it is over. +type wsIn struct { + frames []proto.InFrame + gone bool +} + +func (l *wsLink) send(m proto.WSMsg) error { + ctx, cancel := context.WithTimeout(l.ctx, 30*time.Second) + defer cancel() + return wsjson.Write(ctx, l.c, m) +} + +// subscribe asks the server for a session's input on this link, once, and +// returns where it arrives; nil when the link cannot take it. +func (l *wsLink) subscribe(sid string) chan wsIn { + l.mu.Lock() + ch := l.subs[sid] + if ch == nil { + ch = make(chan wsIn, 64) + l.subs[sid] = ch + } + fresh := !l.subscribed(sid) + l.mu.Unlock() + if fresh && l.send(proto.WSMsg{T: "sub", SID: sid}) != nil { + return nil + } + return ch +} + +// subscribed marks a session as asked for; l.mu is held. +func (l *wsLink) subscribed(sid string) bool { + if l.asked == nil { + l.asked = map[string]bool{} + } + if l.asked[sid] { + return true + } + l.asked[sid] = true + return false +} + +func (l *wsLink) unsubscribe(sid string) { + l.mu.Lock() + delete(l.subs, sid) + delete(l.asked, sid) + l.mu.Unlock() +} + +func (l *wsLink) deliver(sid string, in wsIn) { + l.mu.Lock() + ch := l.subs[sid] + l.mu.Unlock() + if ch == nil { + return + } + select { + case ch <- in: + case <-l.done: + } +} + +// runWS holds the WebSocket until it ends, dispatching what comes on it. +func (a *Agent) runWS(ctx context.Context) error { + u := a.cfg.Host + "/api/agent/ws" + u = "ws" + strings.TrimPrefix(u, "http") + hdr := http.Header{} + hdr.Set("X-Agent-Id", a.state.AgentID) + hdr.Set("Authorization", "Bearer "+a.state.Token) + hdr.Set("X-Agent-Version", Version) + hdr.Set("X-Agent-User", currentUser()) + dctx, dcancel := context.WithTimeout(ctx, 15*time.Second) + c, resp, err := websocket.Dial(dctx, u, &websocket.DialOptions{ + HTTPClient: a.stream, HTTPHeader: hdr, CompressionMode: websocket.CompressionDisabled, + }) + dcancel() + if err != nil { + if resp != nil { + a.applyTransport(resp.Header.Get(proto.TransportHeader)) + return fmt.Errorf("server answered %s", resp.Status) + } + return err + } + c.SetReadLimit(8 << 20) + lctx, lcancel := context.WithCancel(ctx) + l := &wsLink{c: c, ctx: lctx, subs: map[string]chan wsIn{}, done: make(chan struct{})} + a.link.Store(l) + defer func() { + a.link.CompareAndSwap(l, nil) + close(l.done) + lcancel() + c.CloseNow() + }() + log.Printf("connected by websocket") + go func() { + t := time.NewTicker(wsPing) + defer t.Stop() + for { + select { + case <-lctx.Done(): + return + case <-t.C: + pctx, pcancel := context.WithTimeout(lctx, 15*time.Second) + err := c.Ping(pctx) + pcancel() + if err != nil { + lcancel() + return + } + } + } + }() + for { + var m proto.WSMsg + if err := wsjson.Read(lctx, c, &m); err != nil { + if errors.Is(err, context.Canceled) && ctx.Err() == nil { + return errors.New("no answer to ping") + } + return err + } + switch m.T { + case "ctrl": + if m.Ctrl != nil { + // Sessions live on the agent's context: they outlast the link. + a.dispatch(ctx, *m.Ctrl) + } + case "cfg": + a.applyTunnelAddr(ctx, m.Tunnel) + a.applyTransport(m.Transport) + if a.mode() == proto.TransportLP { + return errors.New("set to long polling") + } + case "in": + l.deliver(m.SID, wsIn{frames: m.Frames}) + case "gone": + l.deliver(m.SID, wsIn{gone: true}) + } + } +} diff --git a/internal/proto/proto.go b/internal/proto/proto.go index 581a154..24de5c6 100644 --- a/internal/proto/proto.go +++ b/internal/proto/proto.go @@ -45,6 +45,32 @@ type ControlMessage struct { User string `json:"user,omitempty"` // CtrlOpenTerminal: OS user for the shell ("" = the agent's own) } +// Transports an agent can reach the server by. +const ( + TransportAuto = "auto" // WebSocket, long polling while it cannot connect + TransportWS = "ws" // WebSocket only + TransportLP = "lp" // long polling only +) + +// TransportHeader carries the agent's transport setting on poll and config +// answers. +const TransportHeader = "X-Transport" + +// WSMsg is one message on an agent's WebSocket, either way. Server to agent: +// "ctrl" (Ctrl, as a control poll would answer), "cfg" (Tunnel and +// Transport, also a keep-alive), "in" (Frames for session SID), "gone" +// (session SID is over). Agent to server: "sub" (send the input of session +// SID here), "out" (Out of session SID). +type WSMsg struct { + T string `json:"t"` + Ctrl *ControlMessage `json:"ctrl,omitempty"` + SID string `json:"sid,omitempty"` + Frames []InFrame `json:"frames,omitempty"` + Out *OutMessage `json:"out,omitempty"` + Tunnel string `json:"tunnel,omitempty"` + Transport string `json:"transport,omitempty"` +} + // Terminal / stream frame types on the session "in" channel (server -> agent). // Reused for both PTY terminals and raw TCP port-forward streams. const ( diff --git a/internal/server/db.go b/internal/server/db.go index 5741034..b8326b7 100644 --- a/internal/server/db.go +++ b/internal/server/db.go @@ -118,6 +118,7 @@ CREATE TABLE IF NOT EXISTS forward_rules ( `ALTER TABLE forward_rules ADD COLUMN no_ufw INTEGER NOT NULL DEFAULT 0`, `ALTER TABLE forward_rules ADD COLUMN label TEXT NOT NULL DEFAULT ''`, `ALTER TABLE admins ADD COLUMN password_off INTEGER NOT NULL DEFAULT 0`, + `ALTER TABLE agents ADD COLUMN transport TEXT NOT NULL DEFAULT 'auto'`, } { if _, err := db.sql.Exec(alter); err != nil && !strings.Contains(err.Error(), "duplicate column") { return err @@ -354,6 +355,7 @@ type Agent struct { Version string RunUser string // OS user the agent process runs as (reported by the agent) LoginUser string // OS user terminals are opened as; "" = RunUser + Transport string // auto, ws or lp: how the agent is to reach us LastSeen time.Time } @@ -383,8 +385,8 @@ func (db *DB) GetAgent(id string) (*Agent, error) { a := &Agent{} var last int64 err := db.sql.QueryRow( - `SELECT id, name, hostname, os, arch, token_hash, version, run_user, login_user, last_seen FROM agents WHERE id = ?`, id). - Scan(&a.ID, &a.Name, &a.Hostname, &a.OS, &a.Arch, &a.TokenHash, &a.Version, &a.RunUser, &a.LoginUser, &last) + `SELECT id, name, hostname, os, arch, token_hash, version, run_user, login_user, transport, last_seen FROM agents WHERE id = ?`, id). + Scan(&a.ID, &a.Name, &a.Hostname, &a.OS, &a.Arch, &a.TokenHash, &a.Version, &a.RunUser, &a.LoginUser, &a.Transport, &last) if err == sql.ErrNoRows { return nil, nil } @@ -415,6 +417,11 @@ func (db *DB) ListAgents() ([]Agent, error) { return out, rows.Err() } +func (db *DB) SetAgentTransport(id, transport string) error { + _, err := db.sql.Exec(`UPDATE agents SET transport = ? WHERE id = ?`, transport, id) + return err +} + func (db *DB) TouchAgent(id string) error { _, err := db.sql.Exec(`UPDATE agents SET last_seen = ? WHERE id = ?`, time.Now().Unix(), id) return err diff --git a/internal/server/handlers_agent.go b/internal/server/handlers_agent.go index cee1d74..0873b10 100644 --- a/internal/server/handlers_agent.go +++ b/internal/server/handlers_agent.go @@ -53,8 +53,9 @@ func (s *Server) handlePoll(w http.ResponseWriter, r *http.Request) { _ = s.db.TouchAgent(a.ID) } msg := s.hub.WaitControl(r.Context().Done(), a.ID, pollHold) - // Set after the hold so a port changed meanwhile reaches the agent now. + // Set after the hold so a setting changed meanwhile reaches the agent now. s.setTunnelHeader(w) + w.Header().Set(proto.TransportHeader, s.transportOf(a.ID)) writeJSON(w, msg) } @@ -71,11 +72,13 @@ func (s *Server) setTunnelHeader(w http.ResponseWriter) { // GET /api/agent/config - settings an agent picks up when it connects, without // waiting out a poll hold. func (s *Server) handleAgentConfig(w http.ResponseWriter, r *http.Request) { - if s.authenticateAgent(r) == nil { + a := s.authenticateAgent(r) + if a == nil { http.Error(w, "unauthorized", http.StatusUnauthorized) return } s.setTunnelHeader(w) + w.Header().Set(proto.TransportHeader, s.transportOf(a.ID)) w.WriteHeader(http.StatusNoContent) } diff --git a/internal/server/handlers_panel.go b/internal/server/handlers_panel.go index fd93b85..3e17591 100644 --- a/internal/server/handlers_panel.go +++ b/internal/server/handlers_panel.go @@ -213,6 +213,9 @@ func (s *Server) handleAgentEdit(w http.ResponseWriter, r *http.Request) { "Updated": r.URL.Query().Get("updated") == "1", "UserOK": r.URL.Query().Get("uok"), "UserErr": r.URL.Query().Get("uerr"), + "Transport": firstNonEmpty(a.Transport, proto.TransportAuto), + "OnWS": s.wsAgents.On(a.ID), + "Online": s.online(a), } // The user list lives on the machine, so it is only known while it is online. if s.online(a) { @@ -225,6 +228,21 @@ func (s *Server) handleAgentEdit(w http.ResponseWriter, r *http.Request) { s.render(w, "edit.html", data) } +// POST /agents/{id}/transport - how the agent is to reach the server: auto, +// ws or lp. A connected agent picks it up within a poll round. +func (s *Server) handleAgentTransport(w http.ResponseWriter, r *http.Request) { + if !checkCSRF(r) { + http.Error(w, "bad csrf", http.StatusForbidden) + return + } + id, t := r.PathValue("id"), r.FormValue("transport") + switch t { + case proto.TransportAuto, proto.TransportWS, proto.TransportLP: + _ = s.db.SetAgentTransport(id, t) + } + http.Redirect(w, r, "/agents/"+id+"/edit", http.StatusSeeOther) +} + // usersTaskTimeout bounds how long the edit page waits for the agent. const usersTaskTimeout = 8 * time.Second diff --git a/internal/server/hub.go b/internal/server/hub.go index 377283c..b6ab05e 100644 --- a/internal/server/hub.go +++ b/internal/server/hub.go @@ -275,6 +275,18 @@ func (s *Session) PushIn(f proto.InFrame) { } } +// Unpoll puts frames taken by PollIn back in front, for frames that could not +// be delivered. +func (s *Session) Unpoll(frames []proto.InFrame) { + s.mu.Lock() + s.inBuf = append(append([]proto.InFrame{}, frames...), s.inBuf...) + s.mu.Unlock() + select { + case s.inSignal <- struct{}{}: + default: + } +} + // PollIn drains queued downstream frames, blocking up to timeout for the first. func (s *Session) PollIn(timeout time.Duration) ([]proto.InFrame, bool) { deadline := time.After(timeout) diff --git a/internal/server/server.go b/internal/server/server.go index 8a48478..1ec53a4 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -60,6 +60,7 @@ type Server struct { meters speedMeters // byte counts of the speed tests under way xfers xfers // files on their way between a browser and an agent challenges challenges // passkey challenges handed out + wsAgents wsAgents // agents on their WebSocket now assets map[string]string // static file -> hash of its content } @@ -224,6 +225,7 @@ func (s *Server) Handler() http.Handler { mux.Handle("POST /agents/{id}/rename", s.requireAdmin(s.handleAgentRename)) mux.Handle("POST /agents/{id}/update", s.requireAdmin(s.handleAgentUpdate)) mux.Handle("POST /agents/{id}/user", s.requireAdmin(s.handleAgentSetUser)) + mux.Handle("POST /agents/{id}/transport", s.requireAdmin(s.handleAgentTransport)) mux.Handle("POST /agents/{id}/users/add", s.requireAdmin(s.handleAgentAddUser)) mux.Handle("POST /agents/{id}/delete", s.requireAdmin(s.handleAgentDelete)) mux.Handle("GET /settings", s.requireAdmin(s.handleSettings)) @@ -274,6 +276,7 @@ func (s *Server) Handler() http.Handler { // Agent API (token/secret auth inside handlers). mux.HandleFunc("POST /api/agent/enroll", s.handleEnroll) mux.HandleFunc("GET /api/agent/poll", s.handlePoll) + mux.HandleFunc("GET /api/agent/ws", s.handleAgentWS) mux.HandleFunc("GET /api/agent/config", s.handleAgentConfig) mux.HandleFunc("GET /api/agent/term/{sid}/in", s.handleTermIn) mux.HandleFunc("POST /api/agent/term/{sid}/out", s.handleTermOut) diff --git a/internal/server/web/templates/edit.html b/internal/server/web/templates/edit.html index b4ac97b..ac25e00 100644 --- a/internal/server/web/templates/edit.html +++ b/internal/server/web/templates/edit.html @@ -72,6 +72,25 @@ {{end}} +
+

Связь с сервером

+

+ Сейчас: {{if not .Online}}офлайн{{else if .OnWS}}WebSocket{{else}}long polling{{end}} +

+
+ + + +
+

WebSocket держит одно соединение: команды и терминал идут по нему без отдельного запроса на каждое нажатие. + В режиме «Авто» клиент, потеряв WebSocket, переходит на long polling и раз в минуту пробует вернуться. + Нужен клиент 0.1.0+ и прокси, пропускающий WebSocket (nginx: заголовки Upgrade и Connection).

+
+

Версия клиента

{{if .Updated}}
Команда обновления отправлена. Клиент скачает свежий бинарник и перезапустится.
{{end}} diff --git a/internal/server/wsagent.go b/internal/server/wsagent.go new file mode 100644 index 0000000..2091170 --- /dev/null +++ b/internal/server/wsagent.go @@ -0,0 +1,202 @@ +package server + +import ( + "context" + "net/http" + "sync" + "time" + + "github.com/coder/websocket" + "github.com/coder/websocket/wsjson" + + "github.com/x-grk/longpolling-remote-access/internal/proto" +) + +// An agent's WebSocket carries what its long polls would: control messages +// pushed the moment they are queued, and the input and output of its +// terminal and http-forward sessions, without a request per keystroke. Tasks, +// file transfers and the speed test stay on plain requests. + +// wsPing keeps the connection alive through proxies that drop quiet ones +// (Cloudflare after 100 s). +const wsPing = 20 * time.Second + +// wsAgents are the agents connected by WebSocket now, for the panel to show. +type wsAgents struct { + mu sync.Mutex + m map[string]*websocket.Conn +} + +// set records the agent's connection, closing one it had before. +func (c *wsAgents) set(id string, conn *websocket.Conn) { + c.mu.Lock() + if c.m == nil { + c.m = map[string]*websocket.Conn{} + } + old := c.m[id] + c.m[id] = conn + c.mu.Unlock() + if old != nil { + old.Close(websocket.StatusGoingAway, "replaced") + } +} + +func (c *wsAgents) drop(id string, conn *websocket.Conn) { + c.mu.Lock() + if c.m[id] == conn { + delete(c.m, id) + } + c.mu.Unlock() +} + +// On reports whether the agent is on its WebSocket now. +func (c *wsAgents) On(id string) bool { + c.mu.Lock() + defer c.mu.Unlock() + return c.m[id] != nil +} + +// transportOf is the agent's setting as stored, "auto" when unset. +func (s *Server) transportOf(agentID string) string { + if a, err := s.db.GetAgent(agentID); err == nil && a != nil && a.Transport != "" { + return a.Transport + } + return proto.TransportAuto +} + +// GET /api/agent/ws - the agent's WebSocket. Refused (409) when the agent is +// set to long polling. +func (s *Server) handleAgentWS(w http.ResponseWriter, r *http.Request) { + a := s.authenticateAgent(r) + if a == nil { + http.Error(w, "unauthorized", http.StatusUnauthorized) + return + } + if a.Transport == proto.TransportLP { + w.Header().Set(proto.TransportHeader, proto.TransportLP) + http.Error(w, "long polling only", http.StatusConflict) + return + } + if v := r.Header.Get("X-Agent-Version"); v != "" { + _ = s.db.TouchAgentInfo(a.ID, v, firstNonEmpty(r.Header.Get("X-Agent-User"), a.RunUser)) + } + conn, err := websocket.Accept(w, r, &websocket.AcceptOptions{CompressionMode: websocket.CompressionDisabled}) + if err != nil { + return + } + conn.SetReadLimit(4 << 20) + ctx, cancel := context.WithCancel(r.Context()) + defer cancel() + s.wsAgents.set(a.ID, conn) + defer s.wsAgents.drop(a.ID, conn) + defer conn.CloseNow() + + send := func(m proto.WSMsg) error { + wctx, wcancel := context.WithTimeout(ctx, 30*time.Second) + defer wcancel() + return wsjson.Write(wctx, conn, m) + } + cfg := func() proto.WSMsg { + return proto.WSMsg{T: "cfg", Tunnel: s.tunnelAdvert(), Transport: s.transportOf(a.ID)} + } + if send(cfg()) != nil { + return + } + + // Control: what a poll would get, as soon as it is queued. An idle round + // sends the settings, which doubles as a sign of life. + go func() { + defer cancel() + for ctx.Err() == nil { + msg := s.hub.WaitControl(ctx.Done(), a.ID, pollHold) + if ctx.Err() != nil { + return + } + _ = s.db.TouchAgent(a.ID) + out := proto.WSMsg{T: "ctrl", Ctrl: &msg} + if msg.Type == proto.CtrlIdle { + out = cfg() + } + if send(out) != nil { + return + } + if out.Transport == proto.TransportLP { + conn.Close(websocket.StatusNormalClosure, "long polling only") + return + } + } + }() + go func() { + t := time.NewTicker(wsPing) + defer t.Stop() + for { + select { + case <-ctx.Done(): + return + case <-t.C: + pctx, pcancel := context.WithTimeout(ctx, 15*time.Second) + err := conn.Ping(pctx) + pcancel() + if err != nil { + cancel() + return + } + } + } + }() + + subs := map[string]bool{} + for { + var m proto.WSMsg + if err := wsjson.Read(ctx, conn, &m); err != nil { + return + } + switch m.T { + case "sub": + if subs[m.SID] { + continue + } + sess := s.hub.Session(m.SID) + if sess == nil || sess.AgentID != a.ID { + _ = send(proto.WSMsg{T: "gone", SID: m.SID}) + continue + } + subs[m.SID] = true + go s.pumpSession(ctx, sess, send) + case "out": + sess := s.hub.Session(m.SID) + if sess == nil || sess.AgentID != a.ID || m.Out == nil { + _ = send(proto.WSMsg{T: "gone", SID: m.SID}) + continue + } + if !sess.PushOut(*m.Out) { + _ = send(proto.WSMsg{T: "gone", SID: m.SID}) + continue + } + if m.Out.Exit { + sess.close() + time.AfterFunc(pollHold, func() { s.hub.CloseSession(sess.ID) }) + } + } + } +} + +// pumpSession sends a session's input down the WebSocket as it comes. Input +// taken just as the connection ends goes back to the session, for the agent's +// long poll to collect. +func (s *Server) pumpSession(ctx context.Context, sess *Session, send func(proto.WSMsg) error) { + for ctx.Err() == nil { + frames, alive := sess.PollIn(pollHold) + if !alive { + _ = send(proto.WSMsg{T: "gone", SID: sess.ID}) + return + } + if len(frames) == 0 { + continue + } + if ctx.Err() != nil || send(proto.WSMsg{T: "in", SID: sess.ID, Frames: frames}) != nil { + sess.Unpoll(frames) + return + } + } +}