io.gravitee.policy.trafficshadowing.invoker.ShadowProxyConnection Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of gravitee-policy-traffic-shadowing Show documentation
Asynchronously copy the traffic to another service
The newest version!
/**
 * Copyright (C) 2015 The Gravitee team (http://gravitee.io)
 *
 * 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 io.gravitee.policy.trafficshadowing.invoker;

import io.gravitee.gateway.api.buffer.Buffer;
import io.gravitee.gateway.api.handler.Handler;
import io.gravitee.gateway.api.http2.HttpFrame;
import io.gravitee.gateway.api.proxy.ProxyConnection;
import io.gravitee.gateway.api.proxy.ProxyResponse;
import io.gravitee.gateway.api.stream.WriteStream;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class ShadowProxyConnection implements ProxyConnection {

    private static final Logger LOGGER = LoggerFactory.getLogger(ShadowProxyConnection.class);

    private final ProxyConnection incomingProxyConnection;
    private final ProxyConnection shadowConnection;

    public ShadowProxyConnection(ProxyConnection incomingProxyConnection, ProxyConnection shadowConnection) {
        this.incomingProxyConnection = incomingProxyConnection;
        this.shadowConnection = shadowConnection;

        shadowConnection.responseHandler(response -> {
            LOGGER.debug("Traffic shadowing status is: {}", response.status());

            response.bodyHandler(noop -> {}).endHandler(noop -> {});

            // Resume the shadow response to read the stream and mark as ended
            response.resume();
        });

        shadowConnection.exceptionHandler(throwable -> {
            LOGGER.error("An error occurs while sending traffic shadowing request", throwable);
        });
    }

    @Override
    public ProxyConnection writeCustomFrame(HttpFrame frame) {
        incomingProxyConnection.writeCustomFrame(frame);
        shadowConnection.writeCustomFrame(frame);
        return this;
    }

    @Override
    public ProxyConnection cancel() {
        incomingProxyConnection.cancel();
        shadowConnection.cancel();
        return this;
    }

    @Override
    public ProxyConnection cancelHandler(Handler cancelHandler) {
        incomingProxyConnection.cancelHandler(cancelHandler);
        return this;
    }

    @Override
    public ProxyConnection exceptionHandler(Handler exceptionHandler) {
        incomingProxyConnection.exceptionHandler(exceptionHandler);
        return this;
    }

    @Override
    public ProxyConnection responseHandler(Handler responseHandler) {
        incomingProxyConnection.responseHandler(responseHandler);
        return this;
    }

    @Override
    public WriteStream write(Buffer buffer) {
        incomingProxyConnection.write(buffer);
        shadowConnection.write(buffer);
        return this;
    }

    @Override
    public void end() {
        incomingProxyConnection.end();
        shadowConnection.end();
    }

    @Override
    public void end(Buffer buffer) {
        incomingProxyConnection.end(buffer);
        shadowConnection.end(buffer);
    }

    @Override
    public WriteStream drainHandler(Handler drainHandler) {
        incomingProxyConnection.drainHandler(drainHandler);
        return this;
    }

    @Override
    public boolean writeQueueFull() {
        return incomingProxyConnection.writeQueueFull();
    }
}