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 19import socket 20# pylint: disable=g-importing-member 21from socket import AF_INET 22from socket import AF_INET6 23from socket import inet_pton 24from socket import SOCK_DGRAM 25from socket import SOL_IP 26# pylint: enable=g-importing-member 27 28import unittest 29 30import csocket 31 32 33LOOPBACK_IFINDEX = 1 34SOL_IPV6 = 41 35 36 37class CsocketTest(unittest.TestCase): 38 39 def _BuildSocket(self, family, addr): 40 s = socket.socket(family, SOCK_DGRAM, 0) 41 s.bind((addr, 0)) 42 return s 43 44 def CheckRecvfrom(self, family, addr): 45 s = self._BuildSocket(family, addr) 46 addr = s.getsockname() 47 sockaddr = csocket.Sockaddr(addr) 48 s.sendto("foo", addr) 49 data, addr = csocket.Recvfrom(s, 4096, 0) 50 self.assertEqual("foo", data) 51 self.assertEqual(sockaddr, addr) 52 53 s.close() 54 55 def testRecvfrom(self): 56 self.CheckRecvfrom(AF_INET, "127.0.0.1") 57 self.CheckRecvfrom(AF_INET6, "::1") 58 59 def CheckRecvmsg(self, family, addr): 60 s = self._BuildSocket(family, addr) 61 62 if family == AF_INET: 63 s.setsockopt(SOL_IP, csocket.IP_PKTINFO, 1) 64 s.setsockopt(SOL_IP, csocket.IP_RECVTTL, 1) 65 pktinfo_addr = inet_pton(AF_INET, addr) 66 pktinfo = (SOL_IP, csocket.IP_PKTINFO, 67 csocket.InPktinfo((LOOPBACK_IFINDEX, 68 pktinfo_addr, pktinfo_addr))) 69 ttl = (SOL_IP, csocket.IP_TTL, 64) 70 elif family == AF_INET6: 71 s.setsockopt(SOL_IPV6, csocket.IPV6_RECVPKTINFO, 1) 72 s.setsockopt(SOL_IPV6, csocket.IPV6_RECVHOPLIMIT, 1) 73 pktinfo_addr = inet_pton(AF_INET6, addr) 74 pktinfo = (SOL_IPV6, csocket.IPV6_PKTINFO, 75 csocket.In6Pktinfo((pktinfo_addr, LOOPBACK_IFINDEX))) 76 ttl = (SOL_IPV6, csocket.IPV6_HOPLIMIT, 64) 77 78 addr = s.getsockname() 79 sockaddr = csocket.Sockaddr(addr) 80 s.sendto("foo", addr) 81 data, addr, cmsg = csocket.Recvmsg(s, 4096, 1024, 0) 82 self.assertEqual("foo", data) 83 self.assertEqual(sockaddr, addr) 84 self.assertEqual([pktinfo, ttl], cmsg) 85 86 s.close() 87 88 def testRecvmsg(self): 89 self.CheckRecvmsg(AF_INET, "127.0.0.1") 90 self.CheckRecvmsg(AF_INET6, "::1") 91 92 93if __name__ == "__main__": 94 unittest.main() 95