forked from opensearch-project/opensearch-migrations
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Experimenting with handlers to add and remove headers from http reque…
…sts. Signed-off-by: Greg Schohn <[email protected]>
- Loading branch information
1 parent
bce5d3d
commit d8acd47
Showing
5 changed files
with
330 additions
and
0 deletions.
There are no files selected for viewing
41 changes: 41 additions & 0 deletions
41
...n/java/org/opensearch/migrations/trafficcapture/proxyserver/netty/HeaderAdderHandler.java
This file contains 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
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,41 @@ | ||
package org.opensearch.migrations.trafficcapture.proxyserver.netty; | ||
|
||
import io.netty.buffer.ByteBuf; | ||
import io.netty.buffer.Unpooled; | ||
import io.netty.channel.ChannelHandlerContext; | ||
import io.netty.channel.ChannelInboundHandlerAdapter; | ||
|
||
public class HeaderAdderHandler extends ChannelInboundHandlerAdapter { | ||
boolean insertedHeader = false; | ||
private final ByteBuf headerLineToAdd; | ||
|
||
public HeaderAdderHandler(ByteBuf headerLineToAdd) { | ||
this.headerLineToAdd = headerLineToAdd.retain(); | ||
} | ||
|
||
@Override | ||
public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception { | ||
if (!(msg instanceof ByteBuf) || insertedHeader) { | ||
super.channelRead(ctx, msg); | ||
return; | ||
} | ||
var buf = (ByteBuf) msg; | ||
buf.markReaderIndex(); | ||
while (buf.isReadable()) { | ||
if (buf.readByte() == '\n') { | ||
final var upToIndex = buf.readerIndex(); | ||
var composite = Unpooled.compositeBuffer(3); | ||
buf.resetReaderIndex(); | ||
composite.addComponent(true, buf.retainedSlice(0, upToIndex)); | ||
composite.addComponent(true, headerLineToAdd.duplicate()); | ||
composite.addComponent(true, buf.retainedSlice(upToIndex, buf.readableBytes()-upToIndex)); | ||
buf.release(); | ||
super.channelRead(ctx, composite); | ||
insertedHeader = true; | ||
return; | ||
} | ||
} | ||
buf.resetReaderIndex(); | ||
super.channelRead(ctx, msg); | ||
} | ||
} |
111 changes: 111 additions & 0 deletions
111
...java/org/opensearch/migrations/trafficcapture/proxyserver/netty/HeaderRemoverHandler.java
This file contains 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
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,111 @@ | ||
package org.opensearch.migrations.trafficcapture.proxyserver.netty; | ||
|
||
import java.util.List; | ||
|
||
import io.netty.buffer.ByteBuf; | ||
import io.netty.buffer.CompositeByteBuf; | ||
import io.netty.buffer.Unpooled; | ||
import io.netty.channel.ChannelHandlerContext; | ||
import io.netty.channel.ChannelInboundHandlerAdapter; | ||
import lombok.SneakyThrows; | ||
|
||
public class HeaderRemoverHandler extends ChannelInboundHandlerAdapter { | ||
final String headerToRemove; | ||
CompositeByteBuf previousRemaining; | ||
// This handler has 3 states - copying, dropping, or testing. when previousRemaining != null, we're testing. | ||
// when dropUntilNewline == true, we're dropping, otherwise, we're copying (when previousRemaining==null) | ||
// The starting state is previousRemaining == null and dropUntilNewline = false | ||
boolean dropUntilNewline; | ||
|
||
public HeaderRemoverHandler(String headerToRemove) { | ||
if (!headerToRemove.endsWith(":")) { | ||
throw new IllegalArgumentException("The headerToRemove must end with a ':'"); | ||
} | ||
this.headerToRemove = headerToRemove; | ||
} | ||
|
||
@SneakyThrows | ||
void lambdaSafeSuperChannelRead(ChannelHandlerContext ctx, ByteBuf bb) { | ||
super.channelRead(ctx, bb); | ||
} | ||
|
||
/** | ||
* @return true if there's a discongruity in the incoming buf and the contents that preceded this call will | ||
* need to be buffered by the caller | ||
*/ | ||
boolean matchNextBytes(ChannelHandlerContext ctx, ByteBuf buf) { | ||
if (!buf.isReadable()) { | ||
return false; | ||
} | ||
buf.markReaderIndex(); | ||
for (int i=previousRemaining.readerIndex(); ; ++i) { | ||
if (!buf.isReadable()) { // partial match | ||
previousRemaining.addComponent(true, buf); | ||
return true; | ||
} | ||
if (i == headerToRemove.length()) { // match! | ||
previousRemaining.release(); // drop those in limbo ... | ||
previousRemaining = null; | ||
dropUntilNewline = true; // ... plus other bytes until we reset | ||
return true; | ||
} | ||
if (Character.toLowerCase(headerToRemove.charAt(i)) != Character.toLowerCase(buf.readByte())) { // no match | ||
previousRemaining.forEach(bb -> lambdaSafeSuperChannelRead(ctx, bb)); | ||
previousRemaining = null; | ||
dropUntilNewline = false; | ||
return false; | ||
} | ||
} | ||
|
||
} | ||
|
||
boolean advanceByteBufUntilNewline(ByteBuf bb) { | ||
while (bb.isReadable()) { // sonar lint doesn't like if the while statement has an empty body | ||
if (bb.readByte() != '\n') { return true; } | ||
} | ||
return false; | ||
} | ||
|
||
@Override | ||
public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception { | ||
if (!(msg instanceof ByteBuf)) { | ||
super.channelRead(ctx, msg); | ||
return; | ||
} | ||
|
||
var sourceBuf = (ByteBuf) msg; | ||
var startForNextSourceSegment = sourceBuf.readerIndex(); | ||
var cleanedIncomingBuf = ctx.alloc().compositeBuffer(4); | ||
|
||
while (true) { | ||
if (previousRemaining != null) { | ||
final var sourceReaderIdx = sourceBuf.readerIndex(); | ||
if (matchNextBytes(ctx, sourceBuf.slice(sourceReaderIdx, sourceBuf.readableBytes())) && | ||
sourceReaderIdx != startForNextSourceSegment) // would be 0-length | ||
{ | ||
cleanedIncomingBuf.addComponent(true, | ||
sourceBuf.retainedSlice(startForNextSourceSegment, sourceReaderIdx)); | ||
startForNextSourceSegment = -1; | ||
} | ||
} else { | ||
var foundNewline = advanceByteBufUntilNewline(sourceBuf); | ||
if (dropUntilNewline) { | ||
if (foundNewline) { | ||
// took care of previous bytes in the source buffer in the previousRemaining != null branch | ||
startForNextSourceSegment = sourceBuf.readerIndex(); | ||
} | ||
} | ||
if (foundNewline) { | ||
previousRemaining = ctx.alloc().compositeBuffer(16); | ||
} else { | ||
break; | ||
} | ||
} | ||
} | ||
if (startForNextSourceSegment >= 0) { | ||
cleanedIncomingBuf.addComponent(true, | ||
sourceBuf.retainedSlice(startForNextSourceSegment, sourceBuf.readerIndex()-startForNextSourceSegment)); | ||
} | ||
super.channelRead(ctx, cleanedIncomingBuf); | ||
} | ||
} |
62 changes: 62 additions & 0 deletions
62
...va/org/opensearch/migrations/trafficcapture/proxyserver/netty/HeaderAdderHandlerTest.java
This file contains 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
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,62 @@ | ||
package org.opensearch.migrations.trafficcapture.proxyserver.netty; | ||
|
||
import java.nio.charset.StandardCharsets; | ||
import java.util.stream.Stream; | ||
|
||
import io.netty.buffer.ByteBuf; | ||
import io.netty.buffer.Unpooled; | ||
import io.netty.channel.embedded.EmbeddedChannel; | ||
import org.junit.jupiter.api.Assertions; | ||
import org.junit.jupiter.api.Test; | ||
import org.junit.jupiter.params.ParameterizedTest; | ||
import org.junit.jupiter.params.provider.Arguments; | ||
import org.junit.jupiter.params.provider.MethodSource; | ||
|
||
import static org.junit.jupiter.api.Assertions.*; | ||
|
||
class HeaderAdderHandlerTest { | ||
|
||
public static Stream<Arguments> makeArgs() { | ||
return Stream.of( | ||
Arguments.of("\n"), | ||
Arguments.of("\r\n")); | ||
} | ||
|
||
@ParameterizedTest | ||
@MethodSource("makeArgs") | ||
public void simpleCheck(String lineEnding) { | ||
var extraHeader = "host: my.host\n"; | ||
var newHeader = Unpooled.wrappedBuffer(extraHeader.getBytes(StandardCharsets.UTF_8)); | ||
final var msg = makeMessage(lineEnding, ""); | ||
|
||
var channel = new EmbeddedChannel(new HeaderAdderHandler(newHeader)); | ||
channel.writeInbound(Unpooled.wrappedBuffer(msg.getBytes(StandardCharsets.UTF_8))); | ||
var output = Unpooled.compositeBuffer(); | ||
channel.inboundMessages().forEach(v -> output.addComponent(true, (ByteBuf) v)); | ||
|
||
Assertions.assertEquals(makeMessage(lineEnding, extraHeader), output.toString(StandardCharsets.UTF_8)); | ||
} | ||
|
||
@ParameterizedTest | ||
@MethodSource("makeArgs") | ||
public void fragmentedCheck(String lineEnding) { | ||
var extraHeader = "host: my.host\n"; | ||
var newHeader = Unpooled.wrappedBuffer(extraHeader.getBytes(StandardCharsets.UTF_8)); | ||
final var msg = makeMessage(lineEnding, ""); | ||
|
||
var channel = new EmbeddedChannel(new HeaderAdderHandler(newHeader)); | ||
msg.chars().forEach(c -> channel.writeInbound(Unpooled.wrappedBuffer(new byte[]{(byte) c}))); | ||
var output = Unpooled.compositeBuffer(); | ||
channel.inboundMessages().forEach(v -> output.addComponent(true, (ByteBuf) v)); | ||
|
||
Assertions.assertEquals(makeMessage(lineEnding, extraHeader), output.toString(StandardCharsets.UTF_8)); | ||
} | ||
|
||
String makeMessage(String lineEnding, String extraHeader) { | ||
return "GET / HTTP/1.1" + lineEnding + | ||
extraHeader + | ||
"NICEHeader: v1" + lineEnding + | ||
"silLYHeader: yyy" + lineEnding + | ||
lineEnding; | ||
} | ||
} |
59 changes: 59 additions & 0 deletions
59
.../org/opensearch/migrations/trafficcapture/proxyserver/netty/HeaderRemoverHandlerTest.java
This file contains 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
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,59 @@ | ||
package org.opensearch.migrations.trafficcapture.proxyserver.netty; | ||
|
||
import java.nio.charset.StandardCharsets; | ||
import java.util.stream.Stream; | ||
|
||
import io.netty.buffer.ByteBuf; | ||
import io.netty.buffer.Unpooled; | ||
import io.netty.channel.embedded.EmbeddedChannel; | ||
import org.junit.jupiter.api.Assertions; | ||
import org.junit.jupiter.params.ParameterizedTest; | ||
import org.junit.jupiter.params.provider.Arguments; | ||
import org.junit.jupiter.params.provider.MethodSource; | ||
|
||
class HeaderRemoverHandlerTest { | ||
|
||
public static Stream<Arguments> makeArgs() { | ||
return Stream.of( | ||
Arguments.of("\n"), | ||
Arguments.of("\r\n")); | ||
} | ||
|
||
@ParameterizedTest | ||
@MethodSource("makeArgs") | ||
public void simpleCheck(String lineEnding) { | ||
var extraHeader = "host: my.host\n"; | ||
var newHeader = Unpooled.wrappedBuffer(extraHeader.getBytes(StandardCharsets.UTF_8)); | ||
final var msg = makeMessage(lineEnding, true); | ||
|
||
var channel = new EmbeddedChannel(new HeaderAdderHandler(newHeader)); | ||
channel.writeInbound(Unpooled.wrappedBuffer(msg.getBytes(StandardCharsets.UTF_8))); | ||
var output = Unpooled.compositeBuffer(); | ||
channel.inboundMessages().forEach(v -> output.addComponent(true, (ByteBuf) v)); | ||
|
||
Assertions.assertEquals(makeMessage(lineEnding, false), output.toString(StandardCharsets.UTF_8)); | ||
} | ||
|
||
@ParameterizedTest | ||
@MethodSource("makeArgs") | ||
public void fragmentedCheck(String lineEnding) { | ||
var headerToRemove = "host"; | ||
final var msg = makeMessage(lineEnding, true); | ||
|
||
var channel = new EmbeddedChannel(new HeaderRemoverHandler(headerToRemove)); | ||
msg.chars().forEach(c -> channel.writeInbound(Unpooled.wrappedBuffer(new byte[]{(byte) c}))); | ||
var output = Unpooled.compositeBuffer(); | ||
channel.inboundMessages().forEach(v -> output.addComponent(true, (ByteBuf) v)); | ||
|
||
Assertions.assertEquals(makeMessage(lineEnding, false), output.toString(StandardCharsets.UTF_8)); | ||
} | ||
|
||
String makeMessage(String lineEnding, boolean withHosts) { | ||
return "GET / HTTP/1.1" + lineEnding + | ||
"hoststays: v1" + lineEnding + | ||
(withHosts ? ("HOST: begone" + lineEnding) : "") + | ||
"different: v1" + lineEnding + | ||
(withHosts ? ("HosT: begone" + lineEnding) : "") + | ||
lineEnding; | ||
} | ||
} |
57 changes: 57 additions & 0 deletions
57
...src/test/java/org/opensearch/migrations/trafficcapture/proxyserver/netty/MatcherTest.java
This file contains 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
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,57 @@ | ||
package org.opensearch.migrations.trafficcapture.proxyserver.netty; | ||
|
||
import java.nio.charset.StandardCharsets; | ||
import java.time.Duration; | ||
import java.util.regex.Pattern; | ||
|
||
import com.google.common.base.Strings; | ||
import io.netty.buffer.ByteBuf; | ||
import io.netty.buffer.Unpooled; | ||
import lombok.extern.slf4j.Slf4j; | ||
import org.junit.jupiter.api.Assertions; | ||
import org.junit.jupiter.api.Test; | ||
|
||
@Slf4j | ||
public class MatcherTest { | ||
|
||
public static final ByteBuf BIG_BUF = | ||
Unpooled.wrappedBuffer(Strings.repeat("ha", 100_000).getBytes(StandardCharsets.UTF_8)); | ||
public static final ByteBuf SMALL_BUF = | ||
Unpooled.wrappedBuffer(Strings.repeat("ha", 1).getBytes(StandardCharsets.UTF_8)); | ||
|
||
@Test | ||
public void test() { | ||
var p = Pattern.compile("^host:.*", Pattern.CASE_INSENSITIVE); | ||
|
||
Assertions.assertTrue( | ||
bufMatches(p, Unpooled.wrappedBuffer("host: MYHOST".getBytes(StandardCharsets.UTF_8)))); | ||
|
||
getMatchTime(p, BIG_BUF, 1000); | ||
getMatchTime(p, BIG_BUF, 1000); | ||
|
||
for (int i=0; i<1; ++i) { | ||
final var MATCH_REPS = 100_000_000; | ||
var smallTime = getMatchTime(p, SMALL_BUF, MATCH_REPS); | ||
var bigTime = getMatchTime(p, BIG_BUF, MATCH_REPS); | ||
log.info("smallTime = "+smallTime); | ||
log.info("bigTime = "+bigTime); | ||
} | ||
} | ||
|
||
private static Duration getMatchTime(Pattern p, ByteBuf input, int i) { | ||
final var start = System.nanoTime(); | ||
boolean didMatch = false; | ||
for (; i > 0; --i) { | ||
didMatch |= bufMatches(p, input); | ||
} | ||
try { | ||
return Duration.ofNanos(System.nanoTime() - start); | ||
} finally { | ||
Assertions.assertFalse(didMatch); | ||
} | ||
} | ||
|
||
public static boolean bufMatches(Pattern p, ByteBuf b) { | ||
return p.matcher(b.getCharSequence(0, b.readableBytes(),StandardCharsets.UTF_8)).matches(); | ||
} | ||
} |