import ballerina/http;
import ballerina/log;final string ASSOCIATED_CONNECTION = "ASSOCIATED_CONNECTION";
final string REMOTE_BACKEND = "ws://echo.websocket.org";@http:WebSocketServiceConfig {
    path: "/proxy/ws"
}
service SimpleProxyService on new http:Listener(9090) {
    resource function onOpen(http:WebSocketCaller caller) {        http:WebSocketClient wsClientEp = new(
            REMOTE_BACKEND,
            {callbackService: ClientService,
            readyOnConnect: false
        });
        wsClientEp.setAttribute(ASSOCIATED_CONNECTION, caller);
        caller.setAttribute(ASSOCIATED_CONNECTION, wsClientEp);
        var err = wsClientEp->ready();
        if (err is http:WebSocketError) {
            log:printError("Error calling ready on client", err);
        }
    }
    resource function onText(http:WebSocketCaller caller, string text,
                                boolean finalFrame) {        http:WebSocketClient clientEp =
                    getAssociatedClientEndpoint(caller);
        var err = clientEp->pushText(text, finalFrame);
        if (err is http:WebSocketError) {
            log:printError("Error occurred when sending text message", err);
        }
    }
    resource function onBinary(http:WebSocketCaller caller, byte[] data,
                                boolean finalFrame) {        http:WebSocketClient clientEp =
                        getAssociatedClientEndpoint(caller);
        var err = clientEp->pushBinary(data, finalFrame);
        if (err is http:WebSocketError) {
            log:printError("Error occurred when sending binary message", err);
        }
    }
    resource function onError(http:WebSocketCaller caller, error err) {        http:WebSocketClient clientEp =
                        getAssociatedClientEndpoint(caller);
        var e = clientEp->close(statusCode = 1011,
                        reason = "Unexpected condition");
        if (e is http:WebSocketError) {
            log:printError("Error occurred when closing the connection", e);
        }
        _ = caller.removeAttribute(ASSOCIATED_CONNECTION);
        log:printError("Unexpected error hence closing the connection", err);
    }
    resource function onClose(http:WebSocketCaller caller, int statusCode,
                                string reason) {        http:WebSocketClient clientEp =
                        getAssociatedClientEndpoint(caller);
        var err = clientEp->close(statusCode = statusCode, reason = reason);
        if (err is http:WebSocketError) {
            log:printError("Error occurred when closing the connection", err);
        }
        _ = caller.removeAttribute(ASSOCIATED_CONNECTION);
    }
}
service ClientService = @http:WebSocketServiceConfig {} service {
    resource function onText(http:WebSocketClient caller, string text,
                                boolean finalFrame) {        http:WebSocketCaller serverEp =
                        getAssociatedServerEndpoint(caller);
        var err = serverEp->pushText(text, finalFrame);
        if (err is http:WebSocketError) {
            log:printError("Error occurred when sending text message", err);
        }
    }
    resource function onBinary(http:WebSocketClient caller, byte[] data,
                                boolean finalFrame) {        http:WebSocketCaller serverEp =
                        getAssociatedServerEndpoint(caller);
        var err = serverEp->pushBinary(data, finalFrame);
        if (err is http:WebSocketError) {
           log:printError("Error occurred when sending binary message", err);
        }
    }
    resource function onError(http:WebSocketClient caller, error err) {        http:WebSocketCaller serverEp =
                        getAssociatedServerEndpoint(caller);
        var e = serverEp->close(statusCode = 1011,
                        reason = "Unexpected condition");
        if (e is http:WebSocketError) {
            log:printError("Error occurred when closing the connection",
                            err = e);
        }
        _ = caller.removeAttribute(ASSOCIATED_CONNECTION);
        log:printError("Unexpected error hense closing the connection", err);
    }
    resource function onClose(http:WebSocketClient caller, int statusCode,
                                string reason) {        http:WebSocketCaller serverEp =
                        getAssociatedServerEndpoint(caller);
        var err = serverEp->close(statusCode = statusCode, reason = reason);
            if (err is http:WebSocketError) {
                log:printError("Error occurred when closing the connection", err);
            }
        _ = caller.removeAttribute(ASSOCIATED_CONNECTION);
    }
};
function getAssociatedClientEndpoint(http:WebSocketCaller ep)
                                        returns (http:WebSocketClient) {
    http:WebSocketClient wsClient =
            <http:WebSocketClient>ep.getAttribute(ASSOCIATED_CONNECTION);
    return wsClient;
}
function getAssociatedServerEndpoint(http:WebSocketClient ep)
                                        returns (http:WebSocketCaller) {
    http:WebSocketCaller wsEndpoint =
            <http:WebSocketCaller>ep.getAttribute(ASSOCIATED_CONNECTION);
    return wsEndpoint;
}

Proxy Server

Ballerina can act as a proxy to another remote service. Ballerina forwards the requests sent by the clients that connect to it via a WebSocket to the respective remote service.

import ballerina/http;
import ballerina/log;
final string ASSOCIATED_CONNECTION = "ASSOCIATED_CONNECTION";
final string REMOTE_BACKEND = "ws://echo.websocket.org";

The Url of the remote backend.

@http:WebSocketServiceConfig {
    path: "/proxy/ws"
}
service SimpleProxyService on new http:Listener(9090) {
    resource function onOpen(http:WebSocketCaller caller) {

This resource gets invoked when a new client connects. Since messages to the server are not read by the service until the execution of the onOpen resource finishes, operations which should happen before reading messages should be done in the onOpen resource.

        http:WebSocketClient wsClientEp = new(
            REMOTE_BACKEND,
            {callbackService: ClientService,
            readyOnConnect: false
        });

When creating client endpoint, if readyOnConnect flag is set to false client endpoint does not start reading frames automatically.

        wsClientEp.setAttribute(ASSOCIATED_CONNECTION, caller);
        caller.setAttribute(ASSOCIATED_CONNECTION, wsClientEp);

Associate connections before starting to read messages.

        var err = wsClientEp->ready();
        if (err is http:WebSocketError) {
            log:printError("Error calling ready on client", err);
        }
    }

Once the client is ready to receive frames the remote function ready of the client need to be called separately.

    resource function onText(http:WebSocketCaller caller, string text,
                                boolean finalFrame) {

This resource gets invoked upon receiving a new text frame from a client.

        http:WebSocketClient clientEp =
                    getAssociatedClientEndpoint(caller);
        var err = clientEp->pushText(text, finalFrame);
        if (err is http:WebSocketError) {
            log:printError("Error occurred when sending text message", err);
        }
    }
    resource function onBinary(http:WebSocketCaller caller, byte[] data,
                                boolean finalFrame) {

This resource gets invoked upon receiving a new binary frame from a client.

        http:WebSocketClient clientEp =
                        getAssociatedClientEndpoint(caller);
        var err = clientEp->pushBinary(data, finalFrame);
        if (err is http:WebSocketError) {
            log:printError("Error occurred when sending binary message", err);
        }
    }
    resource function onError(http:WebSocketCaller caller, error err) {

This resource gets invoked when an error occurs in the connection.

        http:WebSocketClient clientEp =
                        getAssociatedClientEndpoint(caller);
        var e = clientEp->close(statusCode = 1011,
                        reason = "Unexpected condition");
        if (e is http:WebSocketError) {
            log:printError("Error occurred when closing the connection", e);
        }
        _ = caller.removeAttribute(ASSOCIATED_CONNECTION);
        log:printError("Unexpected error hence closing the connection", err);
    }
    resource function onClose(http:WebSocketCaller caller, int statusCode,
                                string reason) {

This resource gets invoked when a client connection is closed from the client side.

        http:WebSocketClient clientEp =
                        getAssociatedClientEndpoint(caller);
        var err = clientEp->close(statusCode = statusCode, reason = reason);
        if (err is http:WebSocketError) {
            log:printError("Error occurred when closing the connection", err);
        }
        _ = caller.removeAttribute(ASSOCIATED_CONNECTION);
    }
}
service ClientService = @http:WebSocketServiceConfig {} service {

Client service to receive frames from the remote server.

    resource function onText(http:WebSocketClient caller, string text,
                                boolean finalFrame) {

This resource gets invoked upon receiving a new text frame from the remote backend.

        http:WebSocketCaller serverEp =
                        getAssociatedServerEndpoint(caller);
        var err = serverEp->pushText(text, finalFrame);
        if (err is http:WebSocketError) {
            log:printError("Error occurred when sending text message", err);
        }
    }
    resource function onBinary(http:WebSocketClient caller, byte[] data,
                                boolean finalFrame) {

This resource gets invoked upon receiving a new binary frame from the remote backend.

        http:WebSocketCaller serverEp =
                        getAssociatedServerEndpoint(caller);
        var err = serverEp->pushBinary(data, finalFrame);
        if (err is http:WebSocketError) {
           log:printError("Error occurred when sending binary message", err);
        }
    }
    resource function onError(http:WebSocketClient caller, error err) {

This resource gets invoked when an error occurs in the connection.

        http:WebSocketCaller serverEp =
                        getAssociatedServerEndpoint(caller);
        var e = serverEp->close(statusCode = 1011,
                        reason = "Unexpected condition");
        if (e is http:WebSocketError) {
            log:printError("Error occurred when closing the connection",
                            err = e);
        }
        _ = caller.removeAttribute(ASSOCIATED_CONNECTION);
        log:printError("Unexpected error hense closing the connection", err);
    }
    resource function onClose(http:WebSocketClient caller, int statusCode,
                                string reason) {

This resource gets invoked when a client connection is closed by the remote backend.

        http:WebSocketCaller serverEp =
                        getAssociatedServerEndpoint(caller);
        var err = serverEp->close(statusCode = statusCode, reason = reason);
            if (err is http:WebSocketError) {
                log:printError("Error occurred when closing the connection", err);
            }
        _ = caller.removeAttribute(ASSOCIATED_CONNECTION);
    }
};
function getAssociatedClientEndpoint(http:WebSocketCaller ep)
                                        returns (http:WebSocketClient) {
    http:WebSocketClient wsClient =
            <http:WebSocketClient>ep.getAttribute(ASSOCIATED_CONNECTION);
    return wsClient;
}

Function to retrieve associated client for a particular caller.

function getAssociatedServerEndpoint(http:WebSocketClient ep)
                                        returns (http:WebSocketCaller) {
    http:WebSocketCaller wsEndpoint =
            <http:WebSocketCaller>ep.getAttribute(ASSOCIATED_CONNECTION);
    return wsEndpoint;
}

Function to retrieve the associated caller for a client.

# To start the service, navigate to the directory that contains the
# `.bal` file and use the `ballerina build` command.
$ ballerina build websocket_proxy_server.bal
# Run the sample using the `run` command on the jar file generated:
$ ballerina run websocket_proxy_server-executable.jar
# Now, this service can be invoked by any WebSocket client using the url "ws://localhost:9090/proxy/ws"
# To check the sample, you can use Chrome or Firefox JavaScript console and run the following commands. <br>
$ var ws = new WebSocket("ws://localhost:9090/proxy/ws");
$ ws.onmessage = function(frame) {console.log(frame.data)};
$ ws.onclose = function(frame) {console.log(frame)};
# Send messages.
$ ws.send("hello world");
#Close the connection.
$ ws.close(1000, "I want to go");