| 129 | } |
| 130 | |
| 131 | func RobustSSHConnectWithOptions(ctx context.Context, ip, keyFile string, port int, maxWait int, callback SSHProgressCallback, opts *SSHConnectOptions) (*SSHClient, error) { |
| 132 | config, err := newSSHConfig("ubuntu", keyFile) |
| 133 | if err != nil { |
| 134 | emitSSHProgress(callback, SSHStatusKeyParse, 0, err, "Failed to parse SSH private key") |
| 135 | return nil, err |
| 136 | } |
| 137 | |
| 138 | address := net.JoinHostPort(ip, strconv.Itoa(port)) |
| 139 | connectConfig := normalizeSSHConnectOptions(opts) |
| 140 | retryState := newSSHRetryState(maxWait) |
| 141 | |
| 142 | for attempt := 1; ; attempt++ { |
| 143 | if err := ctx.Err(); err != nil { |
| 144 | return nil, sshConnectionCancelledError() |
| 145 | } |
| 146 | if retryState.expired() { |
| 147 | return nil, timeoutError(maxWait, retryState.lastErr) |
| 148 | } |
| 149 | |
| 150 | dialTimeout := retryState.dialTimeout() |
| 151 | if dialTimeout <= 0 { |
| 152 | return nil, timeoutError(maxWait, retryState.lastErr) |
| 153 | } |
| 154 | |
| 155 | conn, dialErr := dialSSH(ctx, address, dialTimeout) |
| 156 | if dialErr != nil { |
| 157 | retryState.recordDialError(dialErr) |
| 158 | |
| 159 | if shouldRetryDial(dialErr) { |
| 160 | emitSSHProgress(callback, SSHStatusDialing, attempt, dialErr, "Waiting for instance to be ready...") |
| 161 | if err := sleepWithContext(ctx, retryState.backoffFor(SSHStatusDialing)); err != nil { |
| 162 | return nil, sshConnectionCancelledError() |
| 163 | } |
| 164 | retryState.advanceBackoff(SSHStatusDialing) |
| 165 | continue |
| 166 | } |
| 167 | return nil, fmt.Errorf("SSH dial failed: %w", dialErr) |
| 168 | } |
| 169 | |
| 170 | client, sshErr := handshakeSSHClient(ctx, conn, address, config) |
| 171 | if errors.Is(sshErr, errSSHConnectionCancelled) { |
| 172 | return nil, sshConnectionCancelledError() |
| 173 | } |
| 174 | if sshErr == nil { |
| 175 | emitSSHProgress(callback, SSHStatusSuccess, attempt, nil, "SSH connection established") |
| 176 | return client, nil |
| 177 | } |
| 178 | |
| 179 | errStatus := ClassifySSHError(sshErr) |
| 180 | retryState.recordSSHError(sshErr, errStatus) |
| 181 | |
| 182 | if errStatus == SSHStatusAuth && retryState.persistentAuthFailed(connectConfig) { |
| 183 | emitSSHProgress(callback, SSHStatusAuth, attempt, ErrPersistentAuthFailure, "Persistent authentication failure detected") |
| 184 | return nil, ErrPersistentAuthFailure |
| 185 | } |
| 186 | |
| 187 | if shouldRetrySSH(sshErr) { |
| 188 | emitSSHProgress(callback, errStatus, attempt, sshErr, retryMessageForStatus(errStatus)) |