1#!/usr/bin/python
2#
3# Copyright 2016 The Android Open Source Project
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9# http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16
17"""Unit tests for csocket."""
18
19from socket import *  # pylint: disable=wildcard-import
20import unittest
21
22import csocket
23
24
25LOOPBACK_IFINDEX = 1
26SOL_IPV6 = 41
27
28
29class CsocketTest(unittest.TestCase):
30
31  def _BuildSocket(self, family, addr):
32    s = socket(family, SOCK_DGRAM, 0)
33    s.bind((addr, 0))
34    return s
35
36  def CheckRecvfrom(self, family, addr):
37    s = self._BuildSocket(family, addr)
38    addr = s.getsockname()
39    sockaddr = csocket.Sockaddr(addr)
40    s.sendto("foo", addr)
41    data, addr = csocket.Recvfrom(s, 4096, 0)
42    self.assertEqual("foo", data)
43    self.assertEqual(sockaddr, addr)
44
45    s.close()
46
47  def testRecvfrom(self):
48    self.CheckRecvfrom(AF_INET, "127.0.0.1")
49    self.CheckRecvfrom(AF_INET6, "::1")
50
51  def CheckRecvmsg(self, family, addr):
52    s = self._BuildSocket(family, addr)
53
54    if family == AF_INET:
55      s.setsockopt(SOL_IP, csocket.IP_PKTINFO, 1)
56      s.setsockopt(SOL_IP, csocket.IP_RECVTTL, 1)
57      pktinfo_addr = inet_pton(AF_INET, addr)
58      pktinfo = (SOL_IP, csocket.IP_PKTINFO,
59                 csocket.InPktinfo((LOOPBACK_IFINDEX,
60                                    pktinfo_addr, pktinfo_addr)))
61      ttl = (SOL_IP, csocket.IP_TTL, 64)
62    elif family == AF_INET6:
63      s.setsockopt(SOL_IPV6, csocket.IPV6_RECVPKTINFO, 1)
64      s.setsockopt(SOL_IPV6, csocket.IPV6_RECVHOPLIMIT, 1)
65      pktinfo_addr = inet_pton(AF_INET6, addr)
66      pktinfo = (SOL_IPV6, csocket.IPV6_PKTINFO,
67                 csocket.In6Pktinfo((pktinfo_addr, LOOPBACK_IFINDEX)))
68      ttl = (SOL_IPV6, csocket.IPV6_HOPLIMIT, 64)
69
70    addr = s.getsockname()
71    sockaddr = csocket.Sockaddr(addr)
72    s.sendto("foo", addr)
73    data, addr, cmsg = csocket.Recvmsg(s, 4096, 1024, 0)
74    self.assertEqual("foo", data)
75    self.assertEqual(sockaddr, addr)
76    self.assertEqual([pktinfo, ttl], cmsg)
77
78    s.close()
79
80  def testRecvmsg(self):
81    self.CheckRecvmsg(AF_INET, "127.0.0.1")
82    self.CheckRecvmsg(AF_INET6, "::1")
83
84
85if __name__ == "__main__":
86  unittest.main()
87