summaryrefslogtreecommitdiffstats
path: root/ssh.py
diff options
context:
space:
mode:
Diffstat (limited to 'ssh.py')
-rw-r--r--ssh.py474
1 files changed, 242 insertions, 232 deletions
diff --git a/ssh.py b/ssh.py
index 004fdbad..1d7ebe32 100644
--- a/ssh.py
+++ b/ssh.py
@@ -28,254 +28,264 @@ import platform_utils
28from repo_trace import Trace 28from repo_trace import Trace
29 29
30 30
31PROXY_PATH = os.path.join(os.path.dirname(__file__), 'git_ssh') 31PROXY_PATH = os.path.join(os.path.dirname(__file__), "git_ssh")
32 32
33 33
34def _run_ssh_version(): 34def _run_ssh_version():
35 """run ssh -V to display the version number""" 35 """run ssh -V to display the version number"""
36 return subprocess.check_output(['ssh', '-V'], stderr=subprocess.STDOUT).decode() 36 return subprocess.check_output(
37 ["ssh", "-V"], stderr=subprocess.STDOUT
38 ).decode()
37 39
38 40
39def _parse_ssh_version(ver_str=None): 41def _parse_ssh_version(ver_str=None):
40 """parse a ssh version string into a tuple""" 42 """parse a ssh version string into a tuple"""
41 if ver_str is None: 43 if ver_str is None:
42 ver_str = _run_ssh_version() 44 ver_str = _run_ssh_version()
43 m = re.match(r'^OpenSSH_([0-9.]+)(p[0-9]+)?\s', ver_str) 45 m = re.match(r"^OpenSSH_([0-9.]+)(p[0-9]+)?\s", ver_str)
44 if m: 46 if m:
45 return tuple(int(x) for x in m.group(1).split('.')) 47 return tuple(int(x) for x in m.group(1).split("."))
46 else: 48 else:
47 return () 49 return ()
48 50
49 51
50@functools.lru_cache(maxsize=None) 52@functools.lru_cache(maxsize=None)
51def version(): 53def version():
52 """return ssh version as a tuple""" 54 """return ssh version as a tuple"""
53 try:
54 return _parse_ssh_version()
55 except FileNotFoundError:
56 print('fatal: ssh not installed', file=sys.stderr)
57 sys.exit(1)
58 except subprocess.CalledProcessError:
59 print('fatal: unable to detect ssh version', file=sys.stderr)
60 sys.exit(1)
61
62
63URI_SCP = re.compile(r'^([^@:]*@?[^:/]{1,}):')
64URI_ALL = re.compile(r'^([a-z][a-z+-]*)://([^@/]*@?[^/]*)/')
65
66
67class ProxyManager:
68 """Manage various ssh clients & masters that we spawn.
69
70 This will take care of sharing state between multiprocessing children, and
71 make sure that if we crash, we don't leak any of the ssh sessions.
72
73 The code should work with a single-process scenario too, and not add too much
74 overhead due to the manager.
75 """
76
77 # Path to the ssh program to run which will pass our master settings along.
78 # Set here more as a convenience API.
79 proxy = PROXY_PATH
80
81 def __init__(self, manager):
82 # Protect access to the list of active masters.
83 self._lock = multiprocessing.Lock()
84 # List of active masters (pid). These will be spawned on demand, and we are
85 # responsible for shutting them all down at the end.
86 self._masters = manager.list()
87 # Set of active masters indexed by "host:port" information.
88 # The value isn't used, but multiprocessing doesn't provide a set class.
89 self._master_keys = manager.dict()
90 # Whether ssh masters are known to be broken, so we give up entirely.
91 self._master_broken = manager.Value('b', False)
92 # List of active ssh sesssions. Clients will be added & removed as
93 # connections finish, so this list is just for safety & cleanup if we crash.
94 self._clients = manager.list()
95 # Path to directory for holding master sockets.
96 self._sock_path = None
97
98 def __enter__(self):
99 """Enter a new context."""
100 return self
101
102 def __exit__(self, exc_type, exc_value, traceback):
103 """Exit a context & clean up all resources."""
104 self.close()
105
106 def add_client(self, proc):
107 """Track a new ssh session."""
108 self._clients.append(proc.pid)
109
110 def remove_client(self, proc):
111 """Remove a completed ssh session."""
112 try: 55 try:
113 self._clients.remove(proc.pid) 56 return _parse_ssh_version()
114 except ValueError: 57 except FileNotFoundError:
115 pass 58 print("fatal: ssh not installed", file=sys.stderr)
116 59 sys.exit(1)
117 def add_master(self, proc): 60 except subprocess.CalledProcessError:
118 """Track a new master connection.""" 61 print("fatal: unable to detect ssh version", file=sys.stderr)
119 self._masters.append(proc.pid) 62 sys.exit(1)
120
121 def _terminate(self, procs):
122 """Kill all |procs|."""
123 for pid in procs:
124 try:
125 os.kill(pid, signal.SIGTERM)
126 os.waitpid(pid, 0)
127 except OSError:
128 pass
129
130 # The multiprocessing.list() API doesn't provide many standard list()
131 # methods, so we have to manually clear the list.
132 while True:
133 try:
134 procs.pop(0)
135 except:
136 break
137
138 def close(self):
139 """Close this active ssh session.
140
141 Kill all ssh clients & masters we created, and nuke the socket dir.
142 """
143 self._terminate(self._clients)
144 self._terminate(self._masters)
145
146 d = self.sock(create=False)
147 if d:
148 try:
149 platform_utils.rmdir(os.path.dirname(d))
150 except OSError:
151 pass
152
153 def _open_unlocked(self, host, port=None):
154 """Make sure a ssh master session exists for |host| & |port|.
155 63
156 If one doesn't exist already, we'll create it.
157 64
158 We won't grab any locks, so the caller has to do that. This helps keep the 65URI_SCP = re.compile(r"^([^@:]*@?[^:/]{1,}):")
159 business logic of actually creating the master separate from grabbing locks. 66URI_ALL = re.compile(r"^([a-z][a-z+-]*)://([^@/]*@?[^/]*)/")
160 """
161 # Check to see whether we already think that the master is running; if we
162 # think it's already running, return right away.
163 if port is not None:
164 key = '%s:%s' % (host, port)
165 else:
166 key = host
167
168 if key in self._master_keys:
169 return True
170
171 if self._master_broken.value or 'GIT_SSH' in os.environ:
172 # Failed earlier, so don't retry.
173 return False
174
175 # We will make two calls to ssh; this is the common part of both calls.
176 command_base = ['ssh', '-o', 'ControlPath %s' % self.sock(), host]
177 if port is not None:
178 command_base[1:1] = ['-p', str(port)]
179
180 # Since the key wasn't in _master_keys, we think that master isn't running.
181 # ...but before actually starting a master, we'll double-check. This can
182 # be important because we can't tell that that 'git@myhost.com' is the same
183 # as 'myhost.com' where "User git" is setup in the user's ~/.ssh/config file.
184 check_command = command_base + ['-O', 'check']
185 with Trace('Call to ssh (check call): %s', ' '.join(check_command)):
186 try:
187 check_process = subprocess.Popen(check_command,
188 stdout=subprocess.PIPE,
189 stderr=subprocess.PIPE)
190 check_process.communicate() # read output, but ignore it...
191 isnt_running = check_process.wait()
192
193 if not isnt_running:
194 # Our double-check found that the master _was_ infact running. Add to
195 # the list of keys.
196 self._master_keys[key] = True
197 return True
198 except Exception:
199 # Ignore excpetions. We we will fall back to the normal command and
200 # print to the log there.
201 pass
202
203 command = command_base[:1] + ['-M', '-N'] + command_base[1:]
204 p = None
205 try:
206 with Trace('Call to ssh: %s', ' '.join(command)):
207 p = subprocess.Popen(command)
208 except Exception as e:
209 self._master_broken.value = True
210 print('\nwarn: cannot enable ssh control master for %s:%s\n%s'
211 % (host, port, str(e)), file=sys.stderr)
212 return False
213 67
214 time.sleep(1)
215 ssh_died = (p.poll() is not None)
216 if ssh_died:
217 return False
218 68
219 self.add_master(p) 69class ProxyManager:
220 self._master_keys[key] = True 70 """Manage various ssh clients & masters that we spawn.
221 return True
222
223 def _open(self, host, port=None):
224 """Make sure a ssh master session exists for |host| & |port|.
225 71
226 If one doesn't exist already, we'll create it. 72 This will take care of sharing state between multiprocessing children, and
73 make sure that if we crash, we don't leak any of the ssh sessions.
227 74
228 This will obtain any necessary locks to avoid inter-process races. 75 The code should work with a single-process scenario too, and not add too
76 much overhead due to the manager.
229 """ 77 """
230 # Bail before grabbing the lock if we already know that we aren't going to
231 # try creating new masters below.
232 if sys.platform in ('win32', 'cygwin'):
233 return False
234
235 # Acquire the lock. This is needed to prevent opening multiple masters for
236 # the same host when we're running "repo sync -jN" (for N > 1) _and_ the
237 # manifest <remote fetch="ssh://xyz"> specifies a different host from the
238 # one that was passed to repo init.
239 with self._lock:
240 return self._open_unlocked(host, port)
241
242 def preconnect(self, url):
243 """If |uri| will create a ssh connection, setup the ssh master for it."""
244 m = URI_ALL.match(url)
245 if m:
246 scheme = m.group(1)
247 host = m.group(2)
248 if ':' in host:
249 host, port = host.split(':')
250 else:
251 port = None
252 if scheme in ('ssh', 'git+ssh', 'ssh+git'):
253 return self._open(host, port)
254 return False
255
256 m = URI_SCP.match(url)
257 if m:
258 host = m.group(1)
259 return self._open(host)
260
261 return False
262 78
263 def sock(self, create=True): 79 # Path to the ssh program to run which will pass our master settings along.
264 """Return the path to the ssh socket dir. 80 # Set here more as a convenience API.
265 81 proxy = PROXY_PATH
266 This has all the master sockets so clients can talk to them. 82
267 """ 83 def __init__(self, manager):
268 if self._sock_path is None: 84 # Protect access to the list of active masters.
269 if not create: 85 self._lock = multiprocessing.Lock()
270 return None 86 # List of active masters (pid). These will be spawned on demand, and we
271 tmp_dir = '/tmp' 87 # are responsible for shutting them all down at the end.
272 if not os.path.exists(tmp_dir): 88 self._masters = manager.list()
273 tmp_dir = tempfile.gettempdir() 89 # Set of active masters indexed by "host:port" information.
274 if version() < (6, 7): 90 # The value isn't used, but multiprocessing doesn't provide a set class.
275 tokens = '%r@%h:%p' 91 self._master_keys = manager.dict()
276 else: 92 # Whether ssh masters are known to be broken, so we give up entirely.
277 tokens = '%C' # hash of %l%h%p%r 93 self._master_broken = manager.Value("b", False)
278 self._sock_path = os.path.join( 94 # List of active ssh sesssions. Clients will be added & removed as
279 tempfile.mkdtemp('', 'ssh-', tmp_dir), 95 # connections finish, so this list is just for safety & cleanup if we
280 'master-' + tokens) 96 # crash.
281 return self._sock_path 97 self._clients = manager.list()
98 # Path to directory for holding master sockets.
99 self._sock_path = None
100
101 def __enter__(self):
102 """Enter a new context."""
103 return self
104
105 def __exit__(self, exc_type, exc_value, traceback):
106 """Exit a context & clean up all resources."""
107 self.close()
108
109 def add_client(self, proc):
110 """Track a new ssh session."""
111 self._clients.append(proc.pid)
112
113 def remove_client(self, proc):
114 """Remove a completed ssh session."""
115 try:
116 self._clients.remove(proc.pid)
117 except ValueError:
118 pass
119
120 def add_master(self, proc):
121 """Track a new master connection."""
122 self._masters.append(proc.pid)
123
124 def _terminate(self, procs):
125 """Kill all |procs|."""
126 for pid in procs:
127 try:
128 os.kill(pid, signal.SIGTERM)
129 os.waitpid(pid, 0)
130 except OSError:
131 pass
132
133 # The multiprocessing.list() API doesn't provide many standard list()
134 # methods, so we have to manually clear the list.
135 while True:
136 try:
137 procs.pop(0)
138 except: # noqa: E722
139 break
140
141 def close(self):
142 """Close this active ssh session.
143
144 Kill all ssh clients & masters we created, and nuke the socket dir.
145 """
146 self._terminate(self._clients)
147 self._terminate(self._masters)
148
149 d = self.sock(create=False)
150 if d:
151 try:
152 platform_utils.rmdir(os.path.dirname(d))
153 except OSError:
154 pass
155
156 def _open_unlocked(self, host, port=None):
157 """Make sure a ssh master session exists for |host| & |port|.
158
159 If one doesn't exist already, we'll create it.
160
161 We won't grab any locks, so the caller has to do that. This helps keep
162 the business logic of actually creating the master separate from
163 grabbing locks.
164 """
165 # Check to see whether we already think that the master is running; if
166 # we think it's already running, return right away.
167 if port is not None:
168 key = "%s:%s" % (host, port)
169 else:
170 key = host
171
172 if key in self._master_keys:
173 return True
174
175 if self._master_broken.value or "GIT_SSH" in os.environ:
176 # Failed earlier, so don't retry.
177 return False
178
179 # We will make two calls to ssh; this is the common part of both calls.
180 command_base = ["ssh", "-o", "ControlPath %s" % self.sock(), host]
181 if port is not None:
182 command_base[1:1] = ["-p", str(port)]
183
184 # Since the key wasn't in _master_keys, we think that master isn't
185 # running... but before actually starting a master, we'll double-check.
186 # This can be important because we can't tell that that 'git@myhost.com'
187 # is the same as 'myhost.com' where "User git" is setup in the user's
188 # ~/.ssh/config file.
189 check_command = command_base + ["-O", "check"]
190 with Trace("Call to ssh (check call): %s", " ".join(check_command)):
191 try:
192 check_process = subprocess.Popen(
193 check_command,
194 stdout=subprocess.PIPE,
195 stderr=subprocess.PIPE,
196 )
197 check_process.communicate() # read output, but ignore it...
198 isnt_running = check_process.wait()
199
200 if not isnt_running:
201 # Our double-check found that the master _was_ infact
202 # running. Add to the list of keys.
203 self._master_keys[key] = True
204 return True
205 except Exception:
206 # Ignore excpetions. We we will fall back to the normal command
207 # and print to the log there.
208 pass
209
210 command = command_base[:1] + ["-M", "-N"] + command_base[1:]
211 p = None
212 try:
213 with Trace("Call to ssh: %s", " ".join(command)):
214 p = subprocess.Popen(command)
215 except Exception as e:
216 self._master_broken.value = True
217 print(
218 "\nwarn: cannot enable ssh control master for %s:%s\n%s"
219 % (host, port, str(e)),
220 file=sys.stderr,
221 )
222 return False
223
224 time.sleep(1)
225 ssh_died = p.poll() is not None
226 if ssh_died:
227 return False
228
229 self.add_master(p)
230 self._master_keys[key] = True
231 return True
232
233 def _open(self, host, port=None):
234 """Make sure a ssh master session exists for |host| & |port|.
235
236 If one doesn't exist already, we'll create it.
237
238 This will obtain any necessary locks to avoid inter-process races.
239 """
240 # Bail before grabbing the lock if we already know that we aren't going
241 # to try creating new masters below.
242 if sys.platform in ("win32", "cygwin"):
243 return False
244
245 # Acquire the lock. This is needed to prevent opening multiple masters
246 # for the same host when we're running "repo sync -jN" (for N > 1) _and_
247 # the manifest <remote fetch="ssh://xyz"> specifies a different host
248 # from the one that was passed to repo init.
249 with self._lock:
250 return self._open_unlocked(host, port)
251
252 def preconnect(self, url):
253 """If |uri| will create a ssh connection, setup the ssh master for it.""" # noqa: E501
254 m = URI_ALL.match(url)
255 if m:
256 scheme = m.group(1)
257 host = m.group(2)
258 if ":" in host:
259 host, port = host.split(":")
260 else:
261 port = None
262 if scheme in ("ssh", "git+ssh", "ssh+git"):
263 return self._open(host, port)
264 return False
265
266 m = URI_SCP.match(url)
267 if m:
268 host = m.group(1)
269 return self._open(host)
270
271 return False
272
273 def sock(self, create=True):
274 """Return the path to the ssh socket dir.
275
276 This has all the master sockets so clients can talk to them.
277 """
278 if self._sock_path is None:
279 if not create:
280 return None
281 tmp_dir = "/tmp"
282 if not os.path.exists(tmp_dir):
283 tmp_dir = tempfile.gettempdir()
284 if version() < (6, 7):
285 tokens = "%r@%h:%p"
286 else:
287 tokens = "%C" # hash of %l%h%p%r
288 self._sock_path = os.path.join(
289 tempfile.mkdtemp("", "ssh-", tmp_dir), "master-" + tokens
290 )
291 return self._sock_path