mirror of
https://github.com/juanfont/headscale.git
synced 2026-04-11 03:27:20 +02:00
noise: pass context to sshActionFollowUp
Select on ctx.Done() alongside auth.WaitForAuth() so the goroutine exits promptly when the client disconnects instead of parking until cache eviction.
This commit is contained in:
@@ -1,6 +1,7 @@
|
|||||||
package hscontrol
|
package hscontrol
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"context"
|
||||||
"encoding/binary"
|
"encoding/binary"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
"errors"
|
"errors"
|
||||||
@@ -400,6 +401,7 @@ func (ns *noiseServer) SSHActionHandler(
|
|||||||
reqLog.Trace().Caller().Msg("SSH action request")
|
reqLog.Trace().Caller().Msg("SSH action request")
|
||||||
|
|
||||||
action, err := ns.sshAction(
|
action, err := ns.sshAction(
|
||||||
|
req.Context(),
|
||||||
reqLog,
|
reqLog,
|
||||||
srcNodeID, dstNodeID,
|
srcNodeID, dstNodeID,
|
||||||
req.URL.Query().Get("auth_id"),
|
req.URL.Query().Get("auth_id"),
|
||||||
@@ -437,6 +439,7 @@ func (ns *noiseServer) SSHActionHandler(
|
|||||||
// 3. Follow-up request — an auth_id is present, wait for the auth
|
// 3. Follow-up request — an auth_id is present, wait for the auth
|
||||||
// verdict and accept or reject.
|
// verdict and accept or reject.
|
||||||
func (ns *noiseServer) sshAction(
|
func (ns *noiseServer) sshAction(
|
||||||
|
ctx context.Context,
|
||||||
reqLog zerolog.Logger,
|
reqLog zerolog.Logger,
|
||||||
srcNodeID, dstNodeID types.NodeID,
|
srcNodeID, dstNodeID types.NodeID,
|
||||||
authIDStr string,
|
authIDStr string,
|
||||||
@@ -456,7 +459,7 @@ func (ns *noiseServer) sshAction(
|
|||||||
// Follow-up request with auth_id — wait for the auth verdict.
|
// Follow-up request with auth_id — wait for the auth verdict.
|
||||||
if authIDStr != "" {
|
if authIDStr != "" {
|
||||||
return ns.sshActionFollowUp(
|
return ns.sshActionFollowUp(
|
||||||
reqLog, &action, authIDStr,
|
ctx, reqLog, &action, authIDStr,
|
||||||
srcNodeID, dstNodeID,
|
srcNodeID, dstNodeID,
|
||||||
checkFound,
|
checkFound,
|
||||||
)
|
)
|
||||||
@@ -542,8 +545,10 @@ func (ns *noiseServer) sshActionHoldAndDelegate(
|
|||||||
}
|
}
|
||||||
|
|
||||||
// sshActionFollowUp handles follow-up requests where the client
|
// sshActionFollowUp handles follow-up requests where the client
|
||||||
// provides an auth_id. It blocks until the auth session resolves.
|
// provides an auth_id. It blocks until the auth session resolves or
|
||||||
|
// the request context is cancelled (e.g. the client disconnects).
|
||||||
func (ns *noiseServer) sshActionFollowUp(
|
func (ns *noiseServer) sshActionFollowUp(
|
||||||
|
ctx context.Context,
|
||||||
reqLog zerolog.Logger,
|
reqLog zerolog.Logger,
|
||||||
action *tailcfg.SSHAction,
|
action *tailcfg.SSHAction,
|
||||||
authIDStr string,
|
authIDStr string,
|
||||||
@@ -598,7 +603,21 @@ func (ns *noiseServer) sshActionFollowUp(
|
|||||||
|
|
||||||
reqLog.Trace().Caller().Msg("SSH action follow-up")
|
reqLog.Trace().Caller().Msg("SSH action follow-up")
|
||||||
|
|
||||||
verdict := <-auth.WaitForAuth()
|
var verdict types.AuthVerdict
|
||||||
|
select {
|
||||||
|
case <-ctx.Done():
|
||||||
|
// The client disconnected (or its request timed out) before the
|
||||||
|
// auth session resolved. Return an error so the parked goroutine
|
||||||
|
// is freed; without this select sshActionFollowUp would block
|
||||||
|
// until the cache eviction callback signalled FinishAuth, which
|
||||||
|
// could be up to register_cache_expiration (15 minutes).
|
||||||
|
return nil, NewHTTPError(
|
||||||
|
http.StatusUnauthorized,
|
||||||
|
"ssh action follow-up cancelled",
|
||||||
|
ctx.Err(),
|
||||||
|
)
|
||||||
|
case verdict = <-auth.WaitForAuth():
|
||||||
|
}
|
||||||
|
|
||||||
if !verdict.Accept() {
|
if !verdict.Accept() {
|
||||||
action.Reject = true
|
action.Reject = true
|
||||||
|
|||||||
Reference in New Issue
Block a user