1"""Convenience methods for use to manipulate traffic control settings.
2
3see http://linux.die.net/man/8/tc for details about traffic controls in linux.
4
5Example
6  import common
7  from autotest_lib.client.bin.net.net_tc import *
8  from autotest_lib.client.bin.net.net_utils import *
9
10  class mock_netif(object):
11
12    def __init__(self, name):
13        self._name = name
14
15    def get_name(self):
16        return self._name
17
18
19  netem_qdisc = netem()
20  netem_qdisc.add_param('loss 100%')
21
22  ack_filter = u32filter()
23  ack_filter.add_rule('match ip protocol 6 0xff')
24  ack_filter.add_rule('match u8 0x10 0x10 at nexthdr+13')
25  ack_filter.set_dest_qdisc(netem_qdisc)
26
27  root_qdisc = prio()
28  root_qdisc.get_class(2).set_leaf_qdisc(netem_qdisc)
29  root_qdisc.add_filter(ack_filter)
30
31  lo_if = mock_netif('lo')
32
33  root_qdisc.setup(lo_if)
34
35  # run test here ...
36  root_qdisc.restore(lo_if)
37
38"""
39
40import commands, os, re
41import common
42from autotest_lib.client.common_lib import error
43from autotest_lib.client.bin.net import net_utils
44
45# TODO (chavey) clean up those global here and new_handle()
46handle_counter = 0
47INCR = 100
48
49
50def new_handle():
51    global handle_counter
52    handle_counter += INCR
53    return handle_counter
54
55
56class tcclass(object):
57
58    def __init__(self, handle, minor, leaf_qdisc=None):
59        self._parent_class = None
60        self._children = []
61        self._leaf_qdisc = leaf_qdisc
62        self._handle = handle
63        self._minor = minor
64
65
66    def get_leaf_qdisc(self):
67        return self._leaf_qdisc
68
69
70    def set_leaf_qdisc(self, leaf_qdisc):
71        leaf_qdisc.set_parent_class(self)
72        self._leaf_qdisc = leaf_qdisc
73
74
75    def get_parent_class(self):
76        return self._parent_class
77
78
79    def set_parent_class(self, parent_class):
80        self._parent_class = parent_class
81
82
83    def get_minor(self):
84        return self._minor
85
86
87    def id(self):
88        return '%s:%s' % (self._handle, self._minor)
89
90
91    def add_child(self, child_class):
92        child_class.set_parent_class(self)
93        if child_class not in self._children:
94            self._child.append(child_class)
95
96
97    def setup(self, netif):
98        # setup leaf qdisc
99        if self._leaf_qdisc:
100            self._leaf_qdisc.setup(netif)
101
102        # setup child classes
103        for child in self._children:
104            child.setup()
105
106
107    def restore(self, netif):
108        # restore child classes
109        children_copy = list(self._children)
110        children_copy.reverse()
111        for child in children_copy:
112            child.restore()
113
114        # restore leaf qdisc
115        if self._leaf_qdisc:
116            self._leaf_qdisc.restore(netif)
117
118
119class tcfilter(object):
120
121    _tc_cmd = 'tc filter %(cmd)s dev %(dev)s parent %(parent)s protocol ' \
122               '%(protocol)s prio %(priority)s %(filtertype)s \\\n ' \
123               '%(rules)s \\\n  flowid %(flowid)s'
124
125    conf_device = 'dev'
126    conf_parent = 'parent'
127    conf_type = 'filtertype'
128    conf_protocol = 'protocol'
129    conf_priority = 'priority'
130    conf_flowid = 'flowid'
131    conf_command = 'cmd'
132    conf_rules = 'cmd'
133    conf_qdiscid = 'qdiscid'
134    conf_name = 'name'
135    conf_params = 'params'
136
137
138    def __init__(self):
139        self._parent_qdisc = None
140        self._dest_qdisc = None
141        self._protocol = 'ip'
142        self._priority = 1
143        self._handle = None
144        self._tc_conf = None
145
146
147    def get_parent_qdisc(self):
148        return self._parent_qdisc
149
150
151    def set_parent_qdisc(self, parent_qdisc):
152        self._parent_qdisc = parent_qdisc
153
154
155    def get_dest_qdisc(self):
156        return self._dest_qdisc
157
158
159    def set_dest_qdisc(self, dest_qdisc):
160        self._dest_qdisc = dest_qdisc
161
162
163    def get_protocol(self):
164        return self._protocol
165
166
167    def set_protocol(self, protocol):
168        self._protocol = protocol
169
170
171    def get_priority(self):
172        return self._priority
173
174
175    def set_priority(self, priority):
176        self._priority = priority
177
178
179    def get_handle(self):
180        return self._handle
181
182
183    def set_handle(self, handle):
184        self._handle = handle
185
186
187    def _get_tc_conf(self, netif):
188        if self._tc_conf:
189            return self._tc_conf
190        self._tc_conf = dict()
191        self._tc_conf[tcfilter.conf_device] = netif.get_name()
192        self._tc_conf[tcfilter.conf_parent] = self._parent_qdisc.id()
193        self._tc_conf[tcfilter.conf_type] = self.filtertype
194        self._tc_conf[tcfilter.conf_protocol] = self._protocol
195        self._tc_conf[tcfilter.conf_priotity] = self._priority
196        self._tc_conf[tcfilter.conf_flowid] = (
197            self._dest_qdisc.get_parent_class().id())
198        return self._tc_conf
199
200
201    def tc_cmd(self, tc_conf):
202        print self._tc_cmd % tc_conf
203
204
205    def setup(self, netif):
206        pass
207
208
209    def restore(self, netif):
210        pass
211
212
213class u32filter(tcfilter):
214
215    filtertype = 'u32'
216
217    def __init__(self):
218        super(u32filter, self).__init__()
219        self._rules = []
220
221
222    def _filter_rules(self):
223        return ' \\\n  '.join(self._rules)
224
225
226    def add_rule(self, rule):
227        self._rules.append(rule)
228
229
230    def setup(self, netif):
231        tc_conf = self._get_tc_conf(netif)
232        tc_conf[tcfilter.conf_cmd] = 'add'
233        tc_conf[tcfilter.conf_rules] = self._filter_rules()
234        self.tc_cmd(tc_conf)
235
236
237    def restore(self, netif):
238        tc_conf = self._get_tc_conf(netif)
239        tc_conf[tcfilter.conf_cmd] = 'del'
240        tc_conf[tcfilter.conf_rules] = self._filter_rules()
241        self.tc_cmd(tc_conf)
242
243#TODO (ncrao): generate some typical rules: ack, syn, synack,
244#              dport/sport, daddr/sddr, etc.
245class qdisc(object):
246
247    # tc command
248    _tc_cmd = 'tc qdisc %(cmd)s dev %(dev)s %(parent)s ' \
249              'handle %(qdiscid)s %(name)s %(params)s'
250
251    def __init__(self, handle):
252        self._handle = handle
253        self._parent_class = None
254        self._tc_conf = None
255
256
257    def get_handle(self):
258        return self._handle
259
260
261    def get_parent_class(self):
262        return self._parent_class
263
264
265    def set_parent_class(self, parent_class):
266        self._parent_class = parent_class
267
268
269    def _get_tc_conf(self, netif):
270        if self._tc_conf:
271            return self._tc_conf
272        self._tc_conf = dict()
273        self._tc_conf[tcfilter.conf_device] = netif.get_name()
274        if self._parent_class:
275            self._tc_conf[tcfilter.conf_parent] = ('parent %s' %
276                                                   self._parent_class.id())
277        else:
278            self._tc_conf[tcfilter.conf_parent] = 'root'
279        self._tc_conf[tcfilter.conf_qdiscid] = self.id()
280        self._tc_conf[tcfilter.conf_name] = self.name
281        self._tc_conf[tcfilter.conf_params] = ''
282        return self._tc_conf
283
284
285    def id(self):
286        return '%s:0' % self._handle
287
288
289    def tc_cmd(self, tc_conf):
290        print self._tc_cmd % tc_conf
291
292
293    def setup(self, netif):
294        tc_conf = self._get_tc_conf(netif)
295        tc_conf[tcfilter.conf_command] = 'add'
296        self.tc_cmd(tc_conf)
297
298
299    def restore(self, netif):
300        tc_conf = self._get_tc_conf(netif)
301        tc_conf[tcfilter.conf_command] = 'del'
302        self.tc_cmd(tc_conf)
303
304
305class classful_qdisc(qdisc):
306
307    classful = True
308
309    def __init__(self, handle):
310        super(classful_qdisc, self).__init__(handle)
311        self._classes = []
312        self._filters = []
313
314
315    def add_class(self, child_class):
316        self._classes.append(child_class)
317
318
319    def add_filter(self, filter):
320        filter.set_parent_qdisc(self)
321        self._filters.append(filter)
322
323
324    def setup(self, netif):
325        super(classful_qdisc, self).setup(netif)
326
327        # setup child classes
328        for child in self._classes:
329            child.setup(netif)
330
331        # setup filters
332        for filter in self._filters:
333            filter.setup(netif)
334
335
336    def restore(self, netif):
337        # restore filters
338        filters_copy = list(self._filters)
339        filters_copy.reverse()
340        for filter in filters_copy:
341            filter.restore(netif)
342
343        # restore child classes
344        classes_copy = list(self._classes)
345        classes_copy.reverse()
346        for child in classes_copy:
347            child.restore(netif)
348
349        super(classful_qdisc, self).restore(netif)
350
351
352class prio(classful_qdisc):
353
354    name = 'prio'
355
356    def __init__(self, handle=new_handle(), bands=3):
357        super(prio, self).__init__(handle)
358        self._bands = bands
359        for counter in range(bands):
360            self.add_class(tcclass(handle, counter + 1))
361
362
363    def setup(self, netif):
364        super(prio, self).setup(netif)
365
366
367    def get_class(self, band):
368        if band > self._bands:
369            raise error.TestError('error inserting %s at band %s' % \
370                                  (qdisc.name, band))
371        return self._classes[band]
372
373
374class classless_qdisc(qdisc):
375
376    classful = False
377
378    def __init__(self, handle):
379        super(classless_qdisc, self).__init__(handle)
380
381
382class pfifo(classless_qdisc):
383
384    name = 'pfifo'
385
386    def __init__(self, handle=new_handle()):
387        super(pfifo, self).__init__(handle)
388
389
390    def setup(self, netif):
391        super(pfifo, self).setup(netif)
392
393
394class netem(classless_qdisc):
395
396    name = 'netem'
397
398    def __init__(self, handle=new_handle()):
399        super(netem, self).__init__(handle)
400        self._params = list()
401
402
403    def add_param(self, param):
404        self._params.append(param)
405
406
407    def setup(self, netif):
408        super(netem, self).setup(netif)
409        tc_conf = self._get_tc_conf(netif)
410        tc_conf[tcfilter.conf_command] = 'change'
411        tc_conf[tcfilter.conf_params] = ' '.join(self._params)
412        self.tc_cmd(tc_conf)
413