Skip to content

Commit a59dea0

Browse files
author
Zizhong Zhang
committed
TS-4500: add cookie-rewrite functionality into header-rewrite plugin
1 parent 7ec0d8d commit a59dea0

File tree

8 files changed

+451
-0
lines changed

8 files changed

+451
-0
lines changed

‎ci/tsqa/files/header-rewrite.config‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
cond %{READ_REQUEST_PRE_REMAP_HOOK}
2+
cond %{PATH} /^.*addcookie$/ [AND]
3+
add-cookie testkey testaddvalue
4+
5+
cond %{READ_REQUEST_PRE_REMAP_HOOK}
6+
cond %{PATH} /^.*rmcookie$/ [AND]
7+
rm-cookie testkey
8+
9+
10+
cond %{READ_REQUEST_PRE_REMAP_HOOK}
11+
cond %{PATH} /^.*setcookie$/ [AND]
12+
set-cookie testkey testsetvalue
13+

‎ci/tsqa/tests/test_header_rewrite.py‎

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
'''
2+
Test cookie rewrite
3+
'''
4+
import os
5+
import requests
6+
import time
7+
import logging
8+
import random
9+
import tsqa.test_cases
10+
import helpers
11+
import shutil
12+
import SocketServer
13+
import urllib2
14+
15+
log = logging.getLogger(__name__)
16+
17+
class EchoServerHandler(SocketServer.BaseRequestHandler):
18+
"""
19+
A subclass of RequestHandler which will return all data received back
20+
"""
21+
22+
def handle(self):
23+
# Receive the data in small chunks and retransmit it
24+
while True:
25+
data = self.request.recv(4096).strip()
26+
if data:
27+
log.debug('Sending data back to the client')
28+
else:
29+
log.debug('Client disconnected')
30+
break
31+
cookie = ''
32+
if 'Cookie' in data:
33+
cookie = data.split('Cookie: ')[1].split('\r\n')[0]
34+
35+
resp = ('HTTP/1.1 200 OK\r\n'
36+
'Content-Length: {data_length}\r\n'
37+
'Content-Type: text/html; charset=UTF-8\r\n'
38+
'Connection: keep-alive\r\n'
39+
'\r\n{data_string}'.format(
40+
data_length = len(cookie),
41+
data_string = cookie
42+
))
43+
self.request.sendall(resp)
44+
45+
class TestHeaderRewrite(helpers.EnvironmentCase):
46+
'''
47+
Tests for header rewrite
48+
'''
49+
@classmethod
50+
def setUpEnv(cls, env):
51+
cls.traffic_server_port = int(cls.configs['records.config']['CONFIG']['proxy.config.http.server_ports'])
52+
53+
# create a socket server
54+
cls.socket_server = tsqa.endpoint.SocketServerDaemon(EchoServerHandler)
55+
cls.socket_server.start()
56+
cls.socket_server.ready.wait()
57+
58+
cls.configs['remap.config'].add_line(
59+
'map / http://127.0.0.1:%d' %(cls.socket_server.port)
60+
)
61+
62+
# setup the plugin
63+
cls.config_file = 'header-rewrite.config'
64+
cls.test_config_path = helpers.tests_file_path(cls.config_file)
65+
66+
cls.configs['plugin.config'].add_line('%s/header_rewrite.so %s' % (
67+
cls.environment.layout.plugindir,
68+
cls.test_config_path
69+
))
70+
71+
def test_cookie_rewrite(self):
72+
73+
cookie_test_add_dict = {
74+
'' : 'testkey=testaddvalue',
75+
'testkey=somevalue' : 'testkey=somevalue',
76+
'otherkey=testvalue' : 'otherkey=testvalue;testkey=testaddvalue',
77+
'testkey = "other=value"; a = a' : 'testkey = "other=value"; a = a',
78+
'testkeyx===' : 'testkeyx===;testkey=testaddvalue'
79+
}
80+
for key in cookie_test_add_dict:
81+
opener = urllib2.build_opener()
82+
opener.addheaders.append(('Cookie', key))
83+
f = opener.open("http://127.0.0.1:%d/addcookie" % (self.traffic_server_port))
84+
resp = f.read()
85+
self.assertEqual(resp, cookie_test_add_dict[key])
86+
87+
cookie_test_rm_dict = {
88+
'' : '',
89+
' testkey=somevalue' : '',
90+
'otherkey=testvalue' : 'otherkey=testvalue',
91+
'testkey = "other=value" ; a = a' : ' a = a',
92+
'otherkey=othervalue= ; testkey===' : 'otherkey=othervalue= ',
93+
'firstkey ="firstvalue" ; testkey = =; secondkey=\'\'' : 'firstkey ="firstvalue" ; secondkey=\'\''
94+
}
95+
for key in cookie_test_rm_dict:
96+
opener = urllib2.build_opener()
97+
opener.addheaders.append(('Cookie', key))
98+
f = opener.open("http://127.0.0.1:%d/rmcookie" % (self.traffic_server_port))
99+
resp = f.read()
100+
self.assertEqual(resp, cookie_test_rm_dict[key])
101+
102+
cookie_test_set_dict = {
103+
'' : 'testkey=testsetvalue',
104+
'testkey=somevalue' : 'testkey=testsetvalue',
105+
'otherkey=testvalue' : 'otherkey=testvalue;testkey=testsetvalue',
106+
'testkey = "other=value"; a = a' : 'testkey = testsetvalue; a = a',
107+
'testkeyx===' : 'testkeyx===;testkey=testsetvalue',
108+
'firstkey ="firstvalue" ; testkey = =; secondkey=\'\'' : 'firstkey ="firstvalue" ; testkey = testsetvalue; secondkey=\'\''
109+
}
110+
for key in cookie_test_set_dict:
111+
opener = urllib2.build_opener()
112+
opener.addheaders.append(('Cookie', key))
113+
f = opener.open("http://127.0.0.1:%d/setcookie" % (self.traffic_server_port))
114+
resp = f.read()
115+
self.assertEqual(resp, cookie_test_set_dict[key])

‎doc/admin-guide/plugins/header_rewrite.en.rst‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -458,6 +458,15 @@ occurs first).
458458

459459
The following operators are available:
460460

461+
add-cookie
462+
~~~~~~~~~~
463+
::
464+
465+
add-cookie <name> <value>
466+
467+
Adds a new ``<name>`` cookie line with the contents ``<value>``. Note that this
468+
operator will do nothing if a cookie pair with ``<name>`` already exists.
469+
461470
add-header
462471
~~~~~~~~~~
463472
::
@@ -512,6 +521,14 @@ rm-header
512521

513522
Removes the header ``<name>``.
514523

524+
rm-cookie
525+
~~~~~~~~~
526+
::
527+
528+
rm-cookie <name>
529+
530+
Removes the cookie ``<name>``.
531+
515532
set-config
516533
~~~~~~~~~~
517534
::
@@ -626,6 +643,15 @@ When invoked, and when ``<value>`` is any of ``1``, ``true``, or ``TRUE``, this
626643
operator causes |TS| to abort further request remapping. Any other value and
627644
the operator will effectively be a no-op.
628645

646+
set-cookie
647+
~~~~~~~~~~
648+
::
649+
650+
set-cookie <name> <value>
651+
652+
Replaces the value of cookie ``<name>`` with ``<value>``, creating the cookie
653+
if necessary.
654+
629655
Operator Flags
630656
--------------
631657

‎plugins/header_rewrite/factory.cc‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,12 @@ operator_factory(const std::string &op)
5656
o = new OperatorNoOp();
5757
} else if (op == "counter") {
5858
o = new OperatorCounter();
59+
} else if (op == "rm-cookie") {
60+
o = new OperatorRMCookie();
61+
} else if (op == "set-cookie") {
62+
o = new OperatorSetCookie();
63+
} else if (op == "add-cookie") {
64+
o = new OperatorAddCookie();
5965
} else if (op == "set-conn-dscp") {
6066
o = new OperatorSetConnDSCP();
6167
} else if (op == "set-debug") {

‎plugins/header_rewrite/operator.cc‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,3 +58,14 @@ OperatorHeaders::initialize(Parser &p)
5858
require_resources(RSRC_CLIENT_REQUEST_HEADERS);
5959
require_resources(RSRC_CLIENT_RESPONSE_HEADERS);
6060
}
61+
62+
void
63+
OperatorCookies::initialize(Parser &p)
64+
{
65+
Operator::initialize(p);
66+
67+
_cookie = p.get_arg();
68+
69+
require_resources(RSRC_SERVER_REQUEST_HEADERS);
70+
require_resources(RSRC_CLIENT_REQUEST_HEADERS);
71+
}

‎plugins/header_rewrite/operator.h‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,4 +83,21 @@ class OperatorHeaders : public Operator
8383
DISALLOW_COPY_AND_ASSIGN(OperatorHeaders);
8484
};
8585

86+
///////////////////////////////////////////////////////////////////////////////
87+
// Base class for all Cookie based Operators, this is obviously also an
88+
// Operator interface.
89+
//
90+
class OperatorCookies : public Operator
91+
{
92+
public:
93+
OperatorCookies() : _cookie("") { TSDebug(PLUGIN_NAME_DBG, "Calling CTOR for OperatorCookies"); }
94+
void initialize(Parser &p);
95+
96+
protected:
97+
std::string _cookie;
98+
99+
private:
100+
DISALLOW_COPY_AND_ASSIGN(OperatorCookies);
101+
};
102+
86103
#endif // __OPERATOR_H

0 commit comments

Comments
 (0)