feat(controller): add pod exec api

Author: duanhongyi

charts/controller/templates/controller-clusterrole.yaml

@@ -31,6 +31,9 @@ rules:
 - apiGroups: [""]
   resources: ["pods/log"]
   verbs: ["get"]
+- apiGroups: [""]
+  resources: ["pods/exec"]
+  verbs: ["create"]
 - apiGroups: [""]
   resources: ["pods"]
   verbs: ["get", "list", "create", "delete"]

rootfs/api/asgi.py

@@ -10,7 +10,17 @@
 import os
 
 from django.core.asgi import get_asgi_application
+from channels.routing import ProtocolTypeRouter, URLRouter
+from channels.security.websocket import AllowedHostsOriginValidator
+
 
 os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'api.settings.production')
+http = get_asgi_application()
+
+from .routing import websocket_urlpatterns  # noqa
+from .middleware import ChannelOAuthMiddleware  # noqa
+websocket = AllowedHostsOriginValidator(ChannelOAuthMiddleware(URLRouter(
+    websocket_urlpatterns
+)))
 
-application = get_asgi_application()
+application = ProtocolTypeRouter({"http": http, "websocket": websocket})

rootfs/api/authentication.py

@@ -2,7 +2,7 @@
 from django.conf import settings
 from django.contrib.auth.models import AnonymousUser
 from django.core.cache import cache
-from django.utils.translation import gettext_lazy as _
+from django.utils.translation import gettext_lazy
 from rest_framework import authentication
 from rest_framework.authentication import TokenAuthentication, \
     get_authorization_header
@@ -31,20 +31,22 @@ def authenticate(self, request):
                 return None
 
             if len(auth) == 1:
-                msg = _('Invalid token header. No credentials provided.')
+                msg = gettext_lazy('Invalid token header. No credentials provided.')
                 raise exceptions.AuthenticationFailed(msg)
             elif len(auth) > 2:
-                msg = _('Invalid token header. Token string should not contain spaces.')  # noqa
+                msg = gettext_lazy(
+                    'Invalid token header. Token string should not contain spaces.')
                 raise exceptions.AuthenticationFailed(msg)
 
             try:
                 token = auth[1].decode()
             except UnicodeError:
-                msg = _('Invalid token header. Token string should not contain invalid characters.')  # noqa
+                msg = gettext_lazy(
+                    'Invalid token header. Token string should not contain invalid characters.')
                 raise exceptions.AuthenticationFailed(msg)
             return cache.get_or_set(
-                token, lambda: self._get_user(token), settings.OAUTH_CACHE_USER_TIME), None  # noqa
-        return super(DryccAuthentication, self).authenticate(request)  # noqa
+                token, lambda: self._get_user(token), settings.OAUTH_CACHE_USER_TIME), None
+        return super(DryccAuthentication, self).authenticate(request)
 
     @staticmethod
     def _get_user(key):
@@ -57,4 +59,4 @@ def _get_user(key):
             return user
         except Exception as e:
             logger.info(e)
-            raise exceptions.AuthenticationFailed(_('Verify token fail.'))
+            raise exceptions.AuthenticationFailed(gettext_lazy('Verify token fail.'))

rootfs/api/backend.py

@@ -56,7 +56,6 @@ class DryccOAuth(BaseOAuth2):
 
     def get_user_details(self, response):
         """Return user details from GitHub account"""
-        print(response)
         return {
             'username': response.get('username'),
             'email': response.get('email') or '',

rootfs/api/consumers.py

@@ -0,0 +1,95 @@
+import json
+import asyncio
+import collections
+from django.conf import settings
+
+from asgiref.sync import sync_to_async
+
+from kubernetes.client import Configuration
+from kubernetes.client.api import core_v1_api
+from kubernetes.stream import stream
+
+from channels.db import database_sync_to_async
+from channels.exceptions import DenyConnection
+from channels.generic.websocket import AsyncWebsocketConsumer
+
+from .models.app import App
+from .permissions import has_app_permission
+
+
+Request = collections.namedtuple("Request", ["user", "method"])
+
+
+class AppPodExecConsumer(AsyncWebsocketConsumer):
+
+    @property
+    def kubernetes(self):
+        with open('/var/run/secrets/kubernetes.io/serviceaccount/token') as token_file:
+            token = token_file.read()
+        config = Configuration(host=settings.SCHEDULER_URL)
+        config.api_key = {"authorization": "Bearer " + token}
+        config.verify_ssl = settings.K8S_API_VERIFY_TLS
+        if config.verify_ssl:
+            config.ssl_ca_cert = "/var/run/secrets/kubernetes.io/serviceaccount/ca.crt"
+        Configuration.set_default(config)
+        return core_v1_api.CoreV1Api()
+
+    async def check(self):
+        self.id = self.scope["url_route"]["kwargs"]["id"]
+        self.pod_id = self.scope["url_route"]["kwargs"]["pod_id"]
+        if self.scope["user"] is None:
+            return False, "user not login"
+        request = Request(self.scope["user"], "POST")
+        app = await database_sync_to_async(App.objects.get)(id=self.id)
+        return await database_sync_to_async(has_app_permission)(request, app)
+
+    async def connect(self):
+        self.stream = None
+        self.conneted = True
+        is_ok, message = await self.check()
+        if is_ok:
+            await self.accept()
+        else:
+            raise DenyConnection(message)
+
+    async def send(self, data):
+        if data is None:
+            return
+        elif isinstance(data, bytes):
+            await super().send(bytes_data=data)
+        elif isinstance(data, str):
+            await super().send(text_data=data)
+
+    async def task(self):
+        while self.stream.is_open() and self.conneted:
+            self.stream.update(timeout=9)
+            if await sync_to_async(self.stream.peek_stdout)():
+                data = self.stream.read_stdout()
+            elif await sync_to_async(self.stream.peek_stderr)():
+                data = self.stream.peek_stderr()
+            else:
+                data = None
+            await self.send(data)
+
+    async def disconnect(self, close_code):
+        if self.stream:
+            self.stream.close()
+        self.conneted = False
+
+    async def receive(self, text_data=None, bytes_data=None):
+        if self.stream is None and text_data is not None:
+            args = (self.kubernetes.connect_get_namespaced_pod_exec, self.pod_id, self.id)
+            kwargs = json.loads(text_data)
+            kwargs.update({"stderr": True, "stdout": True})
+            if kwargs["stdin"]:
+                kwargs.update({"_preload_content": False})
+                self.stream = stream(*args, **kwargs)
+                asyncio.create_task(self.task())
+            else:
+                await self.send(stream(*args, **kwargs))
+                await self.close(code=1000)
+        elif self.stream is not None:
+            data = text_data if text_data else bytes_data
+            await sync_to_async(self.stream.write_stdin)(data)
+        else:
+            raise ValueError("This operation is not supported!")

rootfs/api/management/commands/healthchecks.py

@@ -12,6 +12,5 @@ def handle(self, *args, **options):
             django.db.connection.cursor()
             print("Database is alive!")
         except Exception as e:
-            print("There was a problem connecting to the database")
-            print(str(e))
+            print("There was a problem connecting to the database", e)
             sys.exit(1)

rootfs/api/middleware.py

@@ -3,9 +3,13 @@
 
 See https://docs.djangoproject.com/en/1.11/topics/http/middleware/
 """
-
+import collections
 from api import __version__
 
+from channels.middleware import BaseMiddleware
+from channels.db import database_sync_to_async
+from api.authentication import DryccAuthentication
+
 
 class APIVersionMiddleware(object):
     """
@@ -26,3 +30,27 @@ def __call__(self, request):
         response['DRYCC_API_VERSION'] = version
         response['DRYCC_PLATFORM_VERSION'] = __version__
         return response
+
+
+class ChannelOAuthMiddleware(BaseMiddleware):
+    """
+    Middleware which populates scope["user"] from a auth2 token.
+    """
+    Request = collections.namedtuple('Request', ["META"])
+
+    async def get_user(self, scope):
+        headers = {}
+        for header in scope["headers"]:
+            if header[0] in (b"user-agent", b"authorization"):
+                key = "HTTP_%s" % header[0].decode().replace("-", "_").upper()
+                headers[key] = header[1].decode()
+        if len(headers) != 2:
+            return None
+        request = self.Request(headers)
+        user, _ = await database_sync_to_async(DryccAuthentication().authenticate)(request)
+        return user
+
+    async def __call__(self, scope, receive, send):
+        scope = dict(scope)
+        scope["user"] = await self.get_user(scope)
+        return await super().__call__(scope, receive, send)

rootfs/api/routing.py

@@ -0,0 +1,10 @@
+# chat/routing.py
+from django.urls import re_path
+
+from . import consumers
+
+websocket_urlpatterns = [
+    re_path(
+        r'^v2/apps/(?P<id>.*)/pods/(?P<pod_id>.*)/exec/?$',
+        consumers.AppPodExecConsumer.as_asgi())
+]

rootfs/drycc/gunicorn/config.py

@@ -1,6 +1,5 @@
 import os
 from os.path import dirname, realpath
-from multiprocessing import cpu_count
 
 import faulthandler
 faulthandler.enable()
@@ -13,7 +12,7 @@
 else:
     bind = '0.0.0.0:8000'
 
-workers = int(os.environ.get('GUNICORN_WORKERS', cpu_count() * 4 + 1))
+workers = int(os.environ.get('GUNICORN_WORKERS', 2))
 worker_class = "uvicorn.workers.UvicornWorker"
 pythonpath = dirname(dirname(dirname(realpath(__file__))))
 timeout = 1200

rootfs/requirements.txt

@@ -1,12 +1,13 @@
 # Drycc controller requirements
 backoff==1.11.1
 django==4.1
+channels==3.0.5
 django-cors-headers==3.7.0
 django-guardian==2.4.0
 djangorestframework==3.12.4
 docker==5.0.0
 gunicorn==20.1.0
-uvicorn==0.18.2
+uvicorn[standard]==0.18.2
 asgiref==3.5.2
 idna==3.2
 jmespath==0.10.0
@@ -28,3 +29,4 @@ dj-database-url==0.5.0
 social-auth-app-django==5.0.0
 python-jose==3.3.0
 Authlib==0.15.4
+kubernetes==24.2.0