fix #586 / set cpu_affinity() segfault: pre-emptively check if provided CPUs are valid before calling the C function

Author: giampaolo

HISTORY.rst

@@ -15,7 +15,8 @@ Bug tracker at https://github.com/giampaolo/psutil/issues
   "localhost"
 - #579: [Windows] fixed many compiler warnings.
 - #585: [FreeBSD] net_connections() may raise KeyError.
-
+- #586: [FreeBSD] cpu_affinity() segfaults on set in case an invalid CPU
+  number is provided.
 
 2.2.1 - 2015-02-02
 ==================

psutil/_psbsd.py

@@ -406,6 +406,14 @@ def cpu_affinity_get(self):
 
     @wrap_exceptions
     def cpu_affinity_set(self, cpus):
+        # Pre-emptively check if CPUs are valid because the C
+        # function has a weird behavior in case of invalid CPUs
+        #, see: https://github.com/giampaolo/psutil/issues/586
+        allcpus = tuple(range(len(per_cpu_times())))
+        for cpu in cpus:
+            if cpu not in allcpus:
+                raise ValueError("invalid CPU #%i (choose between %s)"
+                                 % (cpu, allcpus))
         try:
             cext.proc_cpu_affinity_set(self.pid, cpus)
         except OSError as err:
@@ -414,7 +422,6 @@ def cpu_affinity_set(self, cpus):
             # on because the set does not overlap with the thread's
             # anonymous mask>>
             if err.errno in (errno.EINVAL, errno.EDEADLK):
-                allcpus = tuple(range(len(per_cpu_times())))
                 for cpu in cpus:
                     if cpu not in allcpus:
                         raise ValueError("invalid CPU #%i (choose between %s)"