diff --git a/pkg/wsman/client/wsman.go b/pkg/wsman/client/wsman.go index bcfbb553..cd544562 100644 --- a/pkg/wsman/client/wsman.go +++ b/pkg/wsman/client/wsman.go @@ -35,6 +35,10 @@ const ( RedirectionTLSPort = "16995" RedirectionNonTLSPort = "16994" WSManPath = "/wsman" + timeout = 30 * time.Second + idleConnTimeout = 90 * time.Second // keeps sockets warm as long as AMT holds them open + maxConcurrentPosts = 2 // caps simultaneous WSMAN Posts to match AMT firmware capacity + maxTransientRetries = 3 ) type Message struct { @@ -64,6 +68,9 @@ type Target struct { useDigest bool logAMTMessages bool challenge *AuthChallenge + challengeMu sync.Mutex // serializes access to challenge digest state + primeMu sync.Mutex // single-flights the initial cold-wake auth handshake + sem chan struct{} conn net.Conn bufferPool sync.Pool UseTLS bool @@ -72,7 +79,100 @@ type Target struct { tlsConfig *tls.Config } -const timeout = 10 * time.Second +// hostSems limits concurrency per device, shared across all Targets for a host. +var ( + hostSemsMu sync.Mutex + hostSems = make(map[string]chan struct{}) +) + +// hostSemaphore returns the shared concurrency limiter for an endpoint. +func hostSemaphore(endpoint string) chan struct{} { + hostSemsMu.Lock() + defer hostSemsMu.Unlock() + + if sem, ok := hostSems[endpoint]; ok { + return sem + } + + sem := make(chan struct{}, maxConcurrentPosts) + hostSems[endpoint] = sem + + return sem +} + +// isTransientTransportError reports whether err is a recoverable transport error. +func isTransientTransportError(err error) bool { + if err == nil { + return false + } + + if errors.Is(err, io.EOF) { + return true + } + + errText := strings.ToLower(err.Error()) + + return strings.Contains(errText, "connection reset by peer") || + strings.Contains(errText, "bad record mac") || + strings.Contains(errText, "tls: internal error") || + strings.Contains(errText, "broken pipe") +} + +// retryOnTransient retries req on transient transport errors until deadline. +func (t *Target) retryOnTransient(msgBody []byte, origReq *http.Request, origErr error, deadline time.Time) (*http.Response, error) { + if !isTransientTransportError(origErr) { + return nil, origErr + } + + authenticated := t.useDigest && origReq.Header.Get("Authorization") != "" + + lastErr := origErr + + for attempt := 1; attempt <= maxTransientRetries; attempt++ { + // All attempts share one budget so a Post cannot hold its slot for + // maxTransientRetries * timeout. + if time.Now().After(deadline) { + return nil, lastErr + } + + if tr, ok := t.Transport.(*http.Transport); ok { + tr.CloseIdleConnections() + } + + retryReq, reqErr := t.newPostRequest(msgBody) + if reqErr != nil { + return nil, reqErr + } + + retryReq.Header = origReq.Header.Clone() + + if authenticated { + auth, ok, authErr := t.digestAuthHeader() + if authErr != nil { + return nil, fmt.Errorf("failed digest auth %w", authErr) + } + + if ok { + retryReq.Header.Set("Authorization", auth) + } + } + + res, err := t.Do(retryReq) + if err == nil { + return res, nil + } + + lastErr = err + + logrus.WithError(err).Warnf("wsman transient retry %d/%d failed", attempt, maxTransientRetries) + + if !isTransientTransportError(err) { + return nil, err + } + } + + return nil, lastErr +} func NewWsman(cp Parameters) *Target { path := WSManPath @@ -97,6 +197,7 @@ func NewWsman(cp Parameters) *Target { InsecureSkipVerify: cp.SelfSignedAllowed, conn: cp.Connection, tlsConfig: cp.TlsConfig, + sem: hostSemaphore(cp.Target + ":" + port), } res.Timeout = timeout @@ -153,9 +254,9 @@ func NewWsman(cp Parameters) *Target { } res.Transport = &http.Transport{ - MaxIdleConns: 10, - IdleConnTimeout: 30 * time.Second, - DisableKeepAlives: true, + MaxIdleConns: maxConcurrentPosts, + IdleConnTimeout: idleConnTimeout, + DisableKeepAlives: false, TLSClientConfig: config, } } @@ -171,6 +272,9 @@ func NewWsman(cp Parameters) *Target { } func (t *Target) IsAuthenticated() bool { + t.challengeMu.Lock() + defer t.challengeMu.Unlock() + return t.challenge != nil && t.challenge.Realm != "" } @@ -220,27 +324,92 @@ func (t *Target) GetServerCertificate() (*tls.Certificate, error) { return capturedCert, nil } +// digestAuthHeader builds a Digest Authorization header under lock. +func (t *Target) digestAuthHeader() (auth string, ok bool, err error) { + t.challengeMu.Lock() + defer t.challengeMu.Unlock() + + if t.challenge.Realm == "" { + return "", false, nil + } + + auth, err = t.challenge.authorize("POST", "/wsman") + if err != nil { + return "", false, err + } + + return auth, true, nil +} + +// primeChallenge parses a 401 challenge and resets nc counter when nonce changes. +func (t *Target) primeChallenge(header string) error { + t.challengeMu.Lock() + defer t.challengeMu.Unlock() + + prevNonce := t.challenge.Nonce + + if err := t.challenge.parseChallenge(header); err != nil { + return err + } + + if t.challenge.Nonce != prevNonce { + t.challenge.NonceCount = 0 + } + + return nil +} + +func (t *Target) newPostRequest(msgBody []byte) (*http.Request, error) { + req, err := http.NewRequest("POST", t.endpoint, bytes.NewReader(msgBody)) + if err != nil { + return nil, err + } + + req.Header.Add("content-type", ContentType) + // GetBody enables auto-retry on stale keep-alive connections. + req.GetBody = func() (io.ReadCloser, error) { + return io.NopCloser(bytes.NewReader(msgBody)), nil + } + + return req, nil +} + // Post overrides http.Client's Post method. func (t *Target) Post(msg string) (response []byte, err error) { - msgBody := []byte(msg) + // Bound concurrent Posts to match AMT firmware connection capacity. + if t.sem != nil { + t.sem <- struct{}{} + defer func() { <-t.sem }() + } - var auth string + // Single-flight the first authentication handshake when digest is enabled. + if t.useDigest && !t.IsAuthenticated() { + t.primeMu.Lock() + if t.IsAuthenticated() { + t.primeMu.Unlock() + } else { + defer t.primeMu.Unlock() + } + } - bodyReader := bytes.NewReader(msgBody) + // Cap the whole Post, retries included, so a slot is never held for longer. + deadline := time.Now().Add(timeout) - req, err := http.NewRequest("POST", t.endpoint, bodyReader) + msgBody := []byte(msg) + + req, err := t.newPostRequest(msgBody) if err != nil { return nil, err } if t.username != "" && t.password != "" { if t.useDigest { - auth, err = t.challenge.authorize("POST", "/wsman") - if err != nil { - return nil, fmt.Errorf("failed digest auth %w", err) + auth, ok, authErr := t.digestAuthHeader() + if authErr != nil { + return nil, fmt.Errorf("failed digest auth %w", authErr) } - if t.challenge.Realm != "" { + if ok { req.Header.Set("Authorization", auth) } } else { @@ -248,40 +417,44 @@ func (t *Target) Post(msg string) (response []byte, err error) { } } - req.Header.Add("content-type", ContentType) - if t.logAMTMessages { logrus.Trace(msg) } res, err := t.Do(req) if err != nil { - return nil, err + res, err = t.retryOnTransient(msgBody, req, err, deadline) + if err != nil { + return nil, err + } } if t.useDigest && res.StatusCode == 401 { - if err := t.challenge.parseChallenge(res.Header.Get("WWW-Authenticate")); err != nil { - return nil, err - } + _, _ = io.Copy(io.Discard, res.Body) + _ = res.Body.Close() - auth, err = t.challenge.authorize("POST", "/wsman") - if err != nil { - return nil, fmt.Errorf("failed digest auth %w", err) + if parseErr := t.primeChallenge(res.Header.Get("WWW-Authenticate")); parseErr != nil { + return nil, parseErr } - bodyReader = bytes.NewReader(msgBody) + auth, _, authErr := t.digestAuthHeader() + if authErr != nil { + return nil, fmt.Errorf("failed digest auth %w", authErr) + } - req, err = http.NewRequest("POST", t.endpoint, bodyReader) + req, err = t.newPostRequest(msgBody) if err != nil { return nil, err } req.Header.Set("Authorization", auth) - req.Header.Add("content-type", ContentType) res, err = t.Do(req) - if err != nil && err.Error() != io.EOF.Error() { - return nil, err + if err != nil { + res, err = t.retryOnTransient(msgBody, req, err, deadline) + if err != nil { + return nil, err + } } }