• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright (c) 2013 The Chromium OS Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5import dpkt
6import socket
7import struct
8
9from lansim import tools
10
11
12# An initial set of Protocol to Hardware address mappings.
13_ARP_INITIAL_CACHE = {
14    # Broadcast address:
15    socket.inet_aton('255.255.255.255'): tools.inet_hwton('FF:FF:FF:FF:FF:FF'),
16}
17
18
19class SimpleHostError(Exception):
20    """A SimpleHost generic error."""
21
22
23class SimpleHost(object):
24    """A simple host supporting IPv4.
25
26    This class is useful as a base clase to implement other hosts. It supports
27    a single IPv4 address.
28    """
29    def __init__(self, sim, hw_addr, ip_addr):
30        """Creates the host and associates it with the given NetworkBridge.
31
32        @param sim: The Simulator interface where this host lives.
33        @param hw_addr: Hex or binary representation of the Ethernet address.
34        @param ip_addr: The IPv4 address. For example: "10.0.0.1".
35        """
36        self._sim = sim
37        self._hw_addr = hw_addr
38        self._ip_addr = ip_addr
39        self._bin_hw_addr = tools.inet_hwton(hw_addr)
40        self._bin_ip_addr = socket.inet_aton(ip_addr)
41        # arp cache: Protocol to Hardware address resolution cache.
42        self._arp_cache = dict(_ARP_INITIAL_CACHE)
43        # Reply to broadcast ARP requests.
44        rule = {
45            "dst": "\xff" * 6, # Broadcast HW addr.
46            "arp.pln": 4, # Protocol Addres Length is 4 (IP v4).
47            "arp.op": dpkt.arp.ARP_OP_REQUEST,
48            "arp.tpa": self._bin_ip_addr}
49        sim.add_match(rule, self.arp_request)
50
51        # Reply to unicast ARP requests.
52        rule["dst"] = self._bin_hw_addr
53        sim.add_match(rule, self.arp_request)
54
55        # Mappings used for TCP traffic forwarding.
56        self._tcp_fwd_in = {}
57        self._tcp_fwd_out = {}
58        self._tcp_fwd_ports = {}
59
60    @property
61    def ip_addr(self):
62        """Returns the host IPv4 address."""
63        return self._ip_addr
64
65
66    @property
67    def simulator(self):
68        """Returns the Simulator instance where this host runs on."""
69        return self._sim
70
71
72    def arp_request(self, pkt):
73        """Sends the ARP_REPLY matching the request.
74
75        @param pkt: a dpkt.Packet with the ARP_REQUEST.
76        """
77        # Update the local ARP cache whenever we get a request.
78        self.add_arp(hw_addr=pkt.arp.sha, ip_addr=pkt.arp.spa)
79
80        arp_resp = dpkt.arp.ARP(
81            op = dpkt.arp.ARP_OP_REPLY,
82            pln = 4,
83            tpa = pkt.arp.spa, # Target Protocol Address.
84            tha = pkt.arp.sha, # Target Hardware Address.
85            spa = self._bin_ip_addr, # Source Protocol Address.
86            sha = self._bin_hw_addr) # Source Hardware Address.
87        eth_resp = dpkt.ethernet.Ethernet(
88            dst = pkt.arp.sha,
89            src = self._bin_hw_addr,
90            type = dpkt.ethernet.ETH_TYPE_ARP,
91            data = arp_resp)
92        self._sim.write(eth_resp)
93
94
95    def add_arp(self, hw_addr, ip_addr):
96        """Maps the ip_addr to a given hw_addr.
97
98        This is useful to send IP packets with send_ip() to hosts that haven't
99        comunicate with us yet.
100
101        @param hw_addr: The network encoded corresponding Ethernet address.
102        @param ip_addr: The network encoded IPv4 address.
103        """
104        self._arp_cache[ip_addr] = hw_addr
105
106
107    def _resolve_mac_address(self, ip_addr):
108        """Resolves the hw_addr of an IP address locally when it is known.
109
110        This method uses the information gathered from received ARP requests and
111        locally added mappings with add_arp(). It also knows how to resolve
112        multicast addresses.
113
114        @param ip_addr: The IP address to resolve encoded in network format.
115        @return: The Hardware address encoded in network format or None
116        if unknown.
117        @raise SimpleHostError if the MAC address for ip_addr is unknown.
118        """
119        # From RFC 1112 6.4:
120        #  An IP host group address is mapped to an Ethernet multicast address
121        #  by placing the low-order 23-bits of the IP address into the low-order
122        #  23 bits of the Ethernet multicast address 01-00-5E-00-00-00 (hex).
123        #  Because there are 28 significant bits in an IP host group address,
124        #  more than one host group address may map to the same Ethernet
125        #  multicast address.
126        int_ip_addr, = struct.unpack('!I', ip_addr)
127        if int_ip_addr & 0xF0000000 == 0xE0000000: # Multicast IP address
128            int_hw_ending = int_ip_addr & ((1 << 23) - 1) | 0x5E000000
129            return '\x01\x00' + struct.pack('!I', int_hw_ending)
130        if ip_addr in self._arp_cache:
131            return self._arp_cache[ip_addr]
132        # No address found.
133        raise SimpleHostError("Unknown destination IP host.")
134
135
136    def send_ip(self, pkt):
137        """Sends an IP packet.
138
139        The source IP address and the hardware layer is automatically filled.
140        @param pkt: A dpkg.ip.IP packet.
141        @raise SimpleHostError if the MAC address for ip_addr is unknown.
142        """
143        hw_dst = self._resolve_mac_address(pkt.dst)
144
145        pkt.src = self._bin_ip_addr
146        # Set the packet length and force to recompute the checksum.
147        pkt.len = len(pkt)
148        pkt.sum = 0
149        hw_pkt = dpkt.ethernet.Ethernet(
150            dst = hw_dst,
151            src = self._bin_hw_addr,
152            type = dpkt.ethernet.ETH_TYPE_IP,
153            data = pkt)
154        return self._sim.write(hw_pkt)
155
156
157    def tcp_forward(self, port, dest_addr, dest_port):
158        """Forwards all the TCP/IP traffic on a given port to another host.
159
160        This method makes all the incoming traffic for this host on a particular
161        port be redirected to dest_addr:dest_port. This allows us to use the
162        kernel's network stack to handle that traffic.
163
164        @param port: The TCP port on this simulated host.
165        @param dest_addr: A host IP address on the same network in plain text.
166        @param dest_port: The TCP port on the destination host.
167        """
168        if not self._tcp_fwd_ports:
169            # Lazy initialization.
170            self._sim.add_match({
171                'ip.dst': self._bin_ip_addr,
172                'ip.p': dpkt.ip.IP_PROTO_TCP}, self._handle_tcp_forward)
173
174        self._tcp_fwd_ports[port] = socket.inet_aton(dest_addr), dest_port
175
176
177    def _tcp_pick_port(self, dhost, dport):
178        """Picks a new unused source TCP port on the host."""
179        for p in range(1024, 65536):
180            if (dhost, dport, p) in self._tcp_fwd_out:
181                continue
182            if p in self._tcp_fwd_ports:
183                continue
184            return p
185        raise SimpleHostError("Too many connections.")
186
187
188    def _handle_tcp_forward(self, pkt):
189        # Source from:
190        shost = pkt.ip.src
191        sport = pkt.ip.tcp.sport
192        dport = pkt.ip.tcp.dport
193
194        ### Handle responses from forwarded traffic back to the sender (out).
195        if (shost, sport, dport) in self._tcp_fwd_out:
196            fhost, fport, oport = self._tcp_fwd_out[(shost, sport, dport)]
197            # Redirect the packet
198            pkt.ip.tcp.sport = oport
199            pkt.ip.tcp.dport = fport
200            pkt.ip.dst = fhost
201            pkt.ip.tcp.sum = 0 # Force checksum
202            self.send_ip(pkt.ip)
203            return
204
205        ### Handle incoming traffic to a local forwarded port (in).
206        if dport in self._tcp_fwd_ports:
207            # Forward to:
208            fhost, fport = self._tcp_fwd_ports[dport]
209
210            ### Check if it is an existing connection.
211            # lport: The port from where we send data out.
212            if (shost, sport, dport) in self._tcp_fwd_in:
213                lport = self._tcp_fwd_in[(shost, sport, dport)]
214            else:
215                # Pick a new local port on our side.
216                lport = self._tcp_pick_port(fhost, fport)
217                self._tcp_fwd_in[(shost, sport, dport)] = lport
218                self._tcp_fwd_out[(fhost, fport, lport)] = (shost, sport, dport)
219
220            # Redirect the packet
221            pkt.ip.tcp.sport = lport
222            pkt.ip.tcp.dport = fport
223            pkt.ip.dst = fhost
224            pkt.ip.tcp.sum = 0 # Force checksum
225            self.send_ip(pkt.ip)
226            return
227
228
229    def socket(self, family, sock_type):
230        """Creates an asynchronous socket on the simulated host.
231
232        This method creates an asynchronous socket object that can be used to
233        receive and send packets. This module only supports UDP sockets.
234
235        @param family: The socket family, only AF_INET is supported.
236        @param sock_type: The socket type, only SOCK_DGRAM is supported.
237        @return: an UDPSocket object. See UDPSocket documentation for details.
238        @raise SimpleHostError if socket family and type is not supported.
239        """
240        if family != socket.AF_INET:
241            raise SimpleHostError("socket family not supported.")
242        if sock_type != socket.SOCK_DGRAM:
243            raise SimpleHostError("socket type not supported.")
244
245        return UDPSocket(self)
246
247
248class UDPSocket(object):
249    """An asynchronous UDP socket interface.
250
251    This UDP socket interface provides a way to send and received UDP messages
252    on an asynchronous way. This means that the socket doesn't have a recv()
253    method as a normal socket would have, since the simulation is event driven
254    and a callback should not block its execution. See the listen() method to
255    see how to receive messages from this socket.
256
257    This interface is used by modules outside lansim to interact with lansim
258    in a way that can be ported to other different backends. For example, this
259    same interface could be implemented using the Python's socket module and
260    the real kernel stack.
261    """
262    def __init__(self, host):
263        """Initializes the UDP socket.
264
265        To be used for receiving packets, listen() must be called.
266
267        @param host: A SimpleHost object.
268        """
269        self._host = host
270        self._sim = host.simulator
271        self._port = None
272
273
274    def __del__(self):
275        self.close()
276
277
278    def listen(self, ip_addr, port, recv_callback):
279        """Bind and listen on the ip_addr:port.
280
281        Calls recv_callback(pkt, src_addr, src_port) every time an UDP frame
282        is received. src_addr and src_port are passed with the source IPv4
283        (as in '192.168.0.2') and the sender port number.
284
285        This function can only be called once, since the socket can't be
286        reassigned.
287
288        @param ip_addr: Local destination ip_addr. If None, the Host's IPv4
289        address is used, for example '224.0.0.251' or '192.168.0.1'.
290        @param port: Local destination port number.
291        @param recv_callback: A callback function that accepts three
292        arguments, the received string, the sender IPv4 address and the
293        sender port number.
294        """
295        if ip_addr is None:
296            ip_addr = self._host.ip_addr()
297        self._port = port
298
299        # Binds all the traffic to the provided callback converting the
300        # single argument callback to the multiple argument.
301        self._sim.add_match({
302            "ip.dst": socket.inet_aton(ip_addr),
303            "ip.udp.dport": port},
304            lambda pkt: recv_callback(pkt.ip.udp.data,
305                                      socket.inet_ntoa(pkt.ip.src),
306                                      pkt.ip.udp.sport))
307
308
309    def send(self, data, ip_addr, port):
310        """Send an UDP message with the data string to ip_addr:port.
311
312        @param data: Any string small enough to fit in a single UDP packet.
313        @param ip_addr: Destination IPv4 address.
314        @param port: Destination UDP port number.
315        """
316        pkt_udp = dpkt.udp.UDP(
317            dport = port,
318            sport = self._port if self._port != None else 0,
319            data = data
320        )
321        # dpkt doesn't set the Length field on UDP packets according to RFC 768.
322        pkt_udp.ulen = len(pkt_udp.pack_hdr()) + len(str(pkt_udp.data))
323
324        pkt_ip = dpkt.ip.IP(
325            dst = socket.inet_aton(ip_addr),
326            ttl = 255, # The comon value for IP packets.
327            off = dpkt.ip.IP_DF, # Don't frag.
328            p = dpkt.ip.IP_PROTO_UDP,
329            data = pkt_udp
330        )
331        self._host.send_ip(pkt_ip)
332
333
334    def close(self):
335        """Closes the socket and disconnects the bound callback."""
336        #TODO(deymo): Remove the add_match rule added on listen().
337