From 18495a0954404bef9786000ef18ae344d5e5eb8e Mon Sep 17 00:00:00 2001
From: Oliver Goodman <oag@optusnet.com.au>
Date: Sat, 20 Aug 2016 10:16:41 +0000
Subject: [PATCH] Adds an API call to identify WebSocket requests.

---
 .../api/functions/TSHttpTxnIsWebsocket.en.rst | 45 +++++++++++++++++++
 lib/atscppapi/src/InterceptPlugin.cc          |  5 +++
 lib/atscppapi/src/TransactionPlugin.cc        |  5 +++
 .../src/include/atscppapi/TransactionPlugin.h |  2 +
 proxy/InkAPI.cc                               |  9 ++++
 proxy/api/ts/ts.h                             |  1 +
 6 files changed, 67 insertions(+)
 create mode 100644 doc/developer-guide/api/functions/TSHttpTxnIsWebsocket.en.rst

diff --git a/doc/developer-guide/api/functions/TSHttpTxnIsWebsocket.en.rst b/doc/developer-guide/api/functions/TSHttpTxnIsWebsocket.en.rst
new file mode 100644
index 00000000000..c60dc9dbf72
--- /dev/null
+++ b/doc/developer-guide/api/functions/TSHttpTxnIsWebsocket.en.rst
@@ -0,0 +1,45 @@
+.. Licensed to the Apache Software Foundation (ASF) under one
+   or more contributor license agreements.  See the NOTICE file
+   distributed with this work for additional information
+   regarding copyright ownership.  The ASF licenses this file
+   to you 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.
+
+.. default-domain:: c
+
+TSHttpTxnIsWebsocket
+********************
+
+Test whether a request is attempting to initiate Websocket connection.
+
+Synopsis
+========
+
+`#include <ts/ts.h>`
+
+.. function:: int TSHttpTxnIsWebsocket(TSHttpTxn txnp)
+
+Description
+===========
+
+:func:`TSHttpTxnIsWebsocket` tests whether the transaction
+is a WebSocket upgrade request.
+
+Return Values
+=============
+
+A non-zero value is returned if the relevant header value is found.
+
+See also
+========
+
+:manpage:`TSAPI(3ts)`
diff --git a/lib/atscppapi/src/InterceptPlugin.cc b/lib/atscppapi/src/InterceptPlugin.cc
index 2221d120c3a..274cb1fb2c1 100644
--- a/lib/atscppapi/src/InterceptPlugin.cc
+++ b/lib/atscppapi/src/InterceptPlugin.cc
@@ -254,6 +254,11 @@ InterceptPlugin::doRead()
   // Modify the input VIO to reflect how much data we've completed.
   TSVIONDoneSet(state_->input_.vio_, TSVIONDoneGet(state_->input_.vio_) + consumed);
 
+  if (isWebsocket()) {
+    TSVIOReenable(state_->input_.vio_);
+    return true;
+  }
+
   if ((state_->hdr_parsed_) && (state_->num_body_bytes_read_ >= state_->expected_body_size_)) {
     LOG_DEBUG("Completely read body");
     if (state_->num_body_bytes_read_ > state_->expected_body_size_) {
diff --git a/lib/atscppapi/src/TransactionPlugin.cc b/lib/atscppapi/src/TransactionPlugin.cc
index 778aca93d0b..7710d7df843 100644
--- a/lib/atscppapi/src/TransactionPlugin.cc
+++ b/lib/atscppapi/src/TransactionPlugin.cc
@@ -72,6 +72,11 @@ TransactionPlugin::getMutex()
   return state_->mutex_;
 }
 
+bool TransactionPlugin::isWebsocket() const
+{
+  return TSHttpTxnIsWebsocket(state_->ats_txn_handle_);
+}
+
 TransactionPlugin::~TransactionPlugin()
 {
   LOG_DEBUG("Destroying TransactionPlugin=%p", this);
diff --git a/lib/atscppapi/src/include/atscppapi/TransactionPlugin.h b/lib/atscppapi/src/include/atscppapi/TransactionPlugin.h
index c47a164039c..959bb503546 100644
--- a/lib/atscppapi/src/include/atscppapi/TransactionPlugin.h
+++ b/lib/atscppapi/src/include/atscppapi/TransactionPlugin.h
@@ -94,6 +94,8 @@ class TransactionPlugin : public Plugin
   void registerHook(Plugin::HookType hook_type);
   virtual ~TransactionPlugin();
 
+  bool isWebsocket() const;
+
 protected:
   TransactionPlugin(Transaction &transaction);
 
diff --git a/proxy/InkAPI.cc b/proxy/InkAPI.cc
index 949d18d4ced..38cb5d380d0 100644
--- a/proxy/InkAPI.cc
+++ b/proxy/InkAPI.cc
@@ -4982,6 +4982,15 @@ TSHttpTxnInfoIntGet(TSHttpTxn txnp, TSHttpTxnInfoKey key, TSMgmtInt *value)
   return TS_SUCCESS;
 }
 
+int
+TSHttpTxnIsWebsocket(TSHttpTxn txnp)
+{
+  sdk_assert(sdk_sanity_check_txn(txnp) == TS_SUCCESS);
+
+  HttpSM *sm = (HttpSM *)txnp;
+  return sm->t_state.is_websocket;
+}
+
 TSReturnCode
 TSHttpTxnCacheLookupUrlGet(TSHttpTxn txnp, TSMBuffer bufp, TSMLoc obj)
 {
diff --git a/proxy/api/ts/ts.h b/proxy/api/ts/ts.h
index 294d495739c..c1bc1d79605 100644
--- a/proxy/api/ts/ts.h
+++ b/proxy/api/ts/ts.h
@@ -2336,6 +2336,7 @@ tsapi TSReturnCode TSHttpTxnCacheLookupUrlGet(TSHttpTxn txnp, TSMBuffer bufp, TS
 tsapi TSReturnCode TSHttpTxnCacheLookupUrlSet(TSHttpTxn txnp, TSMBuffer bufp, TSMLoc obj);
 tsapi TSReturnCode TSHttpTxnPrivateSessionSet(TSHttpTxn txnp, int private_session);
 tsapi int TSHttpTxnBackgroundFillStarted(TSHttpTxn txnp);
+tsapi int TSHttpTxnIsWebsocket(TSHttpTxn txnp);
 
 /* Get the Txn's (HttpSM's) unique identifier, which is a sequence number since server start) */
 tsapi uint64_t TSHttpTxnIdGet(TSHttpTxn txnp);