fix: interaction data corruption causing unmarshal errors#1341
Open
dogancanbakir wants to merge 1 commit intodevfrom
Open
fix: interaction data corruption causing unmarshal errors#1341dogancanbakir wants to merge 1 commit intodevfrom
dogancanbakir wants to merge 1 commit intodevfrom
Conversation
Neo - PR Security ReviewNo security issues found Highlights
Hardening Notes
Comment |
WalkthroughThis PR refactors JSON encoding across server modules by replacing Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (8)
pkg/server/smb_server.go (1)
103-111: Encoding change is correct.The
jsoniter.Marshal()approach properly avoids the trailing newline issue.Same minor nit as responder_server.go: Line 109 has an incorrect error message saying "dns interaction" instead of "smb interaction".
🔧 Fix error message
if err := h.options.Storage.AddInteractionWithId(h.options.Token, data); err != nil { - gologger.Warning().Msgf("Could not store dns interaction: %s\n", err) + gologger.Warning().Msgf("Could not store smb interaction: %s\n", err) }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@pkg/server/smb_server.go` around lines 103 - 111, The warning message uses the wrong protocol name; when jsoniter.Marshal(interaction) succeeds and h.options.Storage.AddInteractionWithId(h.options.Token, data) returns an error, update the gologger.Warning().Msgf call to reference "smb interaction" instead of "dns interaction" (i.e., adjust the message text in the AddInteractionWithId error handling branch so it correctly says "Could not store smb interaction: %s\n").pkg/server/responder_server.go (1)
91-99: Encoding change looks correct.The switch from
jsoniter.NewEncoder().Encode()tojsoniter.Marshal()properly eliminates the trailing newline that was causing unmarshal errors. TheMarshalfunction returns compact JSON without a trailing newline.Minor nit: Line 97 has an incorrect error message saying "dns interaction" instead of "responder interaction".
🔧 Fix error message
if err := h.options.Storage.AddInteractionWithId(h.options.Token, data); err != nil { - gologger.Warning().Msgf("Could not store dns interaction: %s\n", err) + gologger.Warning().Msgf("Could not store responder interaction: %s\n", err) }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@pkg/server/responder_server.go` around lines 91 - 99, Update the log message string in the interaction storage error to correctly refer to "responder interaction" instead of "dns interaction": locate the block that marshals the interaction (calls jsoniter.Marshal(interaction)) and then calls h.options.Storage.AddInteractionWithId(h.options.Token, data); change the gologger.Warning().Msgf call that currently says "Could not store dns interaction: %s" to "Could not store responder interaction: %s" so the message accurately reflects the operation and use the same error variable.pkg/storage/roundtrip_test.go (5)
173-178: Same encoding inconsistency in disk test.Same suggestion as above: consider using
jsoniter.Marshal()to match production code.♻️ Suggested change
- buffer := &bytes.Buffer{} - err = jsoniter.NewEncoder(buffer).Encode(inter) - require.NoError(t, err, "encode interaction %d", i) - - err = db.AddInteraction(correlationID, buffer.Bytes()) + data, err := jsoniter.Marshal(inter) + require.NoError(t, err, "encode interaction %d", i) + + err = db.AddInteraction(correlationID, data)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@pkg/storage/roundtrip_test.go` around lines 173 - 178, The test uses jsoniter.NewEncoder(buffer).Encode(inter) to produce the bytes fed into db.AddInteraction(correlationID, buffer.Bytes()), causing encoding differences from production; replace the encoder usage with jsoniter.Marshal(inter) and pass that result into AddInteraction so the test serializes interactions exactly like production code (look for variables/identifiers buffer, inter, jsoniter.NewEncoder, and db.AddInteraction).
119-124: Test usesEncoder.Encode()but production code now usesMarshal().The production code (e.g.,
http_server.go,dns_server.go) was changed to usejsoniter.Marshal()to avoid trailing newlines, but this test still usesjsoniter.NewEncoder(buffer).Encode(). WhileUnmarshaltolerates trailing newlines, using the same encoding method as production would make this test more representative of the actual code path.♻️ Suggested change to match production code
- buffer := &bytes.Buffer{} - err = jsoniter.NewEncoder(buffer).Encode(inter) - require.NoError(t, err, "encode interaction %d", i) - - err = mem.AddInteraction(correlationID, buffer.Bytes()) + data, err := jsoniter.Marshal(inter) + require.NoError(t, err, "encode interaction %d", i) + + err = mem.AddInteraction(correlationID, data)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@pkg/storage/roundtrip_test.go` around lines 119 - 124, The test currently uses jsoniter.NewEncoder(buffer).Encode(inter) which writes a trailing newline unlike production; update the test to use jsoniter.Marshal(inter) and pass the resulting bytes to mem.AddInteraction(correlationID, data) (handling the returned bytes and error inline) so the test encodes interactions the same way as production code paths that use jsoniter.Marshal().
353-358: Test relies on timing for cache eviction - potential flakiness.The test sleeps for 200ms to wait for cache eviction with a 100ms TTL. While this should generally work, timing-based tests can be flaky under CI load. The assertion
require.False(t, found, "cache entry should be evicted")helps catch this, but if it becomes flaky, consider:
- Increasing the sleep margin
- Using a retry loop to poll for eviction
This is minor since the current margin (2x the TTL) is reasonable.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@pkg/storage/roundtrip_test.go` around lines 353 - 358, The test in roundtrip_test.go relies on time.Sleep(200 * time.Millisecond) to wait for a 100ms TTL eviction which can be flaky; replace this fixed sleep with a short polling/retry loop that calls db.cache.GetIfPresent(correlationID) until it returns not found or a reasonable timeout elapses (e.g., 500–1000ms), or alternatively increase the sleep margin; update the assertion to fail only after the timeout so GetIfPresent is polled rather than assumed evicted after a single sleep.
230-234: Same encoding inconsistency in PollResponse test.Consider using
jsoniter.Marshal()here as well for consistency with production.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@pkg/storage/roundtrip_test.go` around lines 230 - 234, In the PollResponse test in roundtrip_test.go, replace the jsoniter.NewEncoder usage with jsoniter.Marshal for consistency with production: instead of creating buffer := &bytes.Buffer{} and calling jsoniter.NewEncoder(buffer).Encode(inter) then mem.AddInteraction(correlationID, buffer.Bytes()), call jsoniter.Marshal(inter) to get the payload bytes and pass those to mem.AddInteraction(correlationID, payload) so the test uses jsoniter.Marshal like other tests.
261-286: Test validates encoder behavior but production no longer uses encoder.This test explicitly verifies that
Encoder.Encode()appends a trailing newline. However, since production code now usesMarshal()(which doesn't append a newline), this test documents historical behavior rather than current production behavior.Consider either:
- Updating the test to verify
Marshal()does not append a trailing newline (the actual production behavior)- Renaming/documenting it as a regression test to ensure
Unmarshaltolerates trailing newlines if encountered🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@pkg/storage/roundtrip_test.go` around lines 261 - 286, TestTrailingNewlineHandling currently asserts Encoder.Encode appends a trailing newline but production uses jsoniter.Marshal; update the test to call jsoniter.Marshal on the interaction (or otherwise use Marshal instead of Encoder.Encode), assert that the resulting byte slice does NOT end with '\n' (reflecting current production behavior), and then (optionally) keep a small regression subtest that feeds an encoder-produced payload to jsoniter.Unmarshal to ensure Unmarshal tolerates trailing newlines; reference the interaction struct, jsoniter.Marshal/jsoniter.Unmarshal and TestTrailingNewlineHandling when making changes.pkg/client/client.go (1)
456-468: Whitespace trimming is fine, but likely redundant vs JSON whitespace handling; consider simplifying and/or clarifying intent.
bytes.TrimRight(plaintext, " \t\r\n")is safe as a defensive normalization, but it’s worth confirming whether the client actually needs this now that the server moved away fromEncoder.Encode()(newline) and because JSON parsers typically tolerate trailing whitespace anyway. If the goal is purely “drop encoder newline”, consider using the idiomaticbytes.TrimSpacefor clarity.If you still need to tolerate non-whitespace trailing bytes, this approach won’t address that—so I’d either (a) keep this strictly as “whitespace normalization” and add a short comment, or (b) add an explicit fallback path with clear logging so corruption isn’t silently masked.
♻️ Possible simplification
- plaintext = bytes.TrimRight(plaintext, " \t\r\n") + // Normalize trailing whitespace (e.g. legacy encoder newline) before JSON unmarshal. + plaintext = bytes.TrimSpace(plaintext)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@pkg/client/client.go` around lines 456 - 468, The bytes.TrimRight call in the loop over response.Data is defensive whitespace normalization but is redundant given JSON parsers tolerate trailing whitespace; replace bytes.TrimRight(plaintext, " \t\r\n") with bytes.TrimSpace(plaintext) and add a one-line comment above it explaining this is purely defensive whitespace normalization (not an attempt to mask non-whitespace corruption). Keep the rest of the flow (decryptMessage, jsoniter.Unmarshal, callback) unchanged; if you actually need to tolerate non-whitespace trailing bytes instead, add an explicit fallback branch that logs a warning before attempting a more lenient parse so corruption isn’t silently masked.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@pkg/client/client.go`:
- Around line 456-468: The bytes.TrimRight call in the loop over response.Data
is defensive whitespace normalization but is redundant given JSON parsers
tolerate trailing whitespace; replace bytes.TrimRight(plaintext, " \t\r\n") with
bytes.TrimSpace(plaintext) and add a one-line comment above it explaining this
is purely defensive whitespace normalization (not an attempt to mask
non-whitespace corruption). Keep the rest of the flow (decryptMessage,
jsoniter.Unmarshal, callback) unchanged; if you actually need to tolerate
non-whitespace trailing bytes instead, add an explicit fallback branch that logs
a warning before attempting a more lenient parse so corruption isn’t silently
masked.
In `@pkg/server/responder_server.go`:
- Around line 91-99: Update the log message string in the interaction storage
error to correctly refer to "responder interaction" instead of "dns
interaction": locate the block that marshals the interaction (calls
jsoniter.Marshal(interaction)) and then calls
h.options.Storage.AddInteractionWithId(h.options.Token, data); change the
gologger.Warning().Msgf call that currently says "Could not store dns
interaction: %s" to "Could not store responder interaction: %s" so the message
accurately reflects the operation and use the same error variable.
In `@pkg/server/smb_server.go`:
- Around line 103-111: The warning message uses the wrong protocol name; when
jsoniter.Marshal(interaction) succeeds and
h.options.Storage.AddInteractionWithId(h.options.Token, data) returns an error,
update the gologger.Warning().Msgf call to reference "smb interaction" instead
of "dns interaction" (i.e., adjust the message text in the AddInteractionWithId
error handling branch so it correctly says "Could not store smb interaction:
%s\n").
In `@pkg/storage/roundtrip_test.go`:
- Around line 173-178: The test uses jsoniter.NewEncoder(buffer).Encode(inter)
to produce the bytes fed into db.AddInteraction(correlationID, buffer.Bytes()),
causing encoding differences from production; replace the encoder usage with
jsoniter.Marshal(inter) and pass that result into AddInteraction so the test
serializes interactions exactly like production code (look for
variables/identifiers buffer, inter, jsoniter.NewEncoder, and
db.AddInteraction).
- Around line 119-124: The test currently uses
jsoniter.NewEncoder(buffer).Encode(inter) which writes a trailing newline unlike
production; update the test to use jsoniter.Marshal(inter) and pass the
resulting bytes to mem.AddInteraction(correlationID, data) (handling the
returned bytes and error inline) so the test encodes interactions the same way
as production code paths that use jsoniter.Marshal().
- Around line 353-358: The test in roundtrip_test.go relies on time.Sleep(200 *
time.Millisecond) to wait for a 100ms TTL eviction which can be flaky; replace
this fixed sleep with a short polling/retry loop that calls
db.cache.GetIfPresent(correlationID) until it returns not found or a reasonable
timeout elapses (e.g., 500–1000ms), or alternatively increase the sleep margin;
update the assertion to fail only after the timeout so GetIfPresent is polled
rather than assumed evicted after a single sleep.
- Around line 230-234: In the PollResponse test in roundtrip_test.go, replace
the jsoniter.NewEncoder usage with jsoniter.Marshal for consistency with
production: instead of creating buffer := &bytes.Buffer{} and calling
jsoniter.NewEncoder(buffer).Encode(inter) then mem.AddInteraction(correlationID,
buffer.Bytes()), call jsoniter.Marshal(inter) to get the payload bytes and pass
those to mem.AddInteraction(correlationID, payload) so the test uses
jsoniter.Marshal like other tests.
- Around line 261-286: TestTrailingNewlineHandling currently asserts
Encoder.Encode appends a trailing newline but production uses jsoniter.Marshal;
update the test to call jsoniter.Marshal on the interaction (or otherwise use
Marshal instead of Encoder.Encode), assert that the resulting byte slice does
NOT end with '\n' (reflecting current production behavior), and then
(optionally) keep a small regression subtest that feeds an encoder-produced
payload to jsoniter.Unmarshal to ensure Unmarshal tolerates trailing newlines;
reference the interaction struct, jsoniter.Marshal/jsoniter.Unmarshal and
TestTrailingNewlineHandling when making changes.
ℹ️ Review info
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (1)
go.sumis excluded by!**/*.sum
📒 Files selected for processing (12)
go.modpkg/client/client.gopkg/server/dns_server.gopkg/server/ftp_server.gopkg/server/http_server.gopkg/server/http_server_test.gopkg/server/ldap_server.gopkg/server/responder_server.gopkg/server/smb_server.gopkg/server/smtp_server.gopkg/storage/roundtrip_test.gopkg/storage/storagedb.go
knakul853
approved these changes
Feb 26, 2026
|
@dogancanbakir TestFullRoundTripInMemory and TestFullRoundTripDisk these test still using older decoder logic good to update to new logic |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #1340
OnCacheRemovalCallbacktype assertion (value.([]byte)→key.(string)) so LevelDB entries are cleaned on cache evictionSetIDPublicKeybefore creating new cache entry on re-registrationjsoniter.NewEncoder().Encode()withjsoniter.Marshal()across all server protocol handlers to avoid trailing\nin encrypted databytes.TrimRightin client before unmarshalSummary by CodeRabbit
Bug Fixes
Tests
Chores