BED-5717: Create BHE Client& Fix GOAWAY AWS Error Response #141
Conversation
WalkthroughThis change introduces a dedicated BloodHound Enterprise (BHE) client library in Go, refactors the main application to use this abstraction, and moves all HTTP, proxy, and signing logic into modular packages. It adds comprehensive testing, configuration for connection limits, and mock implementations for testing, while removing legacy HTTP and proxy handling from the main codebase. Changes
Sequence Diagram(s)sequenceDiagram
participant Main as Main Application
participant BHEClient as BHEClient
participant Server as BHE Server
Main->>BHEClient: NewBHEClient(params)
loop Job Polling
Main->>BHEClient: UpdateClient()
BHEClient->>Server: PATCH /client
Server-->>BHEClient: UpdateClientResponse
Main->>BHEClient: GetAvailableJobs()
BHEClient->>Server: GET /jobs/available
Server-->>BHEClient: [Job List]
Main->>BHEClient: StartJob(jobId)
BHEClient->>Server: POST /jobs/{jobId}/start
Main->>BHEClient: Ingest(data channel)
BHEClient->>Server: POST /ingest (with retries)
Main->>BHEClient: EndJob(status, message)
BHEClient->>Server: POST /jobs/{jobId}/end
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Poem
Note ⚡️ Unit Test Generation is now available in beta!Learn more here, or try it out under "Finishing Touches" below. ✨ Finishing Touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
mvlipka
changed the title
mvlipka
changed the title
mvlipka
force-pushed
the
BED-5717/BloodHound-Client
branch
from
6eeaac1 to
a44db36
Compare
mvlipka
changed the title
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 7
🧹 Nitpick comments (14)
client/rest/utils.go (1)
149-152: Consider adding bounds checking to prevent overflow.The exponential backoff calculation could result in very large values for high retry counts or base values, potentially causing integer overflow when converting to
time.Duration.func VariableExponentialBackoff(base, retry int) { backoff := math.Pow(float64(base), float64(retry+1)) + // Cap the backoff to prevent overflow (e.g., 5 minutes) + const maxBackoff = 300 + if backoff > maxBackoff { + backoff = maxBackoff + } time.Sleep(time.Second * time.Duration(backoff)) }client/bloodhound/reader.go (1)
8-22: Consider documenting interface implementation.While the implementation is correct, it would be clearer to document that this type implements
io.ReadCloserand provides additionalRewindfunctionality for retry scenarios.+// rewindableByteReader implements io.ReadCloser with the ability to rewind +// the read position to the beginning. This is useful for retrying HTTP requests +// that need to re-read the body. type rewindableByteReader struct { data *bytes.Reader }config/config.go (1)
275-283: Consider using consistent naming convention.The configuration name
maxReqsPerConnuses an abbreviation while other configs use full words. ConsidermaxRequestsPerConnectionfor consistency with naming patterns likemaxConnsPerHostandmaxIdleConnsPerHost.BHEMaxReqPerConn = Config{ - Name: "maxReqsPerConn", + Name: "maxRequestsPerConnection", Shorthand: "", Usage: "The number of requests a single HTTP connection can make, when this limit is reached, a new HTTP connection will be established with the server", Required: false, Default: 10_000, MinValue: 1, MaxValue: 10_000, }client/rest/dialer.go (1)
73-78: Remove redundant response body close.The response body is already closed in the deferred function (lines 64-68), making the explicit close on lines 74-76 redundant.
} else if res.StatusCode != 200 { - if res.Body != nil { - res.Body.Close() - } conn.Close() return nil, fmt.Errorf("unable to connect to %s via proxy (%s): statusCode %d", addr, s.host, res.StatusCode)client/bloodhound/response.go (1)
3-5: Consider adding type documentation.While the implementation is clean and correct, adding documentation would clarify its purpose as a generic wrapper for BHE API responses.
+// bloodhoundResponse is a generic wrapper for BloodHound Enterprise API responses +// that contain data in a standardized JSON structure with a "data" field. type bloodhoundResponse[T any] struct { Data T `json:"data"` }client/bloodhound/client_test.go (2)
15-61: Add test coverage for successful request scenarios.The test suite covers error scenarios well but lacks coverage for successful HTTP requests. Consider adding a test case that verifies the client correctly handles successful responses.
+ t.Run("successful request", func(t *testing.T) { + testServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + w.Write([]byte("success")) + })) + defer testServer.Close() + + testUrl, _ := url.Parse(testServer.URL) + + client, err := NewBHEClient(*testUrl, "tokenId", "token", "", 10, 1, logr.Discard()) + require.NoError(t, err) + + req, _ := http.NewRequest("GET", testUrl.String(), nil) + resp, err := client.SendRequest(req) + + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode) + })
63-136: Well-structured ingest tests. Consider adding context cancellation test.The existing tests provide good coverage for ingest scenarios. The tests correctly verify error handling and retry behavior.
Consider adding a test for context cancellation during ingest:
+ t.Run("context cancellation", func(t *testing.T) { + testServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + time.Sleep(100 * time.Millisecond) + w.WriteHeader(http.StatusAccepted) + })) + defer testServer.Close() + + testUrl, _ := url.Parse(testServer.URL) + client, err := NewBHEClient(*testUrl, "tokenId", "token", "", 1, 1, logr.Discard()) + require.NoError(t, err) + + data := make(chan []any, 1) + data <- []any{"test"} + close(data) + + ctx, cancel := context.WithCancel(context.Background()) + go func() { + time.Sleep(50 * time.Millisecond) + cancel() + }() + + hadErrors := client.Ingest(ctx, data) + require.True(t, hadErrors) + })client/bloodhound/transport.go (2)
46-51: Document the datetime truncation logic.The code truncates the datetime to 13 characters (up to the hour), but this behavior is not documented and could be confusing.
// datetime datetime := time.Now().Format(time.RFC3339) digester = hmac.New(sha256.New, digester.Sum(nil)) - // hash the substring of the current datetime excluding minutes, seconds, microseconds and timezone + // Truncate to hour precision (YYYY-MM-DDTHH) for signature stability within the same hour + // This allows replay of requests within the same hour window if _, err := digester.Write([]byte(datetime[:13])); err != nil {
95-97: Consider logging errors from discard operation.While ignoring errors during cleanup is often acceptable, logging them could help with debugging.
-func discard(reader io.Reader) { - io.Copy(io.Discard, reader) +func discard(reader io.Reader) { + if _, err := io.Copy(io.Discard, reader); err != nil { + // Log error but don't fail the operation + // log.V(2).Info("error discarding reader", "error", err) + } }cmd/start.go (1)
74-77: Consider making retry count configurable.The retry count is hardcoded to 3. Consider making this configurable for better flexibility in different environments.
- } else if bheClient, err := bloodhound.NewBHEClient(*bheInstance, config.BHETokenId.Value().(string), config.BHEToken.Value().(string), config.Proxy.Value().(string), config.BHEMaxReqPerConn.Value().(int), 3, log); err != nil { + } else if bheClient, err := bloodhound.NewBHEClient(*bheInstance, config.BHETokenId.Value().(string), config.BHEToken.Value().(string), config.Proxy.Value().(string), config.BHEMaxReqPerConn.Value().(int), config.BHEMaxRetries.Value().(int), log); err != nil {client/bloodhound/client.go (4)
78-78: Consider makingretryDelayconfigurable.The retry delay is hardcoded to 5. This could be made configurable to allow users to adjust the backoff behavior based on their specific needs.
287-287: Usedeferfor response body closure.To ensure the response body is closed even if an unexpected error occurs, use
defer.} else if res, err := s.SendRequest(req); err != nil { return err } else { - res.Body.Close() + defer res.Body.Close() return nil }
307-307: Usedeferfor response body closure.} else if res, err := s.SendRequest(req); err != nil { return err } else { - res.Body.Close() + defer res.Body.Close() return nil }
326-326: Usedeferfor response body closure.} else if res, err := s.SendRequest(req); err != nil { return err } else { - res.Body.Close() + defer res.Body.Close() return nil }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (11)
client/bloodhound/client.go(1 hunks)client/bloodhound/client_test.go(1 hunks)client/bloodhound/mocks/client.go(1 hunks)client/bloodhound/reader.go(1 hunks)client/bloodhound/response.go(1 hunks)client/bloodhound/transport.go(1 hunks)client/rest/dialer.go(1 hunks)client/rest/utils.go(4 hunks)cmd/start.go(7 hunks)cmd/utils.go(3 hunks)config/config.go(2 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (4)
config/config.go (1)
config/internal/config.go (1)
Config(32-41)
client/rest/dialer.go (1)
config/config.go (1)
Proxy(117-123)
client/bloodhound/mocks/client.go (2)
models/job.go (2)
JobStatus(11-11)ClientJob(52-62)models/update-client-request.go (1)
UpdateClientResponse(28-35)
cmd/utils.go (2)
client/rest/dialer.go (2)
NewProxyDialer(22-34)Dial(97-118)config/config.go (3)
AzAuthUrl(189-195)AzGraphUrl(196-202)AzMgmtUrl(203-209)
🪛 ast-grep (0.38.6)
client/rest/dialer.go
[warning] 18-18: MinVersionis missing from this TLS configuration. By default, TLS 1.2 is currently used as the minimum when acting as a client, and TLS 1.0 when acting as a server. General purpose web applications should default to TLS 1.3 with all other protocols disabled. Only where it is known that a web server must support legacy clients with unsupported an insecure browsers (such as Internet Explorer 10), it may be necessary to enable TLS 1.0 to provide support. AddMinVersion: tls.VersionTLS13' to the TLS configuration to bump the minimum version to TLS 1.3.
Context: tls.Config{}
Note: [CWE-327]: Use of a Broken or Risky Cryptographic Algorithm [OWASP A03:2017]: Sensitive Data Exposure [OWASP A02:2021]: Cryptographic Failures [REFERENCES]
https://owasp.org/Top10/A02_2021-Cryptographic_Failures
(missing-ssl-minversion-go)
🔇 Additional comments (11)
client/rest/utils.go (1)
138-142: LGTM! Good use of error type assertion.The implementation correctly uses
errors.Asto check forhttp2.GoAwayError, which aligns with the PR objective of handling GOAWAY frames from AWS load balancers.client/bloodhound/transport.go (1)
14-18: Clean adapter pattern implementation.Good use of the adapter pattern to allow functional RoundTripper implementations for testing.
cmd/utils.go (2)
43-46: Good refactoring to centralize proxy handling.The registration of proxy dialers in the init function properly delegates proxy handling to the rest package, improving code organization.
91-101: Clean migration to rest.Dial for connection testing.The refactored connection testing correctly uses the new rest.Dial function while maintaining the same error handling pattern.
client/bloodhound/mocks/client.go (1)
1-8: Generated mock file - no issues.This is a properly generated GoMock file that implements the BloodHoundClient interface for testing purposes.
cmd/start.go (2)
103-105: Good resource cleanup with defer statements.Proper use of defer statements to ensure idle connections are closed after job processing, preventing resource leaks.
145-146: Clean integration with the new ingest method.The refactored code cleanly uses the new BHE client's Ingest method, properly handling the error flag return value.
client/bloodhound/client.go (4)
257-276: LGTM!The method properly uses
SendRequestfor HTTP communication and follows good error handling practices.
331-364: LGTM!The method properly handles client updates with appropriate error handling and resource cleanup.
366-374: LGTM!Clean implementation for handling orphaned jobs.
376-379: LGTM!Simple and correct wrapper implementation.
| // Ingest sends the ingest data to the BHE server and returns true if there were any errors while making the request | ||
| func (s *BHEClient) Ingest(ctx context.Context, in <-chan []any) bool { | ||
| endpoint := s.bheUrl.ResolveReference(&url.URL{Path: "/api/v2/ingest"}) | ||
|
|
||
| var ( | ||
| hasErrors = false | ||
| unrecoverableErrMsg = fmt.Sprintf("ending current ingest job due to unrecoverable error while requesting %v", endpoint) | ||
| ) | ||
|
|
||
| for data := range pipeline.OrDone(ctx.Done(), in) { | ||
| var ( | ||
| body bytes.Buffer | ||
| gw = gzip.NewWriter(&body) | ||
| ) | ||
|
|
||
| ingestData := models.IngestRequest{ | ||
| Meta: models.Meta{ | ||
| Type: "azure", | ||
| }, | ||
| Data: data, | ||
| } | ||
|
|
||
| err := json.NewEncoder(gw).Encode(ingestData) | ||
| if err != nil { | ||
| s.log.Error(err, unrecoverableErrMsg) | ||
| } | ||
| gw.Close() | ||
|
|
||
| if req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint.String(), &body); err != nil { | ||
| s.log.Error(err, unrecoverableErrMsg) | ||
| return true | ||
| } else { | ||
| req.Header.Set("User-Agent", constants.UserAgent()) | ||
| req.Header.Set("Accept", "application/json") | ||
| req.Header.Set("Content-Encoding", "gzip") | ||
|
|
||
| for currentAttempt := 0; currentAttempt <= s.maxRetries; currentAttempt++ { | ||
| // No retries on regular err cases, only on HTTP 504 Gateway Timeout and HTTP 503 Service Unavailable | ||
| if response, err := s.httpClient.Do(req); err != nil { | ||
| if rest.IsClosedConnectionErr(err) { | ||
| // try again on force closed connection | ||
| s.log.Error(err, fmt.Sprintf("remote host force closed connection while requesting %s; attempt %d/%d; trying again", req.URL, currentAttempt+1, s.maxRetries)) | ||
|
|
||
| if currentAttempt == s.maxRetries { | ||
| s.log.Error(ErrExceededRetryLimit, "") | ||
| hasErrors = true | ||
| } else { | ||
| rest.VariableExponentialBackoff(s.retryDelay, currentAttempt) | ||
| } | ||
|
|
||
| continue | ||
| } else if rest.IsGoAwayErr(err) { | ||
| // AWS currently has a 10,000 request per connection limitation, retry in case AWS changes this limitation | ||
| s.log.Error(err, fmt.Sprintf("received GOAWAY from from AWS load balancer while requesting %s; attempt %d/%d; trying again", req.URL, currentAttempt+1, s.maxRetries)) | ||
|
|
||
| if currentAttempt == s.maxRetries { | ||
| s.log.Error(ErrExceededRetryLimit, "") | ||
| hasErrors = true | ||
| } else { | ||
| rest.VariableExponentialBackoff(s.retryDelay, currentAttempt) | ||
| } | ||
|
|
||
| continue | ||
| } | ||
|
|
||
| s.log.Error(err, unrecoverableErrMsg) | ||
| return true | ||
| } else if response.StatusCode == http.StatusGatewayTimeout || response.StatusCode == http.StatusServiceUnavailable || response.StatusCode == http.StatusBadGateway { | ||
| serverError := fmt.Errorf("received server error %d while requesting %v; attempt %d/%d; trying again", response.StatusCode, endpoint, currentAttempt+1, s.maxRetries) | ||
| s.log.Error(serverError, "") | ||
|
|
||
| if currentAttempt == s.maxRetries { | ||
| s.log.Error(ErrExceededRetryLimit, "") | ||
| hasErrors = true | ||
| } else { | ||
| rest.VariableExponentialBackoff(s.retryDelay, currentAttempt) | ||
| } | ||
|
|
||
| if err := response.Body.Close(); err != nil { | ||
| s.log.Error(fmt.Errorf("failed to close ingest body: %w", err), unrecoverableErrMsg) | ||
| } | ||
|
|
||
| continue | ||
| } else if response.StatusCode != http.StatusAccepted { | ||
| if bodyBytes, err := io.ReadAll(response.Body); err != nil { | ||
| s.log.Error(fmt.Errorf("received unexpected response code from %v: %s; failure reading response body", endpoint, response.Status), unrecoverableErrMsg) | ||
| } else { | ||
| s.log.Error(fmt.Errorf("received unexpected response code from %v: %s %s", req.URL, response.Status, bodyBytes), unrecoverableErrMsg) | ||
| } | ||
|
|
||
| if err := response.Body.Close(); err != nil { | ||
| s.log.Error(fmt.Errorf("failed to close ingest body: %w", err), unrecoverableErrMsg) | ||
| } | ||
|
|
||
| return true | ||
| } else { | ||
| if err := response.Body.Close(); err != nil { | ||
| s.log.Error(fmt.Errorf("failed to close ingest body: %w", err), unrecoverableErrMsg) | ||
| } | ||
| } | ||
| } | ||
| } | ||
| } | ||
| return hasErrors | ||
| } |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Consider using SendRequest method to avoid duplicating retry logic.
The Ingest method reimplements retry logic that already exists in SendRequest. This creates code duplication and misses the connection reset functionality.
Consider refactoring to use SendRequest:
- for currentAttempt := 0; currentAttempt <= s.maxRetries; currentAttempt++ {
- // No retries on regular err cases, only on HTTP 504 Gateway Timeout and HTTP 503 Service Unavailable
- if response, err := s.httpClient.Do(req); err != nil {
+ if response, err := s.SendRequest(req); err != nil {
+ s.log.Error(err, "error during ingest request")
+ hasErrors = true
+ } else {
+ defer response.Body.Close()
+ if response.StatusCode != http.StatusAccepted {
+ s.log.Error(fmt.Errorf("unexpected status code: %d", response.StatusCode), "ingest request failed")
+ hasErrors = true
+ }
+ }This would:
- Reuse existing retry logic
- Benefit from connection reset functionality
- Reduce code duplication
- Simplify the method significantly
Committable suggestion skipped: line range outside the PR's diff.
🤖 Prompt for AI Agents
In client/bloodhound/client.go from lines 151 to 255, the Ingest method
duplicates retry logic that is already implemented in the SendRequest method,
missing connection reset handling and causing code duplication. Refactor Ingest
to construct the HTTP request and then call SendRequest to handle sending with
retries and error handling. This will reuse existing retry logic, include
connection reset functionality, reduce duplication, and simplify the method.
| func (s signingTransport) RoundTrip(req *http.Request) (*http.Response, error) { | ||
| // The http client may try to call RoundTrip more than once to replay the same request; in which case rewind the request | ||
| if rbr, ok := req.Body.(*rewindableByteReader); ok { | ||
| if _, err := rbr.Rewind(); err != nil { | ||
| return nil, err | ||
| } | ||
| } | ||
|
|
||
| if req.Header.Get("Signature") == "" { | ||
|
|
||
| // token | ||
| digester := hmac.New(sha256.New, []byte(s.token)) | ||
|
|
||
| // path | ||
| if _, err := digester.Write([]byte(req.Method + req.URL.Path)); err != nil { | ||
| return nil, err | ||
| } | ||
|
|
||
| // datetime | ||
| datetime := time.Now().Format(time.RFC3339) | ||
| digester = hmac.New(sha256.New, digester.Sum(nil)) | ||
| // hash the substring of the current datetime excluding minutes, seconds, microseconds and timezone | ||
| if _, err := digester.Write([]byte(datetime[:13])); err != nil { | ||
| return nil, err | ||
| } | ||
|
|
||
| // body | ||
| digester = hmac.New(sha256.New, digester.Sum(nil)) | ||
| if req.Body != nil { | ||
| var ( | ||
| body = &bytes.Buffer{} | ||
| hashBuf = make([]byte, 64*1024) // 64KB buffer, consider benchmarking and optimizing this value | ||
| tee = io.TeeReader(req.Body, body) | ||
| ) | ||
|
|
||
| defer req.Body.Close() | ||
| defer discard(tee) | ||
| defer discard(body) | ||
|
|
||
| for { | ||
| numRead, err := tee.Read(hashBuf) | ||
| if numRead > 0 { | ||
| if _, err := digester.Write(hashBuf[:numRead]); err != nil { | ||
| return nil, err | ||
| } | ||
| } | ||
|
|
||
| // exit loop on EOF or error | ||
| if err != nil { | ||
| if err != io.EOF { | ||
| return nil, err | ||
| } | ||
| break | ||
| } | ||
| } | ||
|
|
||
| req.Body = &rewindableByteReader{data: bytes.NewReader(body.Bytes())} | ||
| } | ||
|
|
||
| signature := digester.Sum(nil) | ||
|
|
||
| req.Header.Set("Authorization", fmt.Sprintf("%s %s", s.signature, s.tokenId)) | ||
| req.Header.Set("RequestDate", datetime) | ||
| req.Header.Set("Signature", base64.StdEncoding.EncodeToString(signature)) | ||
| } | ||
| return s.base.RoundTrip(req) | ||
| } |
There was a problem hiding this comment.
Potential memory issues with large request bodies.
The current implementation buffers the entire request body in memory, which could cause out-of-memory issues for large payloads.
Consider implementing a size limit or streaming approach for large bodies:
if req.Body != nil {
+ // Check Content-Length if available
+ const maxBodySize = 100 * 1024 * 1024 // 100MB limit
+ if req.ContentLength > 0 && req.ContentLength > maxBodySize {
+ return nil, fmt.Errorf("request body too large: %d bytes exceeds %d byte limit", req.ContentLength, maxBodySize)
+ }
+
var (
body = &bytes.Buffer{}
hashBuf = make([]byte, 64*1024) // 64KB buffer, consider benchmarking and optimizing this value
tee = io.TeeReader(req.Body, body)
)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| func (s signingTransport) RoundTrip(req *http.Request) (*http.Response, error) { | |
| // The http client may try to call RoundTrip more than once to replay the same request; in which case rewind the request | |
| if rbr, ok := req.Body.(*rewindableByteReader); ok { | |
| if _, err := rbr.Rewind(); err != nil { | |
| return nil, err | |
| } | |
| } | |
| if req.Header.Get("Signature") == "" { | |
| // token | |
| digester := hmac.New(sha256.New, []byte(s.token)) | |
| // path | |
| if _, err := digester.Write([]byte(req.Method + req.URL.Path)); err != nil { | |
| return nil, err | |
| } | |
| // datetime | |
| datetime := time.Now().Format(time.RFC3339) | |
| digester = hmac.New(sha256.New, digester.Sum(nil)) | |
| // hash the substring of the current datetime excluding minutes, seconds, microseconds and timezone | |
| if _, err := digester.Write([]byte(datetime[:13])); err != nil { | |
| return nil, err | |
| } | |
| // body | |
| digester = hmac.New(sha256.New, digester.Sum(nil)) | |
| if req.Body != nil { | |
| var ( | |
| body = &bytes.Buffer{} | |
| hashBuf = make([]byte, 64*1024) // 64KB buffer, consider benchmarking and optimizing this value | |
| tee = io.TeeReader(req.Body, body) | |
| ) | |
| defer req.Body.Close() | |
| defer discard(tee) | |
| defer discard(body) | |
| for { | |
| numRead, err := tee.Read(hashBuf) | |
| if numRead > 0 { | |
| if _, err := digester.Write(hashBuf[:numRead]); err != nil { | |
| return nil, err | |
| } | |
| } | |
| // exit loop on EOF or error | |
| if err != nil { | |
| if err != io.EOF { | |
| return nil, err | |
| } | |
| break | |
| } | |
| } | |
| req.Body = &rewindableByteReader{data: bytes.NewReader(body.Bytes())} | |
| } | |
| signature := digester.Sum(nil) | |
| req.Header.Set("Authorization", fmt.Sprintf("%s %s", s.signature, s.tokenId)) | |
| req.Header.Set("RequestDate", datetime) | |
| req.Header.Set("Signature", base64.StdEncoding.EncodeToString(signature)) | |
| } | |
| return s.base.RoundTrip(req) | |
| } | |
| if req.Body != nil { | |
| // Check Content-Length if available | |
| const maxBodySize = 100 * 1024 * 1024 // 100MB limit | |
| if req.ContentLength > 0 && req.ContentLength > maxBodySize { | |
| return nil, fmt.Errorf( | |
| "request body too large: %d bytes exceeds %d byte limit", | |
| req.ContentLength, maxBodySize, | |
| ) | |
| } | |
| var ( | |
| body = &bytes.Buffer{} | |
| hashBuf = make([]byte, 64*1024) // 64KB buffer, consider benchmarking and optimizing this value | |
| tee = io.TeeReader(req.Body, body) | |
| ) | |
| defer req.Body.Close() | |
| defer discard(tee) | |
| defer discard(body) | |
| for { | |
| numRead, err := tee.Read(hashBuf) | |
| if numRead > 0 { | |
| if _, err := digester.Write(hashBuf[:numRead]); err != nil { | |
| return nil, err | |
| } | |
| } | |
| if err != nil { | |
| if err != io.EOF { | |
| return nil, err | |
| } | |
| break | |
| } | |
| } | |
| req.Body = &rewindableByteReader{data: bytes.NewReader(body.Bytes())} | |
| } |
🤖 Prompt for AI Agents
In client/bloodhound/transport.go between lines 27 and 93, the current RoundTrip
method reads and buffers the entire request body into memory, which risks
out-of-memory errors for large payloads. To fix this, implement a size limit on
the amount of data read from the body or switch to a streaming approach that
processes the body in chunks without full buffering. This can be done by
tracking the total bytes read and stopping or returning an error if a predefined
maximum size is exceeded, or by redesigning the signature calculation to work
incrementally with streamed data.
| func (s HttpsDialer) Dial(network string, addr string) (net.Conn, error) { | ||
| return tls.Dial(network, addr, &tls.Config{}) | ||
| } |
There was a problem hiding this comment.
Security: Set minimum TLS version to prevent downgrade attacks.
The TLS configuration is missing MinVersion, which could allow connections using outdated and insecure TLS versions. This is a security vulnerability.
func (s HttpsDialer) Dial(network string, addr string) (net.Conn, error) {
- return tls.Dial(network, addr, &tls.Config{})
+ return tls.Dial(network, addr, &tls.Config{
+ MinVersion: tls.VersionTLS12,
+ })
}🧰 Tools
🪛 ast-grep (0.38.6)
[warning] 18-18: MinVersionis missing from this TLS configuration. By default, TLS 1.2 is currently used as the minimum when acting as a client, and TLS 1.0 when acting as a server. General purpose web applications should default to TLS 1.3 with all other protocols disabled. Only where it is known that a web server must support legacy clients with unsupported an insecure browsers (such as Internet Explorer 10), it may be necessary to enable TLS 1.0 to provide support. AddMinVersion: tls.VersionTLS13' to the TLS configuration to bump the minimum version to TLS 1.3.
Context: tls.Config{}
Note: [CWE-327]: Use of a Broken or Risky Cryptographic Algorithm [OWASP A03:2017]: Sensitive Data Exposure [OWASP A02:2021]: Cryptographic Failures [REFERENCES]
https://owasp.org/Top10/A02_2021-Cryptographic_Failures
(missing-ssl-minversion-go)
🤖 Prompt for AI Agents
In client/rest/dialer.go around lines 18 to 20, the TLS configuration in the
Dial method lacks a MinVersion setting, which can allow insecure TLS versions
and downgrade attacks. Update the tls.Config to include MinVersion set to
tls.VersionTLS12 or higher to enforce a minimum TLS version for secure
connections.
| req.Header.Set("Accept", "application/json") | ||
| req.Header.Set("Content-Encoding", "gzip") | ||
|
|
||
| for currentAttempt := 0; currentAttempt <= s.maxRetries; currentAttempt++ { |
There was a problem hiding this comment.
I renamed retry to currentAttempt as it made more sense
Since this logic is guaranteed to run at least one time before any retry can occur, it made more sense to me to rename this
client/bloodhound/client.go
Outdated
| if body, err := rest.CopyBody(req); err != nil { | ||
| return nil, err | ||
| } else { | ||
| for currentAttept := 0; currentAttept <= s.maxRetries; currentAttept++ { |
There was a problem hiding this comment.
I renamed retry to currentAttempt as it made more sense
Since this logic is guaranteed to run at least one time before any retry can occur, it made more sense to me to rename this
| } | ||
|
|
||
| // BHEClient implements the BloodHoundClient interface to communicate with a BloodHound Enterprise instance | ||
| type BHEClient struct { |
There was a problem hiding this comment.
Moving the code from utils.go to here resolves this TODO
Line 271 in fc3ed0f
| @@ -0,0 +1,118 @@ | |||
| package rest | |||
There was a problem hiding this comment.
This code is unchanged, simply lifted and moved to a more appropriate place from utils.go
| @@ -0,0 +1,97 @@ | |||
| package bloodhound | |||
There was a problem hiding this comment.
This code is unchanged, other than exporting and adding the roundtripperFunc as it was moved from utils.go
| @@ -0,0 +1,22 @@ | |||
| package bloodhound | |||
There was a problem hiding this comment.
This was lifted from utils.go and remains unchanged
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (1)
client/bloodhound/client.go (1)
151-264: Consider usingSendRequestmethod to avoid duplicating retry logic.The
Ingestmethod reimplements retry logic that already exists inSendRequest. This creates code duplication and misses the connection reset functionality.Consider refactoring to use
SendRequestas suggested in previous reviews to eliminate code duplication and ensure consistent error handling across all HTTP operations.
🧹 Nitpick comments (1)
client/bloodhound/client.go (1)
44-83: Consider making retryDelay configurable.The struct design is solid and the constructor properly initializes all fields. The mutex addition addresses thread safety concerns from previous reviews.
Consider making the hardcoded
retryDelayof 5 seconds configurable:-func NewBHEClient(bheUrl url.URL, tokenId, token, proxy string, maxReqPerConn, maxRetries int, logger logr.Logger) (BloodHoundClient, error) { +func NewBHEClient(bheUrl url.URL, tokenId, token, proxy string, maxReqPerConn, maxRetries, retryDelay int, logger logr.Logger) (BloodHoundClient, error) { // ... existing code ... return &BHEClient{ // ... existing fields ... - retryDelay: 5, + retryDelay: retryDelay, // ... remaining fields ... }, nil }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
client/bloodhound/client.go(1 hunks)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: test
🔇 Additional comments (6)
client/bloodhound/client.go (6)
1-30: LGTM! Clean package structure and imports.The package declaration, imports, and constants are well-organized and follow Go best practices. The mock generation directive is properly placed.
31-42: LGTM! Well-designed interface.The
BloodHoundClientinterface provides a clean abstraction with appropriate method signatures for all BHE operations. This design enables easy testing and mocking.
85-149: LGTM! Comprehensive retry logic with proper error handling.The
SendRequestmethod implements robust retry logic with appropriate handling for different error types (closed connections, GOAWAY frames, server errors). The request body copying and thread-safe request counting are well-implemented.
266-383: LGTM! Consistent and well-implemented job management methods.All job management methods (
GetAvailableJobs,Checkin,StartJob,EndJob,UpdateClient,EndOrphanedJob) consistently use theSendRequestmethod, properly handle errors, and include appropriate resource cleanup.
385-388: LGTM! Clean delegation to underlying HTTP client.The
CloseIdleConnectionsmethod properly delegates to the underlying HTTP client's method.
390-427: LGTM! Thread safety issues properly addressed.Both
resetConnectionandincrementRequestmethods correctly implement thread safety using the mutex. The connection reset logic properly closes old connections before replacement, addressing the concerns from previous reviews.
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (2)
client/rest/dialer.go (1)
18-21: Address the TLS security vulnerability.The TLS configuration lacks a minimum version specification, which could allow connections using insecure TLS versions. This is a known security vulnerability.
client/bloodhound/client.go (1)
151-264: Consider refactoring to reduce code duplication.The Ingest method duplicates retry logic that already exists in SendRequest. This creates maintenance overhead and inconsistency.
Consider refactoring to use SendRequest for consistency:
- for currentAttempt := 0; currentAttempt <= s.maxRetries; currentAttempt++ { - // No retries on regular err cases, only on HTTP 504 Gateway Timeout and HTTP 503 Service Unavailable - response, err := s.httpClient.Do(req) - - if err != nil { - if rest.IsClosedConnectionErr(err) { - // try again on force closed connection - s.log.Error(err, fmt.Sprintf("remote host force closed connection while requesting %s; attempt %d/%d; trying again", req.URL, currentAttempt+1, s.maxRetries)) - - if currentAttempt == s.maxRetries { - s.log.Error(ErrExceededRetryLimit, "") - hasErrors = true - } else { - rest.VariableExponentialBackoff(s.retryDelay, currentAttempt) - } - - continue - } else if rest.IsGoAwayErr(err) { - // AWS currently has a 10,000 request per connection limitation, retry in case AWS changes this limitation - s.log.Error(err, fmt.Sprintf("received GOAWAY from from AWS load balancer while requesting %s; attempt %d/%d; trying again", req.URL, currentAttempt+1, s.maxRetries)) - - if currentAttempt == s.maxRetries { - s.log.Error(ErrExceededRetryLimit, "") - hasErrors = true - } else { - rest.VariableExponentialBackoff(s.retryDelay, currentAttempt) - } - - continue - } - - s.log.Error(err, unrecoverableErrMsg) - return true - } - - if err := s.incrementRequest(); err != nil { - return true - } - - if response.StatusCode == http.StatusGatewayTimeout || response.StatusCode == http.StatusServiceUnavailable || response.StatusCode == http.StatusBadGateway { - serverError := fmt.Errorf("received server error %d while requesting %v; attempt %d/%d; trying again", response.StatusCode, endpoint, currentAttempt+1, s.maxRetries) - s.log.Error(serverError, "") - - if currentAttempt == s.maxRetries { - s.log.Error(ErrExceededRetryLimit, "") - hasErrors = true - } else { - rest.VariableExponentialBackoff(s.retryDelay, currentAttempt) - } - - if err := response.Body.Close(); err != nil { - s.log.Error(fmt.Errorf("failed to close ingest body: %w", err), unrecoverableErrMsg) - } - - continue - } else if response.StatusCode != http.StatusAccepted { - if bodyBytes, err := io.ReadAll(response.Body); err != nil { - s.log.Error(fmt.Errorf("received unexpected response code from %v: %s; failure reading response body", endpoint, response.Status), unrecoverableErrMsg) - } else { - s.log.Error(fmt.Errorf("received unexpected response code from %v: %s %s", req.URL, response.Status, bodyBytes), unrecoverableErrMsg) - } - - if err := response.Body.Close(); err != nil { - s.log.Error(fmt.Errorf("failed to close ingest body: %w", err), unrecoverableErrMsg) - } - - return true - } - - if err := response.Body.Close(); err != nil { - s.log.Error(fmt.Errorf("failed to close ingest body: %w", err), unrecoverableErrMsg) - } - } + if response, err := s.SendRequest(req); err != nil { + s.log.Error(err, "error during ingest request") + hasErrors = true + } else { + defer response.Body.Close() + if response.StatusCode != http.StatusAccepted { + s.log.Error(fmt.Errorf("unexpected status code: %d", response.StatusCode), "ingest request failed") + hasErrors = true + } + }This would:
- Reuse existing retry logic
- Benefit from connection reset functionality
- Reduce code duplication significantly
- Simplify maintenance
🧹 Nitpick comments (1)
client/rest/dialer.go (1)
44-84: Improve error handling and resource management.The proxy connection logic has good error handling, but there are a few areas for improvement:
- The defer function for closing response body should handle potential nil pointer
- Resource cleanup could be more consistent
res, err := http.ReadResponse(bufio.NewReader(conn), req) - defer func() { - if res.Body != nil { - res.Body.Close() - } - }() if err != nil { conn.Close() return nil, fmt.Errorf("unable to connect to %s: %w", addr, err) } else if res.StatusCode != 200 { - if res.Body != nil { - res.Body.Close() - } + res.Body.Close() conn.Close() return nil, fmt.Errorf("unable to connect to %s via proxy (%s): statusCode %d", addr, s.host, res.StatusCode) } else { + res.Body.Close() return conn, nil }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
client/bloodhound/client.go(1 hunks)client/rest/dialer.go(1 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (1)
client/rest/dialer.go (1)
config/config.go (1)
Proxy(117-123)
🪛 ast-grep (0.38.6)
client/rest/dialer.go
[warning] 19-19: MinVersionis missing from this TLS configuration. By default, TLS 1.2 is currently used as the minimum when acting as a client, and TLS 1.0 when acting as a server. General purpose web applications should default to TLS 1.3 with all other protocols disabled. Only where it is known that a web server must support legacy clients with unsupported an insecure browsers (such as Internet Explorer 10), it may be necessary to enable TLS 1.0 to provide support. AddMinVersion: tls.VersionTLS13' to the TLS configuration to bump the minimum version to TLS 1.3.
Context: tls.Config{}
Note: [CWE-327]: Use of a Broken or Risky Cryptographic Algorithm [OWASP A03:2017]: Sensitive Data Exposure [OWASP A02:2021]: Cryptographic Failures [REFERENCES]
https://owasp.org/Top10/A02_2021-Cryptographic_Failures
(missing-ssl-minversion-go)
🔇 Additional comments (11)
client/rest/dialer.go (3)
23-35: LGTM: Clean proxy dialer constructor.The constructor properly handles URL parsing and credential extraction. The error handling is appropriate.
86-96: LGTM: Proper proxy configuration handling.The function correctly handles different proxy schemes and falls back to direct connection when no proxy is configured.
98-119: LGTM: Clean connection establishment with proper cleanup.The Dial function correctly handles URL parsing, port defaulting, and connection cleanup. The logging level is appropriate.
client/bloodhound/client.go (8)
25-30: LGTM: Clear constants and error definitions.The BHE authentication signature constant and retry limit error are well-defined.
31-42: Well-designed interface for BloodHound client operations.The interface provides a clean abstraction for all necessary BloodHound operations with appropriate method signatures.
44-57: LGTM: Comprehensive client structure with proper synchronization.The BHEClient struct includes all necessary fields and the mutex for thread-safe request counting. Good to see the race condition concerns from previous reviews have been addressed.
59-83: Constructor properly initializes client with sensible defaults.The constructor correctly sets up the HTTP client with signing transport and provides reasonable default values like 5-second retry delay.
85-149: Robust request handling with comprehensive retry logic.The SendRequest method implements proper retry logic for different error conditions:
- Handles closed connections and GOAWAY errors appropriately
- Implements exponential backoff for server errors
- Properly manages request body rewinding for retries
- Thread-safe request counting
The error handling covers all relevant HTTP status code ranges and provides detailed error messages.
266-285: LGTM: Clean job retrieval with proper error handling.The GetAvailableJobs method correctly uses the SendRequest method and handles JSON decoding properly.
287-383: LGTM: Comprehensive job management methods.All job management methods (Checkin, StartJob, EndJob, UpdateClient, EndOrphanedJob) follow consistent patterns:
- Proper endpoint construction
- Use of SendRequest for consistency
- Appropriate error handling and resource cleanup
- Good logging where needed
385-428: LGTM: Thread-safe connection management.The connection management methods properly handle:
- Closing idle connections
- Thread-safe connection reset with mutex protection
- Thread-safe request counting
- Proper cleanup of old HTTP client before replacement
Good to see the thread safety issues from previous reviews have been resolved.
cweidenkeller
left a comment
cweidenkeller
left a comment
There was a problem hiding this comment.
Had a video peer review session and everything looks good to me =)
2 participants
This PR fixes a bug with AWS' load balancer sending a GOAWAY frame when sending more than 10,000 requests over a single connection
When AzureHound sends a user definable amount of requests to the BloodHound server, it will reset the HTTP client to establish a new connection, resetting the limit.
As part of this work, there was a
TODOto move the BHE client logic into a new client structure. This PR resolves thatTODOby creating aBloodHoundClientinterface and aBHEClientstructure that implements the interface.Most of the code changes here are simply re-structuring the code and moving them to the structure
Testing
I ran an on-demand collection off of
mainand took screenshots of both the Attack Paths page and Posture pages, screenshot the results of the collectionI then ran an on-demand collection off of this branch, and repeated the process
I was then able to compare the numbers between the two and confirmed that the collection stats all matched
Summary by CodeRabbit
New Features
Bug Fixes
Tests
Chores