Introduction

The WebSockets extension was developed within the Google Summer of Code 2012. This page should give developers an insight how it is structured.

The extension is built upon RFC6455, featuring version 13 of the WebSocket-protocol. The focus of this implementation lies on the payloads. As a result the user interface is transparent regarding WebSocket-frames.

Originally I wanted to build the extension upon Java's NIO features, that allows non-blocking reads. It worked fine for non-SSL connections, but I was not able to transform the javax.net.ssl.SSLSocket, that comes out of the modified commons-httpclient-3.1.jar library to a java.nio.channels.SocketChannel with some instance of javax.net.ssl.SSLEngine, as this would have been the way to go with NIO.

As a result each WebSocket channel consists of two threads:

• one listener on the outgoing connection from your browser to ZAP
• another listener on the incoming connection from the web server to ZAP

Database

Messages received are stored into the database. There are 3 tables so far:

• websocket_channel: stores information about each connection
• websocket_message: contains information about each message received & sent
• websocket_message_fuzz: if WebSocket messages are issued with the fuzz-extension, additional information is stored here

CREATE CACHED TABLE websocket_channel (
    channel_id BIGINT PRIMARY KEY,
    host VARCHAR(255) NOT NULL,
    port INTEGER NOT NULL,
    url VARCHAR(255) NOT NULL,
    start_timestamp TIMESTAMP NOT NULL,
    end_timestamp TIMESTAMP NULL,
    history_id INTEGER NULL,
    FOREIGN KEY (history_id) REFERENCES HISTORY(HISTORYID) ON DELETE SET NULL ON UPDATE SET NULL
);

CREATE CACHED TABLE websocket_message (
    message_id BIGINT NOT NULL,
    channel_id BIGINT NOT NULL,
    timestamp TIMESTAMP NOT NULL,
    opcode TINYINT NOT NULL,
    payload_utf8 CLOB NULL,
    payload_bytes BLOB NULL,
    payload_length BIGINT NOT NULL,
    is_outgoing BOOLEAN NOT NULL,
    PRIMARY KEY (message_id, channel_id),
    FOREIGN KEY (channel_id) REFERENCES websocket_channel(channel_id)
);

ALTER TABLE websocket_message ADD CONSTRAINT websocket_message_payload CHECK (
    payload_utf8 IS NOT NULL
      OR
    payload_bytes IS NOT NULL
);

CREATE CACHED TABLE websocket_message_fuzz (
    fuzz_id BIGINT NOT NULL,
    message_id BIGINT NOT NULL,
    channel_id BIGINT NOT NULL,
    state VARCHAR(50) NOT NULL,
    fuzz LONGVARCHAR NOT NULL,
    PRIMARY KEY (fuzz_id, message_id, channel_id),
    FOREIGN KEY (message_id, channel_id) REFERENCES websocket_message(message_id, channel_id) ON DELETE CASCADE
);

These tables are created in the class TableWebSocket, if not existed before.

Things to note:

• Primary key values are created within the application with instances of java.util.concurrent.atomic.AtomicInteger, see WebSocketProxy.channelIdGenerator, WebSocketProxy.messageIdGenerator & WebSocketFuzzableTextMessage.fuzzIdGenerator.
• websocket_channel.history_id may link to the HTTP message of the WebSocket handshake.
• websocket_channel.host and websocket_channel.url are not the same. The first field contains the result of Socket.getInetAddress().getHostName(), while the latter contains the requested URL of the WebSocket handshake.
• websocket_message contains two columns for payloads, namely payload_utf8 and payload_bytes. For binary-opcode messages the column payload_bytes is filled. For all other types of messages, the column payload_utf8 is set with the readable representation. This way, integration into the search-extension should be easier. The constraint websocket_message_payload ensures that at least one of these columns is set. An upgrade from HSQLDB version 1.8.0 to 2.2.9 was made to take advantage of the CLOB/BLOB fields:
  • Only a reference to the large object's content is returned, allowing you to retrieve only a substring, respectively only some bytes. This is used in the payload preview of the WebSockets-tab.

Class Diagram

Core

The first class diagram contains the core part without integration into brk- or fuzz-extension nor with UI classes.

Let us start with ExtensionWebSocket, which is the starting point of my contribution. It initializes all components and hooks them into ZAP. When a new WebSocket-connection is detected in the ProxyThread class, the following call takes place:

ExtensionWebSocket extWs = (ExtensionWebSocket) Control.getSingleton().getExtensionLoader().getExtension(ExtensionWebSocket.NAME);
extWs.addWebSocketsChannel(msg, inSocket, outSocket, outReader);

It takes the incoming & outgoing socket, the HttpMessage of the WebSocket-handshake and the current InputStream of the outgoing connection, which was used to read the HTTP response. This is of importance, as first WebSocket-messages are allowed to appear in the same TCP packet after the HTTP response. As it may buffer bytes, first messages would be lost if opening another InputStream on outSocket.

The ExtensionWebSocket creates a new instance of WebSocketProxy via the factory method WebSocketProxy.create(...), that returns a version specific WebSocketProxy instance. For now WebSocketProxyV13 is the only implementation of the abstract class WebSocketProxy. It contains an inner class WebSocketMessageV13 extending the abstract class WebSocketMessage.

Each WebSocketProxy instance creates two instances of WebSocketListener. These instances are threads listening to one of the given Socket's. If the first byte arrives, it calls WebSocketProxy.processRead(...) that handles the received WebSocket frame.

The WebSocketProxy class implements the Observer-pattern, allowing instances of WebSocketObserver to get notified about new frames or a change of the WebSocketProxy's state. The following observers are used so far (with order value in parenthesis):

• ExtensionFilter (0): Calls all enabled Filter instances, allowing them to change e.g.: the payload. There is a WebSocket-specific filter called FilterWebSocketPayload, that is added to the Filter-extension in the ExtensionWebSocket.hook(...) method.
• WebSocketProxyListenerBreak (95): Halts if a breakpoint applies and possibly changes payload.
• WebSocketStorage (100): Utilizes TableWebSocket to store channels and messages into the database.
• WebSocketPanel (105): Shows channels and their messages in the user interface under the WebSockets-tab. *WebSocketFuzzerHandler(110): Shows fuzzed messages in the user interface under the_fuzz-tab._

As you can see, this mechanism is a very powerful way to get informed about what is going on. In the class diagram you can see that each instance of WebSocketProxy has got its own observerList. If you want to observe all instances you have to add your WebSocketObserver implementation to the ExtensionWebSocket.allChannelObservers list. Do the following in your Extension* class:

@Override
public void hook(ExtensionHook extensionHook) {
	// 'this' implements WebSocketObserver
	extensionHook.addWebSocketObserver(this);
}

With the first WebSocket-connection arriving, the hooked observers are added to the ExtensionWebSocket.allChannelObservers list. Each time a new WebSocketProxy instance is created, every observer from this list is added to the WebSocketProxy.observerList.

WebSocket messages are processed in WebSocketProxy.processRead(...) as mentioned before. There are several types of messages, which is specified by the 4-bits opcode header:

• non-control frames
  • binary
  • text
• control frames
  • close
  • ping
  • pong

A non-control message may be split up across several frames. For this purpose a continuation frame is sent, resuming the last binary- or text-frame. In between arbitrary control frames are allowed to occur.

To achieve some loose coupling, I have introduced the *DTO classes, namely WebSocketChannelDTO & WebSocketMessageDTO. DTO stands for Data Transfer Object. They can be retrieved via:

• public WebSocketMessageDTO WebSocketMessage.getDTO();
• public WebSocketChannelDTO WebSocketProxy.getDTO();
• with various methods from the TableWebSocket These DTO-objects are in use across the WebSocket-extension.

User Interface

The main class is WebSocketPanel, which represents the WebSockets-tab. It contains all the UI elements visible there. The most important ones are:

• WebSocketPanel.channelSelectModel which is filled with all WebSocket channels. Via WebSocketPanel.getChannelComboBoxModel() you can retrieve an instance of ClonedComboBoxModel, whose items are backed by the original model, i.e. if the original ComboBox changes, also the cloned version changes. The ClonedComboBoxModel is used for various dialogues.
• handshakeButton: When a channel is selected in the ComboBox, this button is enabled. It allows the HttpMessage from the handshake to be shown in Request/Response tab.
• brkButton: See _brk_-extension integration for more information.
• filterButton: Opens up the filter.FilterWebSocketReplaceDialog allowing to change the type of messages shown in the WebSockets-tab.
• optionsButton: Opens up the options dialogue defined by OptionsWebSocketPanel. It is backed by the OptionsParamWebSocket, which is the interface to the saved settings.
• messagesView: This instance of WebSocketMessagesView wraps a JTable containing all WebSocketMessages. The model behind the JTable is given by an instance of WebSocketMessagesViewModel.

WebSocketMessagesViewModel extends the PagingTableModel, which holds only PagingTableModel.MAX_PAGE_SIZE entries in cache at any point in time, but the row count returns the total number of messages to be shown, resulting in a scrollbar that reflects a table containing all entries. When scrolling, or when new messages arrive, a new page is loaded from database. While in load, place-holder values are shown in the rows. In WebSocketMessagesViewModel.getRowCount() the number of rows is cached to save some queries. The WebSocketFuzzMessagesViewModel does also extend WebSocketMessagesViewModel as its entries are also stored in the database. See the _fuzz_-extension integration for more information.

There is another useful helper class, named WebSocketUiHelper. It has got methods that create UI elements for selecting channels, opcodes and direction. It is used by various dialogues and came into existence to bring up more consistency across dialogues:

• WebSocketBreakDialog: specify custom conditions for breakpoints
• FilterWebSocketReplaceDialog: allows to replace WebSocket payload using defined pattern
• WebSocketMessagesViewFilterDialog: restrict types of messages to be shown in the WebSockets tab

extension integration

filter

The FilterWebSocketPayload class allows for modification of WebSocket-payloads on specific messages. It is set up in the ExtensionWebSocket.hook(...) method. It overwrites the method onWebSocketPayload(...) and modifies a messages' payload if criteria are met. The ExtensionFilter implements WebSocketObserver and calls onWebSocketPayload(...) when a message arrives.

brk

There are several options for the break-behaviour of WebSocket messages. These options are enforced in the WebSocketBreakpointMessageHandler class. The decision if ZAP should hold on the arrival of a specific message, i.e. if a breakpoint applies, is reached in WebSocketBreakpointMessage.match(...). Beforehand WebSocketProxyListenerBreak.onMessageFrame(...) does some initial checks before passing on the power of decision.

fuzz

The Fuzzer-tab is able to show a messages view that inherits from the view in the WebSockets-tab. The correspondent classes are WebSocketFuzzMessagesView with its model class WebSocketFuzzMessagesViewModel. The view model is also backed by the database. The table websocket_message_fuzz is used to provide more information on the fuzzed messages. Unsuccessful fuzzed messages do not pass the WebSocketStorage class, that is responsible for saving messages into database. As a result there is an extra list for failed messages in WebSocketFuzzMessagesViewModel.erroneousMessages. A reason for unsuccessful fuzzing attempts may be closed WebSocket-channels.

WebSocketFuzzMessageDTO extends WebSocketMessageDTO and holds additional information on the fuzzing process. When an instance of WebSocketFuzzMessageDTO arrives at the WebSocketStorage class, additional information is saved to the websocket_message_fuzz-table. You can not only retrieve a DTO-object from a WebSocketMessage, but also create a WebSocketMessage from a WebSocketMessageDTO. The given DTO-object is saved as base-DTO in the WebSocketMessage. When you retrieve the DTO-object from a WebSocketMessage no new WebSocketMessageDTO instance is created, but the base-DTO is returned with current values. This mechanism is used to integrate the fuzzing of WebSocket messages.