1import os
2import signal
3from string import Template
4import subprocess
5import time
6from multiprocessing import Pool
7from functools import cached_property
8from TdcPlugin import TdcPlugin
9
10from tdc_config import *
11
12try:
13    from pyroute2 import netns
14    from pyroute2 import IPRoute
15    netlink = True
16except ImportError:
17    netlink = False
18    print("!!! Consider installing pyroute2 !!!")
19
20def prepare_suite(obj, test):
21    original = obj.args.NAMES
22
23    if 'skip' in test and test['skip'] == 'yes':
24        return
25
26    if 'nsPlugin' not in test['plugins']:
27        return
28
29    shadow = {}
30    shadow['IP'] = original['IP']
31    shadow['TC'] = original['TC']
32    shadow['NS'] = '{}-{}'.format(original['NS'], test['random'])
33    shadow['DEV0'] = '{}id{}'.format(original['DEV0'], test['id'])
34    shadow['DEV1'] = '{}id{}'.format(original['DEV1'], test['id'])
35    shadow['DUMMY'] = '{}id{}'.format(original['DUMMY'], test['id'])
36    shadow['DEV2'] = original['DEV2']
37    obj.args.NAMES = shadow
38
39    if netlink == True:
40        obj._nl_ns_create()
41    else:
42        obj._ns_create()
43
44    # Make sure the netns is visible in the fs
45    while True:
46        obj._proc_check()
47        try:
48            ns = obj.args.NAMES['NS']
49            f = open('/run/netns/{}'.format(ns))
50            f.close()
51            break
52        except:
53            time.sleep(0.1)
54            continue
55
56    obj.args.NAMES = original
57
58class SubPlugin(TdcPlugin):
59    def __init__(self):
60        self.sub_class = 'ns/SubPlugin'
61        super().__init__()
62
63    def pre_suite(self, testcount, testlist):
64        from itertools import cycle
65
66        super().pre_suite(testcount, testlist)
67
68        print("Setting up namespaces and devices...")
69
70        with Pool(self.args.mp) as p:
71            it = zip(cycle([self]), testlist)
72            p.starmap(prepare_suite, it)
73
74    def pre_case(self, caseinfo, test_skip):
75        if self.args.verbose:
76            print('{}.pre_case'.format(self.sub_class))
77
78        if test_skip:
79            return
80
81    def post_case(self):
82        if self.args.verbose:
83            print('{}.post_case'.format(self.sub_class))
84
85        self._ns_destroy()
86
87    def post_suite(self, index):
88        if self.args.verbose:
89            print('{}.post_suite'.format(self.sub_class))
90
91        # Make sure we don't leak resources
92        for f in os.listdir('/run/netns/'):
93            cmd = self._replace_keywords("$IP netns del {}".format(f))
94
95            if self.args.verbose > 3:
96                print('_exec_cmd:  command "{}"'.format(cmd))
97
98            subprocess.run(cmd, shell=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
99
100    def adjust_command(self, stage, command):
101        super().adjust_command(stage, command)
102        cmdform = 'list'
103        cmdlist = list()
104
105        if self.args.verbose:
106            print('{}.adjust_command'.format(self.sub_class))
107
108        if not isinstance(command, list):
109            cmdform = 'str'
110            cmdlist = command.split()
111        else:
112            cmdlist = command
113        if stage == 'setup' or stage == 'execute' or stage == 'verify' or stage == 'teardown':
114            if self.args.verbose:
115                print('adjust_command:  stage is {}; inserting netns stuff in command [{}] list [{}]'.format(stage, command, cmdlist))
116            cmdlist.insert(0, self.args.NAMES['NS'])
117            cmdlist.insert(0, 'exec')
118            cmdlist.insert(0, 'netns')
119            cmdlist.insert(0, self.args.NAMES['IP'])
120        else:
121            pass
122
123        if cmdform == 'str':
124            command = ' '.join(cmdlist)
125        else:
126            command = cmdlist
127
128        if self.args.verbose:
129            print('adjust_command:  return command [{}]'.format(command))
130        return command
131
132    def _nl_ns_create(self):
133        ns = self.args.NAMES["NS"];
134        dev0 = self.args.NAMES["DEV0"];
135        dev1 = self.args.NAMES["DEV1"];
136        dummy = self.args.NAMES["DUMMY"];
137
138        if self.args.verbose:
139            print('{}._nl_ns_create'.format(self.sub_class))
140
141        netns.create(ns)
142        netns.pushns(newns=ns)
143        with IPRoute() as ip:
144            ip.link('add', ifname=dev1, kind='veth', peer={'ifname': dev0, 'net_ns_fd':'/proc/1/ns/net'})
145            ip.link('add', ifname=dummy, kind='dummy')
146            while True:
147                try:
148                    dev1_idx = ip.link_lookup(ifname=dev1)[0]
149                    dummy_idx = ip.link_lookup(ifname=dummy)[0]
150                    ip.link('set', index=dev1_idx, state='up')
151                    ip.link('set', index=dummy_idx, state='up')
152                    break
153                except:
154                    time.sleep(0.1)
155                    continue
156        netns.popns()
157
158        with IPRoute() as ip:
159            while True:
160                try:
161                    dev0_idx = ip.link_lookup(ifname=dev0)[0]
162                    ip.link('set', index=dev0_idx, state='up')
163                    break
164                except:
165                    time.sleep(0.1)
166                    continue
167
168    def _ns_create_cmds(self):
169        cmds = []
170
171        ns = self.args.NAMES['NS']
172
173        cmds.append(self._replace_keywords('netns add {}'.format(ns)))
174        cmds.append(self._replace_keywords('link add $DEV1 type veth peer name $DEV0'))
175        cmds.append(self._replace_keywords('link set $DEV1 netns {}'.format(ns)))
176        cmds.append(self._replace_keywords('link add $DUMMY type dummy'.format(ns)))
177        cmds.append(self._replace_keywords('link set $DUMMY netns {}'.format(ns)))
178        cmds.append(self._replace_keywords('netns exec {} $IP link set $DEV1 up'.format(ns)))
179        cmds.append(self._replace_keywords('netns exec {} $IP link set $DUMMY up'.format(ns)))
180        cmds.append(self._replace_keywords('link set $DEV0 up'.format(ns)))
181
182        if self.args.device:
183            cmds.append(self._replace_keywords('link set $DEV2 netns {}'.format(ns)))
184            cmds.append(self._replace_keywords('netns exec {} $IP link set $DEV2 up'.format(ns)))
185
186        return cmds
187
188    def _ns_create(self):
189        '''
190        Create the network namespace in which the tests will be run and set up
191        the required network devices for it.
192        '''
193        self._exec_cmd_batched('pre', self._ns_create_cmds())
194
195    def _ns_destroy_cmd(self):
196        return self._replace_keywords('netns delete {}'.format(self.args.NAMES['NS']))
197
198    def _ns_destroy(self):
199        '''
200        Destroy the network namespace for testing (and any associated network
201        devices as well)
202        '''
203        self._exec_cmd('post', self._ns_destroy_cmd())
204
205    @cached_property
206    def _proc(self):
207        ip = self._replace_keywords("$IP -b -")
208        proc = subprocess.Popen(ip,
209            shell=True,
210            stdin=subprocess.PIPE,
211            env=ENVIR)
212
213        return proc
214
215    def _proc_check(self):
216        proc = self._proc
217
218        proc.poll()
219
220        if proc.returncode is not None and proc.returncode != 0:
221            raise RuntimeError("iproute2 exited with an error code")
222
223    def _exec_cmd(self, stage, command):
224        '''
225        Perform any required modifications on an executable command, then run
226        it in a subprocess and return the results.
227        '''
228
229        if self.args.verbose > 3:
230            print('_exec_cmd:  command "{}"'.format(command))
231
232        proc = self._proc
233
234        proc.stdin.write((command + '\n').encode())
235        proc.stdin.flush()
236
237        if self.args.verbose > 3:
238            print('_exec_cmd proc: {}'.format(proc))
239
240        self._proc_check()
241
242    def _exec_cmd_batched(self, stage, commands):
243        for cmd in commands:
244            self._exec_cmd(stage, cmd)
245
246    def _replace_keywords(self, cmd):
247        """
248        For a given executable command, substitute any known
249        variables contained within NAMES with the correct values
250        """
251        tcmd = Template(cmd)
252        subcmd = tcmd.safe_substitute(self.args.NAMES)
253        return subcmd
254