Added custom http headers to MqttWebSocket (#500)

Co-authored-by: preto <utopreto@gmail.com>

Author: hctemmel

src/main/java/com/hivemq/client/internal/mqtt/MqttWebSocketConfigImpl.java

@@ -20,6 +20,9 @@
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
+import java.util.Map;
+import java.util.Objects;
+
 /**
  * @author David Katz
  * @author Christian Hoff
@@ -28,23 +31,26 @@ public class MqttWebSocketConfigImpl implements MqttWebSocketConfig {
 
     static final @NotNull MqttWebSocketConfigImpl DEFAULT =
             new MqttWebSocketConfigImpl(DEFAULT_SERVER_PATH, DEFAULT_QUERY_STRING, DEFAULT_MQTT_SUBPROTOCOL,
-                    DEFAULT_HANDSHAKE_TIMEOUT_MS);
+                    DEFAULT_HANDSHAKE_TIMEOUT_MS, DEFAULT_HTTP_HEADERS);
 
     private final @NotNull String serverPath;
     private final @NotNull String queryString;
     private final @NotNull String subprotocol;
     private final int handshakeTimeoutMs;
+    private final Map<String, String> httpHeaders;
 
     MqttWebSocketConfigImpl(
             final @NotNull String serverPath,
             final @NotNull String queryString,
             final @NotNull String subprotocol,
-            final int handshakeTimeoutMs) {
+            final int handshakeTimeoutMs,
+            final @NotNull Map<String, String> httpHeaders) {
 
         this.serverPath = serverPath;
         this.queryString = queryString;
         this.subprotocol = subprotocol;
         this.handshakeTimeoutMs = handshakeTimeoutMs;
+        this.httpHeaders = httpHeaders;
     }
 
     @Override
@@ -67,6 +73,11 @@ public int getHandshakeTimeoutMs() {
         return handshakeTimeoutMs;
     }
 
+    @Override
+    public @NotNull Map<String, String> getHttpHeaders() {
+        return httpHeaders;
+    }
+
     @Override
     public MqttWebSocketConfigImplBuilder.@NotNull Default extend() {
         return new MqttWebSocketConfigImplBuilder.Default(this);
@@ -83,7 +94,8 @@ public boolean equals(final @Nullable Object o) {
         final MqttWebSocketConfigImpl that = (MqttWebSocketConfigImpl) o;
 
         return serverPath.equals(that.serverPath) && queryString.equals(that.queryString) &&
-                subprotocol.equals(that.subprotocol) && (handshakeTimeoutMs == that.handshakeTimeoutMs);
+                subprotocol.equals(that.subprotocol) && (handshakeTimeoutMs == that.handshakeTimeoutMs) &&
+                Objects.equals(httpHeaders, that.httpHeaders);
     }
 
     @Override
@@ -92,6 +104,7 @@ public int hashCode() {
         result = 31 * result + queryString.hashCode();
         result = 31 * result + subprotocol.hashCode();
         result = 31 * result + Integer.hashCode(handshakeTimeoutMs);
+        result = 31 * result + httpHeaders.hashCode();
         return result;
     }
 }

src/main/java/com/hivemq/client/internal/mqtt/MqttWebSocketConfigImplBuilder.java

@@ -21,6 +21,7 @@
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
+import java.util.Map;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Function;
 
@@ -33,6 +34,7 @@ public abstract class MqttWebSocketConfigImplBuilder<B extends MqttWebSocketConf
     private @NotNull String queryString = MqttWebSocketConfigImpl.DEFAULT_QUERY_STRING;
     private @NotNull String subprotocol = MqttWebSocketConfigImpl.DEFAULT_MQTT_SUBPROTOCOL;
     private int handshakeTimeoutMs = MqttWebSocketConfigImpl.DEFAULT_HANDSHAKE_TIMEOUT_MS;
+    private @NotNull Map<String, String> httpHeaders = MqttWebSocketConfigImpl.DEFAULT_HTTP_HEADERS;
 
     MqttWebSocketConfigImplBuilder() {}
 
@@ -70,8 +72,14 @@ public abstract class MqttWebSocketConfigImplBuilder<B extends MqttWebSocketConf
         return self();
     }
 
+    public @NotNull B httpHeaders(final @Nullable Map<String, String> httpHeaders) {
+        Checks.notNull(httpHeaders, "Http headers");
+        this.httpHeaders = httpHeaders;
+        return self();
+    }
+
     public @NotNull MqttWebSocketConfigImpl build() {
-        return new MqttWebSocketConfigImpl(serverPath, queryString, subprotocol, handshakeTimeoutMs);
+        return new MqttWebSocketConfigImpl(serverPath, queryString, subprotocol, handshakeTimeoutMs, httpHeaders);
     }
 
     public static class Default extends MqttWebSocketConfigImplBuilder<Default> implements MqttWebSocketConfigBuilder {

src/main/java/com/hivemq/client/internal/mqtt/handler/websocket/MqttWebSocketHttpHeaders.java

@@ -0,0 +1,44 @@
+/*
+ * Copyright 2018-present HiveMQ and the HiveMQ Community
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.hivemq.client.internal.mqtt.handler.websocket;
+
+import io.netty.channel.ChannelDuplexHandler;
+import io.netty.channel.ChannelHandlerContext;
+import io.netty.channel.ChannelPromise;
+import io.netty.handler.codec.http.HttpRequest;
+import org.jetbrains.annotations.NotNull;
+
+import java.util.Map;
+
+public class MqttWebSocketHttpHeaders extends ChannelDuplexHandler {
+
+    public static final String HTTP_HEADERS = "http.headers";
+    private final @NotNull Map<String, String> httpHeaders;
+
+    public MqttWebSocketHttpHeaders(@NotNull final Map<String, String> httpHeaders) {
+        this.httpHeaders = httpHeaders;
+    }
+
+    @Override
+    public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) throws Exception {
+        if (msg instanceof HttpRequest) {
+            final HttpRequest request = (HttpRequest) msg;
+            this.httpHeaders.forEach((key, value) -> request.headers().set(key, value));
+        }
+        super.write(ctx, msg, promise);
+    }
+}

src/main/java/com/hivemq/client/internal/mqtt/handler/websocket/MqttWebSocketInitializer.java

@@ -79,6 +79,7 @@ public void initChannel(
         channel.pipeline()
                 .addLast(HTTP_CODEC_NAME, new HttpClientCodec())
                 .addLast(HTTP_AGGREGATOR_NAME, new HttpObjectAggregator(65_535))
+                .addLast(MqttWebSocketHttpHeaders.HTTP_HEADERS, new MqttWebSocketHttpHeaders(webSocketConfig.getHttpHeaders()))
                 .addLast(MqttWebsocketHandshakeHandler.NAME,
                         new MqttWebsocketHandshakeHandler(handshaker, webSocketConfig.getHandshakeTimeoutMs(),
                                 onSuccess, onError))

src/main/java/com/hivemq/client/mqtt/MqttWebSocketConfig.java

@@ -20,6 +20,9 @@
 import com.hivemq.client.internal.mqtt.MqttWebSocketConfigImplBuilder;
 import org.jetbrains.annotations.NotNull;
 
+import java.util.LinkedHashMap;
+import java.util.Map;
+
 /**
  * Configuration for a WebSocket transport to use by {@link MqttClient MQTT clients}.
  *
@@ -51,6 +54,12 @@ public interface MqttWebSocketConfig {
      * @since 1.2
      */
     int DEFAULT_HANDSHAKE_TIMEOUT_MS = 10_000;
+    /**
+     * The default map of headers.
+     * @since 1.2.3
+     */
+    @NotNull Map<String, String> DEFAULT_HTTP_HEADERS = new LinkedHashMap<>();
+
 
     /**
      * Creates a builder for a WebSocket configuration.
@@ -82,11 +91,18 @@ public interface MqttWebSocketConfig {
      */
     int getHandshakeTimeoutMs();
 
+    /**
+     * @return map of already set headers.
+     * @since 1.2.3
+     */
+    @NotNull Map<String, String> getHttpHeaders();
+
     /**
      * Creates a builder for extending this WebSocket configuration.
      *
      * @return the created builder.
      * @since 1.1
      */
     @NotNull MqttWebSocketConfigBuilder extend();
+
 }

src/main/java/com/hivemq/client/mqtt/MqttWebSocketConfigBuilderBase.java

@@ -20,6 +20,7 @@
 import com.hivemq.client.annotations.DoNotImplement;
 import org.jetbrains.annotations.NotNull;
 
+import java.util.Map;
 import java.util.concurrent.TimeUnit;
 
 /**
@@ -71,4 +72,13 @@ public interface MqttWebSocketConfigBuilderBase<B extends MqttWebSocketConfigBui
      */
     @CheckReturnValue
     @NotNull B handshakeTimeout(long timeout, @NotNull TimeUnit timeUnit);
+
+    /**
+     * Sets the {@link MqttWebSocketConfig#getHttpHeaders() headers}.
+     *
+     * @param httpHeaders http headers.
+     * @return the builder.
+     */
+    @CheckReturnValue
+    @NotNull B httpHeaders(@NotNull Map<String, String> httpHeaders);
 }